From 9d2c2b12897956943e2ed539c19e7331f8f58acc Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 28 Oct 2026 19:53:59 +0800 Subject: [PATCH 1/2] fix(pi): avoid blocking event handlers on trace delivery --- .../pi/content/src/daemon-adapter.test.ts | 6 +- src/plugins/pi/content/src/index.ts | 89 +++++++++++++------ 2 files changed, 67 insertions(+), 28 deletions(-) diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 09ba906a..435ae5c0 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -129,7 +129,9 @@ describe("Pi daemon adapter", () => { it("forwards the running host name and version instead of an installed peer", async () => { const handlers = await loadExtension(); - await handlers.get("session_start")?.({ reason: "new" }, fakeContext()); + const ctx = fakeContext(); + await handlers.get("session_start")?.({ reason: "new" }, ctx); + await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); expect(mockState.logs[0]).toMatchObject({ source: "pi", source_version: "18.12.6", @@ -242,7 +244,7 @@ describe("Pi daemon adapter", () => { turnEnded = true; }); await Promise.resolve(); - expect(turnEnded).toBe(false); + expect(turnEnded).toBe(true); acknowledge(); await turnEnd; mockState.logGate = undefined; diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 95bc7adf..f9f446da 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -100,6 +100,10 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { pluginVersion: EXTENSION_VERSION, requestTimeoutMs: UI_STATUS_TIMEOUT_MS, }); + // Pi applies a 30 second deadline to extension handlers. Daemon capture is + // durable but can wait behind disk or session work, so serialize forwarding + // in the extension and let ordinary Pi callbacks return immediately. + let forwarding = Promise.resolve(); const updateSession = async ( ctx: ExtensionContext, @@ -189,20 +193,40 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { if (updateUi && ctx) void refreshUi(ctx); }; - pi.on("session_start", async (event, ctx) => { - await forward("session_start", event, ctx); + const enqueueForward = ( + name: string, + event: unknown, + ctx?: ExtensionContext, + updateUi = false, + ): Promise => { + forwarding = forwarding.then(() => forward(name, event, ctx, updateUi)).catch((error) => { + if (lastContext?.hasUI && config.showUi) { + lastContext.ui.setStatus(STATUS_KEY, `Braintrust tracing unavailable: ${String(error)}`); + } + }); + return forwarding; + }; + + pi.on("session_start", (event, ctx) => { + void enqueueForward("session_start", event, 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) => { + pi.on("input", (event) => void enqueueForward("input", event)); + pi.on("before_agent_start", (event, ctx) => { + void enqueueForward("before_agent_start", event, ctx); + }); + pi.on("context", (event, ctx) => { awaitingFirstToken = true; lastStreamingKeepaliveAt = 0; - await forward("context", event, ctx); + void enqueueForward("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) => { + pi.on("before_provider_request", (event) => { + void enqueueForward("before_provider_request", event); + }); + pi.on("after_provider_response", (event) => { + void enqueueForward("after_provider_response", event); + }); + pi.on("message_update", (event) => { const kind = streamingUpdateKind(event); const tokenKinds = ["text_delta", "thinking_delta", "text", "thinking"]; const isFirstToken = awaitingFirstToken && !!kind && tokenKinds.includes(kind); @@ -211,29 +235,42 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { if (!isFirstToken && now - lastStreamingKeepaliveAt < 10_000) return; if (isFirstToken) awaitingFirstToken = false; lastStreamingKeepaliveAt = now; - await forward("message_update", event); + void enqueueForward("message_update", event); + }); + pi.on("thinking_level_select", (event) => { + void enqueueForward("thinking_level_select", event); }); - pi.on("thinking_level_select", async (event) => forward("thinking_level_select", event)); - pi.on("message_end", async (event) => { + pi.on("message_end", (event) => { awaitingFirstToken = false; lastStreamingKeepaliveAt = 0; - 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) => - forward("session_before_compact", event, ctx), - ); - pi.on("session_compact", async (event, ctx) => forward("session_compact", event, ctx, true)); - 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)); + void enqueueForward("message_end", event); + }); + pi.on("tool_execution_start", (event) => { + void enqueueForward("tool_execution_start", event); + }); + pi.on("tool_execution_end", (event, ctx) => { + void enqueueForward("tool_execution_end", event, ctx); + }); + pi.on("session_before_compact", (event, ctx) => { + void enqueueForward("session_before_compact", event, ctx); + }); + pi.on("session_compact", (event, ctx) => { + void enqueueForward("session_compact", event, ctx, true); + }); + pi.on("session_before_tree", (event, ctx) => { + void enqueueForward("session_before_tree", event, ctx); + }); + pi.on("session_tree", (event, ctx) => { + void enqueueForward("session_tree", event, ctx, true); + }); // All delivery flushing belongs to the daemon, including session shutdown. - pi.on("agent_end", async (event, ctx) => { - await forward("agent_end", event); - pollTraceLink(ctx); + pi.on("agent_end", (event, ctx) => { + void enqueueForward("agent_end", event).then(() => pollTraceLink(ctx)); }); pi.on("session_shutdown", async (event, ctx) => { - await forward("session_shutdown", event, ctx); + // Shutdown is the final lifecycle boundary; drain prior events and persist + // this one before releasing the daemon clients. + await enqueueForward("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; From b6ba8ad2b657ce3d1cc775e0c7df497ad9d919e9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 28 Oct 2026 20:15:10 +0800 Subject: [PATCH 2/2] fix(pi): preserve event timing and callback behavior --- .../pi/content/src/daemon-adapter.test.ts | 29 ++++++------------- src/plugins/pi/content/src/index.ts | 19 +++++++----- 2 files changed, 20 insertions(+), 28 deletions(-) diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index 435ae5c0..7a49f1b5 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -128,13 +128,16 @@ describe("Pi daemon adapter", () => { }); it("forwards the running host name and version instead of an installed peer", async () => { + const now = vi.spyOn(Date, "now").mockReturnValue(1_234); const handlers = await loadExtension(); const ctx = fakeContext(); await handlers.get("session_start")?.({ reason: "new" }, ctx); + now.mockReturnValue(5_678); await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); expect(mockState.logs[0]).toMatchObject({ source: "pi", source_version: "18.12.6", + ts_ms: 1_234, payload: { host_name: "@oh-my-pi/pi-coding-agent" }, }); }); @@ -147,17 +150,10 @@ describe("Pi daemon adapter", () => { 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); + expect(handlers.get("session_start")?.({ reason: "new" }, ctx)).toBeUndefined(); + await vi.waitFor(() => expect(mockState.statusCalls).toBe(1)); releaseStatus(); - await startup; + await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); }); it("does not restore UI from a pending status refresh after shutdown", async () => { @@ -237,19 +233,10 @@ describe("Pi daemon adapter", () => { mockState.logGate = new Promise((resolve) => { acknowledge = resolve; }); - let turnEnded = false; - const turnEnd = handlers - .get("agent_end")?.({ messages: [] }, ctx) - .then(() => { - turnEnded = true; - }); - await Promise.resolve(); - expect(turnEnded).toBe(true); + expect(handlers.get("agent_end")?.({ messages: [] }, ctx)).toBeUndefined(); acknowledge(); - await turnEnd; mockState.logGate = undefined; 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", @@ -389,6 +376,7 @@ describe("Pi daemon adapter", () => { now += 1; await update(); + await handlers.get("session_shutdown")?.({ reason: "quit" }, fakeContext()); expect(mockState.logs.filter((log) => log.event === "message_update")).toHaveLength(2); }); @@ -405,6 +393,7 @@ describe("Pi daemon adapter", () => { await handlers.get("session_start")?.({ reason: "resume" }, ctx); await handlers.get("context")?.({}, ctx); await handlers.get("tool_execution_end")?.({}, ctx); + await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); expect(mockState.logs[0]?.payload).toMatchObject({ legacy_resume: mockState.legacyContinuation, diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index f9f446da..2a1c38f0 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -169,6 +169,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { event: unknown, ctx?: ExtensionContext, updateUi = false, + timestamp = Date.now(), ): Promise => { const descriptor = ctx ? await updateSession(ctx) : undefined; if (!sessionId) return; @@ -177,7 +178,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { ...(PI_VERSION ? { source_version: PI_VERSION } : {}), session_id: sessionId, event: name, - ts_ms: Date.now(), + ts_ms: timestamp, payload: { event: nativePayload(event), extension_version: EXTENSION_VERSION, @@ -199,17 +200,19 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { ctx?: ExtensionContext, updateUi = false, ): Promise => { - forwarding = forwarding.then(() => forward(name, event, ctx, updateUi)).catch((error) => { - if (lastContext?.hasUI && config.showUi) { - lastContext.ui.setStatus(STATUS_KEY, `Braintrust tracing unavailable: ${String(error)}`); - } - }); + const timestamp = Date.now(); + forwarding = forwarding + .then(() => forward(name, event, ctx, updateUi, timestamp)) + .catch((error) => { + if (lastContext?.hasUI && config.showUi) { + lastContext.ui.setStatus(STATUS_KEY, `Braintrust tracing unavailable: ${String(error)}`); + } + }); return forwarding; }; pi.on("session_start", (event, ctx) => { - void enqueueForward("session_start", event, ctx); - void refreshUi(ctx); + void enqueueForward("session_start", event, ctx).then(() => refreshUi(ctx)); }); pi.on("input", (event) => void enqueueForward("input", event)); pi.on("before_agent_start", (event, ctx) => {