diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index c742283..f7c9303 100644 --- a/bt-daemon/src/setup.rs +++ b/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() diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 1254992..9608218 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/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 | undefined, + statusGate: undefined as Promise | undefined, })); vi.mock("./runtime/daemon-client.ts", () => ({ @@ -21,6 +22,7 @@ vi.mock("./runtime/daemon-client.ts", () => ({ return true; } async status(sessionId: string): Promise> { + 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 Promise>(); + const pi = { + on: (name: string, handler: (...args: unknown[]) => Promise) => + 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((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 Promise>(); + const statuses: unknown[] = []; + const widgets: unknown[] = []; + const pi = { + on: (name: string, handler: (...args: unknown[]) => Promise) => + 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((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 Promise>(); 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((resolve) => { acknowledge = resolve; @@ -144,6 +238,11 @@ 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); @@ -151,8 +250,12 @@ describe("Pi daemon adapter", () => { 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,6 +277,15 @@ 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", @@ -181,6 +293,6 @@ describe("Pi daemon adapter", () => { { placement: "belowEditor" }, ]); expect(statuses.length).toBeGreaterThan(0); - expect(mockState.closed).toBe(1); + expect(mockState.closed).toBe(2); }); }); diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 0445fe2..7eefaa2 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/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; + +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 => { lastContext = ctx; const descriptor = sessionDescriptor(ctx); + if (sessionId !== descriptor.sessionId) uiGeneration += 1; sessionId = descriptor.sessionId; return descriptor; }; const refreshUi = async (ctx: ExtensionContext): Promise => { 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(); }); }