diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 76e10658..fce86156 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -9,6 +9,8 @@ const mockState = vi.hoisted(() => ({ legacyContinuationFor: vi.fn(), logGate: undefined as Promise | undefined, statusGate: undefined as Promise | undefined, + statusCalls: 0, + permalink: "https://www.braintrust.dev/trace/1" as string | undefined, })); vi.mock("./runtime/daemon-client.ts", () => ({ @@ -24,6 +26,7 @@ vi.mock("./runtime/daemon-client.ts", () => ({ return true; } async status(sessionId: string): Promise> { + mockState.statusCalls += 1; await mockState.statusGate; return { daemon_version: "test", @@ -34,7 +37,7 @@ vi.mock("./runtime/daemon-client.ts", () => ({ source: "pi", queued: 0, spans_emitted: 1, - permalink: "https://www.braintrust.dev/trace/1", + permalink: mockState.permalink, }, ], }; @@ -104,6 +107,8 @@ describe("Pi daemon adapter", () => { mockState.legacyContinuationFor.mockClear(); mockState.logGate = undefined; mockState.statusGate = undefined; + mockState.statusCalls = 0; + mockState.permalink = "https://www.braintrust.dev/trace/1"; }); it("does not register a duplicate managed adapter instance", async () => { @@ -212,7 +217,7 @@ describe("Pi daemon adapter", () => { }); let turnEnded = false; const turnEnd = handlers - .get("agent_end")?.({ messages: [] }) + .get("agent_end")?.({ messages: [] }, ctx) .then(() => { turnEnded = true; }); @@ -281,6 +286,69 @@ describe("Pi daemon adapter", () => { expect(mockState.closed).toBe(2); }); + it("polls for the trace link until the daemon delivers the first turn", async () => { + vi.useFakeTimers(); + try { + const widgets: unknown[] = []; + const ctx = fakeContext({ + setStatus: vi.fn(), + setWidget: (...args: unknown[]) => widgets.push(args), + }); + const handlers = await loadExtension(); + + // Pi starts and finishes the first turn before the daemon uploads it. + mockState.permalink = undefined; + await handlers.get("session_start")?.({ reason: "new" }, ctx); + await handlers.get("agent_end")?.({ messages: [] }, ctx); + await vi.advanceTimersByTimeAsync(10_000); + expect(widgets.at(-1)).toEqual([ + "braintrust-trace-link", + undefined, + { placement: "belowEditor" }, + ]); + + mockState.permalink = "https://www.braintrust.dev/trace/1"; + await vi.advanceTimersByTimeAsync(2_000); + expect(widgets.at(-1)).toEqual([ + "braintrust-trace-link", + ["Braintrust trace", "https://www.braintrust.dev/trace/1"], + { placement: "belowEditor" }, + ]); + + // Once the link is shown, neither the timer nor later turns query again. + const calls = mockState.statusCalls; + await handlers.get("agent_end")?.({ messages: [] }, ctx); + await vi.advanceTimersByTimeAsync(60_000); + expect(mockState.statusCalls).toBe(calls); + } finally { + vi.useRealTimers(); + } + }); + + it("stops polling for the trace link at its deadline and at shutdown", async () => { + vi.useFakeTimers(); + try { + const ctx = fakeContext(); + const handlers = await loadExtension(); + mockState.permalink = undefined; + await handlers.get("session_start")?.({ reason: "new" }, ctx); + + await handlers.get("agent_end")?.({ messages: [] }, ctx); + await vi.advanceTimersByTimeAsync(120_000); + const afterDeadline = mockState.statusCalls; + await vi.advanceTimersByTimeAsync(120_000); + expect(mockState.statusCalls).toBe(afterDeadline); + + await handlers.get("agent_end")?.({ messages: [] }, ctx); + await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); + const atShutdown = mockState.statusCalls; + await vi.advanceTimersByTimeAsync(120_000); + expect(mockState.statusCalls).toBe(atShutdown); + } finally { + vi.useRealTimers(); + } + }); + it("forwards legacy continuation state, read once per Pi session", async () => { mockState.legacyContinuation = { span: "legacy-root", diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 1d169186..be5afa6c 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -10,6 +10,11 @@ import { EXTENSION_VERSION } from "./version.ts"; const STATUS_KEY = "braintrust-tracing"; const WIDGET_KEY = "braintrust-trace-link"; const UI_STATUS_TIMEOUT_MS = 250; +// The daemon uploads a turn after Pi reports agent_end: about a second later +// with flush_on_turn_end, and when its 30 s batch window closes with +// fire_and_forget. Poll for the trace link long enough to cover both. +const TRACE_LINK_POLL_INTERVAL_MS = 2_000; +const TRACE_LINK_POLL_TIMEOUT_MS = 40_000; const PI_VERSION = runningPiPackage().version; function sessionKeyFor( @@ -72,6 +77,11 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { let lastContext: ExtensionContext | undefined; let awaitingFirstToken = false; let uiGeneration = 0; + // The session_start refresh runs before the daemon has delivered a trace, + // so agent runs poll status until the current session shows its link. + let traceLinkShown = false; + let traceLinkPollGeneration: number | undefined; + let traceLinkPollTimer: ReturnType | undefined; const client = new DaemonClient({ source: "pi", pluginVersion: EXTENSION_VERSION, @@ -94,7 +104,10 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { ): Promise> => { lastContext = ctx; const descriptor = sessionDescriptor(ctx); - if (sessionId !== descriptor.sessionId) uiGeneration += 1; + if (sessionId !== descriptor.sessionId) { + uiGeneration += 1; + traceLinkShown = false; + } sessionId = descriptor.sessionId; if (legacyContinuationSessionId !== descriptor.sessionId) { // The legacy file is immutable migration input. Read it asynchronously @@ -119,13 +132,30 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { } else { ctx.ui.setStatus(STATUS_KEY, "Braintrust: tracing"); } - ctx.ui.setWidget( - WIDGET_KEY, - config.showTraceLink && daemonSession?.permalink - ? ["Braintrust trace", daemonSession.permalink] - : undefined, - { placement: "belowEditor" }, - ); + const permalink = config.showTraceLink ? daemonSession?.permalink : undefined; + ctx.ui.setWidget(WIDGET_KEY, permalink ? ["Braintrust trace", permalink] : undefined, { + placement: "belowEditor", + }); + traceLinkShown = Boolean(permalink); + }; + + const pollTraceLink = (ctx: ExtensionContext): void => { + if (!ctx.hasUI || !config.showUi || !config.showTraceLink) return; + if (traceLinkShown || traceLinkPollGeneration === uiGeneration) return; + const generation = uiGeneration; + const deadline = Date.now() + TRACE_LINK_POLL_TIMEOUT_MS; + traceLinkPollGeneration = generation; + const poll = async (): Promise => { + traceLinkPollTimer = undefined; + if (generation === uiGeneration) await refreshUi(ctx); + if (traceLinkShown || generation !== uiGeneration || Date.now() >= deadline) { + if (traceLinkPollGeneration === generation) traceLinkPollGeneration = undefined; + return; + } + traceLinkPollTimer = setTimeout(() => void poll(), TRACE_LINK_POLL_INTERVAL_MS); + traceLinkPollTimer.unref?.(); + }; + void poll(); }; const forward = async ( @@ -189,12 +219,17 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { pi.on("session_before_tree", async (event, ctx) => forward("session_before_tree", event, ctx)); pi.on("session_tree", async (event, ctx) => forward("session_tree", event, ctx, true)); // All delivery flushing belongs to the daemon, including session shutdown. - pi.on("agent_end", async (event) => forward("agent_end", event)); + pi.on("agent_end", async (event, ctx) => { + await forward("agent_end", event); + pollTraceLink(ctx); + }); 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; + clearTimeout(traceLinkPollTimer); + traceLinkPollTimer = undefined; sessionId = undefined; legacyContinuation = undefined; legacyContinuationSessionId = undefined;