@braintrust/pi-extension 0.1.0 → 0.2.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.2.0",
   "description": "Braintrust extension for pi. Includes automatic tracing for pi sessions, turns, LLM calls, and tool executions to Braintrust.",
   "keywords": [
     "braintrust",

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.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 {

package/src/types.ts

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