From 426b5464087296ee7c0a29922873c53a33af1863 Mon Sep 17 00:00:00 2001 From: justcodebruh Date: Thu, 20 Aug 2026 12:52:28 -0400 Subject: [PATCH] fix(pi-coding-agent): support streamFn -> streamFunction rename @earendil-works/pi-coding-agent renamed agent.streamFn to agent.streamFunction starting in v0.81.0. wrapPiCodingAgentSDK and the Pi Coding Agent auto-instrumentation plugin only detected/patched streamFn, so tracing silently stopped emitting spans on pi >= 0.81 (observed on 0.84.2). Resolve whichever property actually exists on the agent and patch that one instead of assuming streamFn. --- .../pi-coding-agent-stream-function-rename.md | 5 +++ .../plugins/pi-coding-agent-plugin.test.ts | 42 ++++++++++++++++++- .../plugins/pi-coding-agent-plugin.ts | 34 ++++++++++++--- js/src/vendor-sdk-types/pi-coding-agent.ts | 6 ++- 4 files changed, 79 insertions(+), 8 deletions(-) create mode 100644 .changeset/pi-coding-agent-stream-function-rename.md diff --git a/.changeset/pi-coding-agent-stream-function-rename.md b/.changeset/pi-coding-agent-stream-function-rename.md new file mode 100644 index 000000000..47f974b76 --- /dev/null +++ b/.changeset/pi-coding-agent-stream-function-rename.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(pi-coding-agent): Support `agent.streamFunction` rename in `@earendil-works/pi-coding-agent` >= 0.81 diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts index 433361974..a589c492f 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts @@ -427,6 +427,41 @@ describe("PiCodingAgentPlugin", () => { await result; expect(taskSpan?.end).toHaveBeenCalledTimes(1); }); + + it("patches agent.streamFunction when streamFn is absent (pi >= 0.81)", async () => { + const interceptor = promptInterceptor(enablePlugin(plugins)); + const finalMessage = makeAssistantMessage("done"); + const originalStreamFn = vi.fn(async () => makeStream(finalMessage)); + const agent = makeAgent(originalStreamFn, "streamFunction"); + const session = makeSession(agent); + + await interceptor( + async function (this: typeof session) { + const stream = await this.agent.streamFunction( + anthropicModel(), + { systemPrompt: "system", messages: [], tools: [] }, + {}, + ); + await stream.result(); + await this.agent.emit({ + message: finalMessage, + toolResults: [], + turnIndex: 0, + type: "turn_end", + }); + }, + session, + ["hello", undefined], + { moduleVersion: "0.84.2" }, + ); + + expect(agent.streamFunction).not.toBe(originalStreamFn); + expect(agent.streamFn).toBeUndefined(); + expect(originalStreamFn).toHaveBeenCalled(); + + const taskSpan = findSpan(spans, "AgentSession.prompt"); + expect(taskSpan?.end).toHaveBeenCalledTimes(1); + }); }); function enablePlugin(plugins: PiCodingAgentPlugin[]): PiCodingAgentPlugin { @@ -521,11 +556,14 @@ function makeIteratorBackedStream(events: any[]) { }; } -function makeAgent(streamFn: any) { +function makeAgent( + streamFn: any, + streamFnKey: "streamFn" | "streamFunction" = "streamFn", +) { const listeners = new Set(); return { state: { model: anthropicModel(), tools: [] as any[] }, - streamFn, + [streamFnKey]: streamFn, subscribe: vi.fn((listener) => { listeners.add(listener); return vi.fn(() => listeners.delete(listener)); diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts index da6c50b30..adab4417a 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts @@ -68,6 +68,7 @@ type PiToolSpanState = { type PiAgentPatchState = { originalStreamFn: PiStreamFn; + streamFnKey: "streamFn" | "streamFunction"; wrappedStreamFn: PiStreamFn; }; @@ -228,11 +229,23 @@ function extractSession( function isPiAgent(value: unknown): value is PiAgent { return ( isObject(value) && - typeof value.streamFn === "function" && + (typeof value.streamFn === "function" || + typeof value.streamFunction === "function") && typeof value.subscribe === "function" ); } +// @earendil-works/pi-coding-agent renamed `agent.streamFn` to +// `agent.streamFunction` starting in v0.81.0. Resolve whichever property +// actually exists so both the pre- and post-rename shapes are patchable. +function resolveStreamFnKey( + agent: PiAgent, +): "streamFn" | "streamFunction" | undefined { + if (typeof agent.streamFn === "function") return "streamFn"; + if (typeof agent.streamFunction === "function") return "streamFunction"; + return undefined; +} + function promptContextStore(): IsoAsyncLocalStorage { piPromptContextStore ??= iso.newAsyncLocalStorage< PiPromptState | undefined @@ -245,17 +258,28 @@ function currentPiPromptState(): PiPromptState | undefined { } function installPiAgentInstrumentation(agent: PiAgent): void { + const streamFnKey = resolveStreamFnKey(agent); + if (!streamFnKey) { + // isPiAgent() already validated that one of these exists; this should be + // unreachable, but stay defensive rather than patching the wrong property. + throw new Error( + "Pi Coding Agent: unable to resolve streamFn/streamFunction property", + ); + } + const existing = piAgentPatchStates.get(agent); - if (!existing || agent.streamFn !== existing.wrappedStreamFn) { + if (!existing || agent[streamFnKey] !== existing.wrappedStreamFn) { + const originalStreamFn = agent[streamFnKey]; const patchState = { - originalStreamFn: agent.streamFn, - wrappedStreamFn: agent.streamFn, + originalStreamFn, + streamFnKey, + wrappedStreamFn: originalStreamFn, } satisfies PiAgentPatchState; patchState.wrappedStreamFn = makeInstrumentedStreamFn( agent, patchState.originalStreamFn, ); - agent.streamFn = patchState.wrappedStreamFn; + agent[streamFnKey] = patchState.wrappedStreamFn; piAgentPatchStates.set(agent, patchState); } diff --git a/js/src/vendor-sdk-types/pi-coding-agent.ts b/js/src/vendor-sdk-types/pi-coding-agent.ts index 13df2a1b9..a33888f97 100644 --- a/js/src/vendor-sdk-types/pi-coding-agent.ts +++ b/js/src/vendor-sdk-types/pi-coding-agent.ts @@ -39,7 +39,11 @@ export interface PiPromptOptions { } export interface PiAgent { - streamFn: PiStreamFn; + // Renamed from `streamFn` to `streamFunction` in + // @earendil-works/pi-coding-agent v0.81.0. Both are optional here since only + // one exists on any given installed version. + streamFn?: PiStreamFn; + streamFunction?: PiStreamFn; subscribe(listener: PiAgentEventListener): () => void; readonly state?: { model?: PiModel;