fix(pi-coding-agent): support streamFn -> streamFunction rename

.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

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<any>();
   return {
     state: { model: anthropicModel(), tools: [] as any[] },
-    streamFn,
+    [streamFnKey]: streamFn,
     subscribe: vi.fn((listener) => {
       listeners.add(listener);
       return vi.fn(() => listeners.delete(listener));

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<PiPromptState | undefined> {
   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);
   }
 

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;