@braintrust/pi-extension 0.1.0 → 0.3.0

package/README.md

@@ -108,6 +108,8 @@ Example:
 | `additional_metadata` | `BRAINTRUST_ADDITIONAL_METADATA` | `{}` |
 | `log_file` | `BRAINTRUST_LOG_FILE` | unset |
 | `state_dir` | `BRAINTRUST_STATE_DIR` | `~/.pi/agent/state/braintrust-pi-extension` |
+| `show_ui` | `BRAINTRUST_SHOW_UI` | `true` |
+| `show_trace_link` | `BRAINTRUST_SHOW_TRACE_LINK` | `true` |
 | `parent_span_id` | `PI_PARENT_SPAN_ID` | unset |
 | `root_span_id` | `PI_ROOT_SPAN_ID` | unset |
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@braintrust/pi-extension",
-  "version": "0.1.0",
+  "version": "0.3.0",
   "description": "Braintrust extension for pi. Includes automatic tracing for pi sessions, turns, LLM calls, and tool executions to Braintrust.",
   "keywords": [
     "braintrust",
@@ -24,17 +24,16 @@
     "access": "public"
   },
   "dependencies": {
-    "braintrust": "^3.7.0",
+    "braintrust": "^3.8.0",
     "valibot": "^1.3.1"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.64.0",
-    "@mariozechner/pi-coding-agent": "^0.64.0",
-    "@types/node": "^25.5.0",
-    "tsx": "^4.21.0",
+    "@mariozechner/pi-ai": "^0.66.1",
+    "@mariozechner/pi-coding-agent": "^0.66.1",
+    "@types/node": "^25.6.0",
     "typescript": "^6.0.2",
-    "vite-plus": "^0.1.14",
-    "vitest": "^4.1.2"
+    "vite-plus": "^0.1.16",
+    "vitest": "^4.1.4"
   },
   "peerDependencies": {
     "@mariozechner/pi-coding-agent": "*"
@@ -66,6 +65,6 @@
     "test:integration": "vitest run src/index.integration.test.ts",
     "test:watch": "vitest",
     "typecheck": "vp check",
-    "smoke": "tsx -e \"import('./src/index.ts')\""
+    "smoke": "vp exec node -e \"import('./src/index.ts')\""
   }
 }

package/src/config.test.ts

@@ -16,6 +16,8 @@ const ENV_KEYS = [
   "BRAINTRUST_DEBUG",
   "BRAINTRUST_LOG_FILE",
   "BRAINTRUST_STATE_DIR",
+  "BRAINTRUST_SHOW_UI",
+  "BRAINTRUST_SHOW_TRACE_LINK",
   "PI_PARENT_SPAN_ID",
   "PI_ROOT_SPAN_ID",
   "BRAINTRUST_ADDITIONAL_METADATA",
@@ -257,6 +259,53 @@ describe("loadConfig", () => {
     });
   });
 
+  it("defaults showUi and showTraceLink to true", () => {
+    const home = makeTempDir("pi-extension-home-");
+    process.env.HOME = home;
+    process.env.BRAINTRUST_STATE_DIR = join(home, "state");
+
+    const config = loadConfig(home);
+
+    expect(config.showUi).toBe(true);
+    expect(config.showTraceLink).toBe(true);
+  });
+
+  it("applies show_ui and show_trace_link from config files", () => {
+    const home = makeTempDir("pi-extension-home-");
+    const cwd = join(home, "workspace");
+
+    process.env.HOME = home;
+    process.env.BRAINTRUST_STATE_DIR = join(home, "state");
+
+    writeJson(join(home, ".pi", "agent", "braintrust.json"), {
+      show_ui: false,
+    });
+
+    const config = loadConfig(cwd);
+
+    expect(config.showUi).toBe(false);
+    expect(config.showTraceLink).toBe(true);
+  });
+
+  it("overrides show_ui and show_trace_link from environment variables", () => {
+    const home = makeTempDir("pi-extension-home-");
+    const cwd = join(home, "workspace");
+
+    process.env.HOME = home;
+    process.env.BRAINTRUST_STATE_DIR = join(home, "state");
+    process.env.BRAINTRUST_SHOW_UI = "true";
+    process.env.BRAINTRUST_SHOW_TRACE_LINK = "false";
+
+    writeJson(join(home, ".pi", "agent", "braintrust.json"), {
+      show_ui: false,
+    });
+
+    const config = loadConfig(cwd);
+
+    expect(config.showUi).toBe(true);
+    expect(config.showTraceLink).toBe(false);
+  });
+
   it("warns when tracing is enabled without an API key", () => {
     const home = makeTempDir("pi-extension-home-");
     process.env.HOME = home;
@@ -290,6 +339,8 @@ describe("createLogger", () => {
       additionalMetadata: {},
       parentSpanId: undefined,
       rootSpanId: undefined,
+      showUi: true,
+      showTraceLink: true,
       configIssues: [],
     };
 

package/src/config.ts

@@ -218,6 +218,17 @@ function applyConfig(
   );
   if (additionalMetadata !== undefined) target.additionalMetadata = additionalMetadata;
 
+  const showUi = validateOptionalBoolean(source.show_ui, issues, path, "show_ui");
+  if (showUi !== undefined) target.showUi = showUi;
+
+  const showTraceLink = validateOptionalBoolean(
+    source.show_trace_link,
+    issues,
+    path,
+    "show_trace_link",
+  );
+  if (showTraceLink !== undefined) target.showTraceLink = showTraceLink;
+
   return {
     parentSpanConfigured: parentSpanId !== undefined,
     rootSpanConfigured: rootSpanId !== undefined,
@@ -238,6 +249,8 @@ export function loadConfig(cwd = process.cwd()): TraceConfig {
     additionalMetadata: {},
     parentSpanId: undefined,
     rootSpanId: undefined,
+    showUi: true,
+    showTraceLink: true,
     configIssues: [],
   };
 
@@ -344,6 +357,22 @@ export function loadConfig(cwd = process.cwd()): TraceConfig {
   );
   if (envStateDir !== undefined) config.stateDir = envStateDir;
 
+  const envShowUi = validateOptionalBoolean(
+    process.env.BRAINTRUST_SHOW_UI,
+    config.configIssues,
+    "BRAINTRUST_SHOW_UI",
+    "BRAINTRUST_SHOW_UI",
+  );
+  if (envShowUi !== undefined) config.showUi = envShowUi;
+
+  const envShowTraceLink = validateOptionalBoolean(
+    process.env.BRAINTRUST_SHOW_TRACE_LINK,
+    config.configIssues,
+    "BRAINTRUST_SHOW_TRACE_LINK",
+    "BRAINTRUST_SHOW_TRACE_LINK",
+  );
+  if (envShowTraceLink !== undefined) config.showTraceLink = envShowTraceLink;
+
   const envParentSpanId = validateOptionalNonEmptyString(
     process.env.PI_PARENT_SPAN_ID,
     config.configIssues,

package/src/index.integration.test.ts

@@ -433,6 +433,9 @@ async function createHarness(options?: {
     return { agentDir, cwd, session, stateDir };
   }
 
+  // TODO: Remove this legacy fallback once our supported pi compatibility window no
+  // longer includes pi <0.65.0, which introduced the session runtime API and the
+  // session_start-only post-transition model.
   const resourceLoader = new DefaultResourceLoader({
     cwd,
     agentDir,
@@ -447,26 +450,32 @@ async function createHarness(options?: {
     resourceLoader,
     sessionManager,
   });
+  const legacyRuntimeSession = legacySession as typeof legacySession & {
+    newSession(): Promise<boolean>;
+    switchSession(sessionPath: string): Promise<boolean>;
+    fork(entryId: string): Promise<{ cancelled: boolean; selectedText?: string }>;
+    sessionManager: SessionManager;
+  };
 
   const session: TestSessionController = {
-    prompt: (text) => legacySession.prompt(text),
-    newSession: () => legacySession.newSession(),
-    switchSession: (sessionPath) => legacySession.switchSession(sessionPath),
+    prompt: (text) => legacyRuntimeSession.prompt(text),
+    newSession: () => legacyRuntimeSession.newSession(),
+    switchSession: (sessionPath) => legacyRuntimeSession.switchSession(sessionPath),
     fork: async (entryId) => {
-      const result = await legacySession.fork(entryId);
+      const result = await legacyRuntimeSession.fork(entryId);
       return {
         cancelled: result.cancelled,
-        selectedText: result.selectedText,
+        selectedText: result.selectedText ?? "",
       };
     },
     dispose: async () => {
-      legacySession.dispose();
+      legacyRuntimeSession.dispose();
     },
     get sessionFile() {
-      return legacySession.sessionFile;
+      return legacyRuntimeSession.sessionFile;
     },
     get sessionManager() {
-      return legacySession.sessionManager;
+      return legacyRuntimeSession.sessionManager;
     },
   };
 

package/src/index.test.ts

@@ -22,6 +22,8 @@ const mockState = vi.hoisted(() => ({
     additionalMetadata: {},
     parentSpanId: undefined,
     rootSpanId: undefined,
+    showUi: true,
+    showTraceLink: true,
     configIssues: [] as Array<{ path: string; message: string; severity: "error" | "warning" }>,
   },
 }));
@@ -117,6 +119,8 @@ beforeEach(() => {
     additionalMetadata: {},
     parentSpanId: undefined,
     rootSpanId: undefined,
+    showUi: true,
+    showTraceLink: true,
     configIssues: [],
   };
   vi.resetModules();
@@ -406,4 +410,49 @@ describe("braintrustPiExtension", () => {
       undefined,
     );
   });
+
+  it("hides all UI when showUi is false", async () => {
+    mockState.config.showUi = false;
+
+    const { emit } = await createHarness();
+
+    await emit("session_start");
+    await emit("before_agent_start", {
+      prompt: "Inspect the package",
+      images: [],
+    });
+
+    const statusUpdates = mockState.statuses.filter(
+      (s) => s.key === "braintrust-tracing" && s.text !== undefined,
+    );
+    const widgetUpdates = mockState.widgets.filter(
+      (w) => w.key === "braintrust-trace-link" && w.content !== undefined,
+    );
+
+    expect(statusUpdates).toEqual([]);
+    expect(widgetUpdates).toEqual([]);
+  });
+
+  it("hides just the trace link when showTraceLink is false", async () => {
+    mockState.config.showTraceLink = false;
+
+    const { emit } = await createHarness();
+
+    await emit("session_start");
+    await emit("before_agent_start", {
+      prompt: "Inspect the package",
+      images: [],
+    });
+
+    const statusUpdates = mockState.statuses.filter(
+      (s) => s.key === "braintrust-tracing" && s.text !== undefined,
+    );
+    expect(statusUpdates.length).toBeGreaterThan(0);
+    expect(statusUpdates[0]?.text).toContain("Braintrust");
+
+    const widgetUpdates = mockState.widgets.filter(
+      (w) => w.key === "braintrust-trace-link" && w.content !== undefined,
+    );
+    expect(widgetUpdates).toEqual([]);
+  });
 });

package/src/index.ts

@@ -193,7 +193,10 @@ function setTracingStatus(
     configIssue?: ConfigIssue;
   },
 ): void {
-  if (!ctx.hasUI) return;
+  if (!ctx.hasUI || !config.showUi) {
+    if (ctx.hasUI) ctx.ui.setStatus(TRACING_STATUS_KEY, undefined);
+    return;
+  }
 
   const theme = ctx.ui.theme;
 
@@ -269,15 +272,19 @@ function displayPath(path: string): string {
 
 function setTraceWidget(
   ctx: ExtensionContext,
+  config: TraceConfig,
   traceUrl: string | undefined,
   configIssue: ConfigIssue | undefined,
 ): void {
-  if (!ctx.hasUI) return;
+  if (!ctx.hasUI || !config.showUi) {
+    if (ctx.hasUI) ctx.ui.setWidget(TRACING_WIDGET_KEY, undefined);
+    return;
+  }
 
   const theme = ctx.ui.theme;
   const lines: string[] = [];
 
-  if (traceUrl) {
+  if (traceUrl && config.showTraceLink) {
     const label = makeHyperlink(
       traceUrl,
       theme.fg("accent", theme.underline("Braintrust trace ↗")),
@@ -334,7 +341,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       missingApiKey: Boolean(config.enabled && !config.apiKey),
       configIssue,
     });
-    setTraceWidget(ctx, activeSession?.traceUrl, configIssue);
+    setTraceWidget(ctx, config, activeSession?.traceUrl, configIssue);
   }
 
   function persistTraceUrl(session: ActiveSession, traceUrl: string): void {
@@ -615,12 +622,21 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
     });
   });
 
-  pi.on("session_switch", async (event, ctx) => {
+  // TODO: Remove these legacy transition listeners once our compatibility window
+  // no longer includes pi <0.65.0.
+  const legacyPi = pi as ExtensionAPI & {
+    on(
+      event: "session_switch" | "session_fork",
+      handler: (event: unknown, ctx: ExtensionContext) => Promise<void> | void,
+    ): void;
+  };
+
+  legacyPi.on("session_switch", async (event, ctx) => {
     refreshTracingUi(ctx);
     await rolloverSession(ctx, "session_switch", getPreviousSessionFile(event));
   });
 
-  pi.on("session_fork", async (event, ctx) => {
+  legacyPi.on("session_fork", async (event, ctx) => {
     refreshTracingUi(ctx);
     await rolloverSession(ctx, "session_fork", getPreviousSessionFile(event));
   });

package/src/types.ts

@@ -26,6 +26,8 @@ export interface TraceConfig {
   additionalMetadata: JsonObject;
   parentSpanId?: string;
   rootSpanId?: string;
+  showUi: boolean;
+  showTraceLink: boolean;
   configIssues: ConfigIssue[];
 }