perf(pi): avoid blocking tracing on streamed updates

bt-daemon/src/setup.rs

@@ -683,7 +683,8 @@ fn disable_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> {
 }
 
 fn update_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> {
-    runner.run("pi", &["update", PI_PACKAGE])
+    let package = format!("npm:{PI_PACKAGE}");
+    runner.run("pi", &["update", &package])
 }
 
 fn antigravity_home(config_dir: &Path) -> anyhow::Result<&Path> {
@@ -1341,12 +1342,12 @@ mod tests {
     }
 
     #[test]
-    fn pi_update_does_not_install_the_extension() {
+    fn pi_updates_the_npm_extension_without_installing_it() {
         let mut runner = FakeRunner::new([]);
 
         update_pi(&mut runner).unwrap();
 
-        assert!(runner.called("pi update @braintrust/pi-extension"));
+        assert!(runner.called("pi update npm:@braintrust/pi-extension"));
         assert!(!runner
             .calls
             .iter()

src/plugins/pi/content/src/daemon-adapter.test.ts

@@ -6,6 +6,7 @@ const mockState = vi.hoisted(() => ({
   closed: 0,
   claim: true,
   logGate: undefined as Promise<void> | undefined,
+  statusGate: undefined as Promise<void> | undefined,
 }));
 
 vi.mock("./runtime/daemon-client.ts", () => ({
@@ -21,6 +22,7 @@ vi.mock("./runtime/daemon-client.ts", () => ({
       return true;
     }
     async status(sessionId: string): Promise<Record<string, unknown>> {
+      await mockState.statusGate;
       return {
         daemon_version: "test",
         uptime_ms: 1,
@@ -66,6 +68,7 @@ describe("Pi daemon adapter", () => {
     mockState.closed = 0;
     mockState.claim = true;
     mockState.logGate = undefined;
+    mockState.statusGate = undefined;
   });
 
   it("does not register a duplicate managed adapter instance", async () => {
@@ -80,6 +83,82 @@ describe("Pi daemon adapter", () => {
     expect(handlers.size).toBe(0);
   });
 
+  it("does not block session startup on daemon status", async () => {
+    const handlers = new Map<string, (...args: unknown[]) => Promise<unknown>>();
+    const pi = {
+      on: (name: string, handler: (...args: unknown[]) => Promise<unknown>) =>
+        handlers.set(name, handler),
+    };
+    const ctx = {
+      cwd: "/tmp/project",
+      hasUI: true,
+      ui: { setStatus: vi.fn(), setWidget: vi.fn() },
+      sessionManager: {
+        getSessionFile: () => "/tmp/session.jsonl",
+        getSessionId: () => "native-session",
+      },
+    };
+    const { default: extension } = await import("./index.ts");
+    extension(pi as never);
+
+    let releaseStatus!: () => void;
+    mockState.statusGate = new Promise<void>((resolve) => {
+      releaseStatus = resolve;
+    });
+    let started = false;
+    const startup = handlers
+      .get("session_start")?.({ reason: "new" }, ctx)
+      .then(() => {
+        started = true;
+      });
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    expect(started).toBe(true);
+    releaseStatus();
+    await startup;
+  });
+
+  it("does not restore UI from a pending status refresh after shutdown", async () => {
+    const handlers = new Map<string, (...args: unknown[]) => Promise<unknown>>();
+    const statuses: unknown[] = [];
+    const widgets: unknown[] = [];
+    const pi = {
+      on: (name: string, handler: (...args: unknown[]) => Promise<unknown>) =>
+        handlers.set(name, handler),
+    };
+    const ctx = {
+      cwd: "/tmp/project",
+      hasUI: true,
+      ui: {
+        setStatus: (...args: unknown[]) => statuses.push(args),
+        setWidget: (...args: unknown[]) => widgets.push(args),
+      },
+      sessionManager: {
+        getSessionFile: () => "/tmp/session.jsonl",
+        getSessionId: () => "native-session",
+      },
+    };
+    const { default: extension } = await import("./index.ts");
+    extension(pi as never);
+
+    let releaseStatus!: () => void;
+    mockState.statusGate = new Promise<void>((resolve) => {
+      releaseStatus = resolve;
+    });
+    await handlers.get("session_start")?.({ reason: "new" }, ctx);
+    await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx);
+
+    expect(statuses.at(-1)).toEqual(["braintrust-tracing", undefined]);
+    expect(widgets.at(-1)).toEqual(["braintrust-trace-link", undefined]);
+
+    releaseStatus();
+    await Promise.resolve();
+    await Promise.resolve();
+
+    expect(statuses).toEqual([["braintrust-tracing", undefined]]);
+    expect(widgets).toEqual([["braintrust-trace-link", undefined]]);
+  });
+
   it("forwards native events and keeps the trace-link UI", async () => {
     const handlers = new Map<string, (...args: unknown[]) => Promise<unknown>>();
     const statuses: unknown[] = [];
@@ -126,6 +205,21 @@ describe("Pi daemon adapter", () => {
 
     await handlers.get("session_start")?.({ reason: "new" }, ctx);
     await handlers.get("before_agent_start")?.({ prompt: "hello" }, ctx);
+    await handlers.get("context")?.({ messages: [{ role: "user", content: "hello" }] }, ctx);
+    await handlers.get("message_update")?.({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_start" },
+    });
+    await handlers.get("message_update")?.({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "first" },
+      message: { role: "assistant", content: "first" },
+    });
+    await handlers.get("message_update")?.({
+      type: "message_update",
+      assistantMessageEvent: { type: "text_delta", delta: "second" },
+      message: { role: "assistant", content: "firstsecond" },
+    });
     let acknowledge!: () => void;
     mockState.logGate = new Promise<void>((resolve) => {
       acknowledge = resolve;
@@ -144,15 +238,24 @@ describe("Pi daemon adapter", () => {
     expect(mockState.flushes).toHaveLength(0);
     await handlers.get("input")?.({ text: "next turn" });
     expect(mockState.logs.at(-1)?.event).toBe("input");
+    await handlers.get("context")?.({ messages: [{ role: "user", content: "next" }] }, ctx);
+    await handlers.get("message_update")?.({
+      type: "message_update",
+      assistantMessageEvent: { type: "thinking_delta", delta: "first thought" },
+    });
     await handlers.get("session_compact")?.({}, ctx);
     await handlers.get("session_tree")?.({}, ctx);
     await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx);
 
     expect(mockState.logs.map((log) => log.event)).toEqual([
       "session_start",
       "before_agent_start",
+      "context",
+      "message_update",
       "agent_end",
       "input",
+      "context",
+      "message_update",
       "session_compact",
       "session_tree",
       "session_shutdown",
@@ -174,13 +277,22 @@ describe("Pi daemon adapter", () => {
       cwd: "/tmp/project",
       model: { provider: "openai", id: "gpt-5" },
     });
+    const streamingUpdates = mockState.logs.filter((log) => log.event === "message_update");
+    expect(streamingUpdates).toHaveLength(2);
+    expect(
+      streamingUpdates.map(
+        (log) =>
+          (log.payload as { event: { assistantMessageEvent: { type: string } } }).event
+            .assistantMessageEvent.type,
+      ),
+    ).toEqual(["text_delta", "thinking_delta"]);
     expect(mockState.flushes).toHaveLength(0);
     expect(widgets).toContainEqual([
       "braintrust-trace-link",
       ["Braintrust trace", "https://www.braintrust.dev/trace/1"],
       { placement: "belowEditor" },
     ]);
     expect(statuses.length).toBeGreaterThan(0);
-    expect(mockState.closed).toBe(1);
+    expect(mockState.closed).toBe(2);
   });
 });

src/plugins/pi/content/src/index.ts

@@ -8,6 +8,7 @@ import { EXTENSION_VERSION } from "./version.ts";
 
 const STATUS_KEY = "braintrust-tracing";
 const WIDGET_KEY = "braintrust-trace-link";
+const UI_STATUS_TIMEOUT_MS = 250;
 const PI_VERSION = loadPiPackageMetadata().version;
 
 function sessionKeyFor(
@@ -28,6 +29,23 @@ function nativePayload(value: unknown): unknown {
   }
 }
 
+type JsonObject = Record<string, unknown>;
+
+function asObject(value: unknown): JsonObject | undefined {
+  return value !== null && typeof value === "object" && !Array.isArray(value)
+    ? (value as JsonObject)
+    : undefined;
+}
+
+function streamingUpdateKind(value: unknown): string | undefined {
+  const event = asObject(value);
+  const update = asObject(event?.assistantMessageEvent);
+  if (typeof update?.type === "string") return update.type;
+  return typeof event?.type === "string" && event.type !== "message_update"
+    ? event.type
+    : undefined;
+}
+
 function sessionDescriptor(ctx: ExtensionContext): {
   sessionId: string;
   sessionFile?: string;
@@ -49,6 +67,8 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
 
   let sessionId: string | undefined;
   let lastContext: ExtensionContext | undefined;
+  let awaitingFirstToken = false;
+  let uiGeneration = 0;
   const client = new DaemonClient({
     source: "pi",
     pluginVersion: EXTENSION_VERSION,
@@ -58,18 +78,31 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       }
     },
   });
+  // status.get waits for the daemon's entire ingress queue to settle. Keep UI
+  // lookups off the event client so a busy daemon cannot block event capture.
+  const statusClient = new DaemonClient({
+    source: "pi",
+    pluginVersion: EXTENSION_VERSION,
+    requestTimeoutMs: UI_STATUS_TIMEOUT_MS,
+  });
 
   const remember = (ctx: ExtensionContext): ReturnType<typeof sessionDescriptor> => {
     lastContext = ctx;
     const descriptor = sessionDescriptor(ctx);
+    if (sessionId !== descriptor.sessionId) uiGeneration += 1;
     sessionId = descriptor.sessionId;
     return descriptor;
   };
 
   const refreshUi = async (ctx: ExtensionContext): Promise<void> => {
     if (!ctx.hasUI || !config.showUi || !sessionId) return;
-    const status = await client.status(sessionId);
-    const daemonSession = status?.sessions.find((session) => session.session_id === sessionId);
+    const refreshGeneration = uiGeneration;
+    const refreshSessionId = sessionId;
+    const status = await statusClient.status(refreshSessionId);
+    if (uiGeneration !== refreshGeneration || sessionId !== refreshSessionId) return;
+    const daemonSession = status?.sessions.find(
+      (session) => session.session_id === refreshSessionId,
+    );
     if (daemonSession?.last_error) {
       ctx.ui.setStatus(STATUS_KEY, `Braintrust: ${daemonSession.last_error}`);
     } else {
@@ -108,21 +141,33 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       },
       route: config.route,
     });
-    if (updateUi && ctx) await refreshUi(ctx);
+    if (updateUi && ctx) void refreshUi(ctx);
   };
 
   pi.on("session_start", async (event, ctx) => {
     await forward("session_start", event, ctx);
-    await refreshUi(ctx);
+    void refreshUi(ctx);
   });
   pi.on("input", async (event) => forward("input", event));
   pi.on("before_agent_start", async (event, ctx) => forward("before_agent_start", event, ctx));
-  pi.on("context", async (event, ctx) => forward("context", event, ctx));
+  pi.on("context", async (event, ctx) => {
+    awaitingFirstToken = true;
+    await forward("context", event, ctx);
+  });
   pi.on("before_provider_request", async (event) => forward("before_provider_request", event));
   pi.on("after_provider_response", async (event) => forward("after_provider_response", event));
-  pi.on("message_update", async (event) => forward("message_update", event));
+  pi.on("message_update", async (event) => {
+    if (!awaitingFirstToken) return;
+    const kind = streamingUpdateKind(event);
+    if (!kind || !["text_delta", "thinking_delta", "text", "thinking"].includes(kind)) return;
+    awaitingFirstToken = false;
+    await forward("message_update", event);
+  });
   pi.on("thinking_level_select", async (event) => forward("thinking_level_select", event));
-  pi.on("message_end", async (event) => forward("message_end", event));
+  pi.on("message_end", async (event) => {
+    awaitingFirstToken = false;
+    await forward("message_end", event);
+  });
   pi.on("tool_execution_start", async (event) => forward("tool_execution_start", event));
   pi.on("tool_execution_end", async (event, ctx) => forward("tool_execution_end", event, ctx));
   pi.on("session_before_compact", async (event, ctx) =>
@@ -135,12 +180,16 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
   pi.on("agent_end", async (event) => forward("agent_end", event));
   pi.on("session_shutdown", async (event, ctx) => {
     await forward("session_shutdown", event, ctx);
+    // Invalidate fire-and-forget refreshes before clearing the UI so a status
+    // request that finishes during shutdown cannot restore stale state.
+    uiGeneration += 1;
+    sessionId = undefined;
+    lastContext = undefined;
     if (ctx.hasUI) {
       ctx.ui.setStatus(STATUS_KEY, undefined);
       ctx.ui.setWidget(WIDGET_KEY, undefined);
     }
     await client.close();
-    sessionId = undefined;
-    lastContext = undefined;
+    await statusClient.close();
   });
 }