@braintrust/pi-extension 0.4.0 → 0.5.0

package/README.md

@@ -47,9 +47,9 @@ pi -e .
 
 ## Compatibility
 
-This package supports the **last three stable pi versions**.
+This package supports the **latest patch release from each of the last six stable pi minor versions**, currently excluding pi versions before `0.65.0`.
 
-Our GitHub Actions compatibility job automatically resolves and tests the latest patch release from each of the last three stable pi minor versions, so new pi releases are picked up without manually updating the matrix.
+Our GitHub Actions compatibility job automatically resolves and tests that compatibility window, so new pi releases are picked up without manually updating the matrix.
 
 ## Quick start
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@braintrust/pi-extension",
-  "version": "0.4.0",
+  "version": "0.5.0",
   "description": "Braintrust extension for pi. Includes automatic tracing for pi sessions, turns, LLM calls, and tool executions to Braintrust.",
   "keywords": [
     "braintrust",
@@ -28,8 +28,8 @@
     "valibot": "^1.3.1"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.68.0",
-    "@mariozechner/pi-coding-agent": "^0.68.0",
+    "@mariozechner/pi-ai": "^0.71.1",
+    "@mariozechner/pi-coding-agent": "^0.71.1",
     "@types/node": "^25.6.0",
     "typescript": "^6.0.2",
     "vite-plus": "^0.1.16",

package/src/index.integration.test.ts

@@ -149,6 +149,28 @@ function makeTempDir(prefix: string): string {
   return dir;
 }
 
+// Injected by CI (see .github/workflows/ci.yml). When unset (e.g. local dev) we
+// assume the currently installed pi is at least as new as any version we branch
+// on below.
+const PI_COMPAT_VERSION = process.env.PI_COMPAT_VERSION;
+
+function piCompatAtLeast(target: string): boolean {
+  if (!PI_COMPAT_VERSION) return true;
+  const parse = (v: string) =>
+    v
+      .split("-")[0]
+      .split(".")
+      .map((part) => Number.parseInt(part, 10) || 0);
+  const actual = parse(PI_COMPAT_VERSION);
+  const wanted = parse(target);
+  for (let i = 0; i < Math.max(actual.length, wanted.length); i += 1) {
+    const a = actual[i] ?? 0;
+    const w = wanted[i] ?? 0;
+    if (a !== w) return a > w;
+  }
+  return true;
+}
+
 function buildAssistantMessage(model: Model<Api>): AssistantMessage {
   return {
     role: "assistant",
@@ -585,9 +607,18 @@ describe("braintrustPiExtension integration", () => {
     const firstLlmSpanId = llmSpans[0]?.spanId;
 
     expect(toolSpans).toHaveLength(2);
+    // pi < 0.68.1 emits `tool_execution_end` in assistant source order, so the
+    // extension logs tool spans as [tool-1, tool-2]. Starting with pi 0.68.1 the
+    // agent emits parallel tool completions eagerly (completion order), so the
+    // fast `tool-2` finishes before the slow `tool-1` and spans are logged as
+    // [tool-2, tool-1]. See pi-coding-agent changelog 0.68.1 / issue #3503.
+    // TODO: drop the pi < 0.68.1 branch once we stop testing against it.
+    const expectedToolCallIdOrder = piCompatAtLeast("0.68.1")
+      ? ["tool-2", "tool-1"]
+      : ["tool-1", "tool-2"];
     expect(
       toolSpans.map((span) => (span.metadata as Record<string, unknown> | undefined)?.tool_call_id),
-    ).toEqual(["tool-1", "tool-2"]);
+    ).toEqual(expectedToolCallIdOrder);
     expect(toolSpans.map((span) => span.parentSpanId)).toEqual([firstLlmSpanId, firstLlmSpanId]);
   });
 

package/src/index.test.ts

@@ -239,6 +239,59 @@ describe("braintrustPiExtension", () => {
     expect(mockState.updateSpans).toEqual([]);
   });
 
+  it("records resolved model, thinking level, and provider response metadata on llm spans", async () => {
+    const { emit } = await createHarness();
+
+    await emit("session_start");
+    await emit("thinking_level_select", { level: "high", previousLevel: "off" });
+    await emit("before_agent_start", {
+      prompt: "Use a routed model",
+      images: [],
+    });
+    await emit("context", { messages: [{ role: "user", content: "Use a routed model" }] });
+    await emit("after_provider_response", {
+      status: 200,
+      headers: {
+        "x-ratelimit-remaining-requests": "42",
+        "retry-after": "5",
+        authorization: "secret",
+      },
+    });
+    await emit("message_end", {
+      message: {
+        role: "assistant",
+        provider: "openrouter",
+        model: "auto",
+        responseModel: "anthropic/claude-sonnet-4-5",
+        timestamp: 1_700_000_000_000,
+        content: [{ type: "text", text: "Done." }],
+      },
+    });
+
+    const turnSpan = mockState.startSpans.find(
+      (span) => span.type === "task" && span.name === "Turn 1",
+    );
+    const llmSpan = mockState.startSpans.find((span) => span.type === "llm");
+
+    expect(turnSpan?.metadata).toMatchObject({ thinking_level: "high" });
+    expect(llmSpan).toMatchObject({ name: "anthropic/claude-sonnet-4-5" });
+    expect(llmSpan?.metadata).toMatchObject({
+      model: "anthropic/claude-sonnet-4-5",
+      requested_model: "auto",
+      response_model: "anthropic/claude-sonnet-4-5",
+      thinking_level: "high",
+      provider_response_status: 200,
+      provider_response_headers: {
+        "x-ratelimit-remaining-requests": "42",
+        "retry-after": "5",
+      },
+    });
+    const llmMetadata = llmSpan?.metadata as
+      | { provider_response_headers?: Record<string, unknown> }
+      | undefined;
+    expect(llmMetadata?.provider_response_headers?.authorization).toBeUndefined();
+  });
+
   it("parents tool spans under the llm span that emitted the matching tool call", async () => {
     const { emit } = await createHarness();
 

package/src/index.ts

@@ -39,9 +39,15 @@ interface SessionDescriptor {
   sessionKey: string;
 }
 
+interface ProviderResponseMetadata {
+  status?: number;
+  headers?: Record<string, string>;
+}
+
 interface PendingLlmCall {
   startedAt: number;
   input: NormalizedAgentMessage[];
+  providerResponse?: ProviderResponseMetadata;
 }
 
 interface TrackedToolStart {
@@ -62,6 +68,7 @@ interface ActiveTurn {
   lastAssistantMessage?: AssistantMessageLike;
   lastOutput?: NormalizedAssistantMessage;
   error?: string;
+  thinkingLevel?: string;
 }
 
 interface ActiveSession {
@@ -80,6 +87,7 @@ interface ActiveSession {
   startedAt?: number;
   totalTurns: number;
   totalToolCalls: number;
+  thinkingLevel?: string;
   currentTurn?: ActiveTurn;
 }
 
@@ -127,6 +135,52 @@ function safeModelName(model: unknown): string | undefined {
   return undefined;
 }
 
+function stringProperty(
+  value: Record<string, unknown>,
+  keys: readonly string[],
+): string | undefined {
+  for (const key of keys) {
+    const item = value[key];
+    if (typeof item === "string" && item.trim()) return item;
+  }
+  return undefined;
+}
+
+function responseModelName(message: AssistantMessageLike): string | undefined {
+  return stringProperty(message as unknown as Record<string, unknown>, [
+    "responseModel",
+    "routedModel",
+    "resolvedModel",
+    "actualModel",
+    "concreteModel",
+    "outputModel",
+  ]);
+}
+
+function providerResponseMetadata(event: unknown): ProviderResponseMetadata | undefined {
+  if (!isPlainObject(event)) return undefined;
+  const metadata: ProviderResponseMetadata = {};
+  if (typeof event.status === "number") metadata.status = event.status;
+
+  const headers = event.headers;
+  if (isPlainObject(headers)) {
+    const allowedHeaders: Record<string, string> = {};
+    for (const [key, value] of Object.entries(headers)) {
+      const normalizedKey = key.toLowerCase();
+      if (!normalizedKey.startsWith("x-ratelimit-") && normalizedKey !== "retry-after") {
+        continue;
+      }
+      if (typeof value === "string") allowedHeaders[normalizedKey] = value;
+      else if (typeof value === "number" || typeof value === "boolean") {
+        allowedHeaders[normalizedKey] = String(value);
+      }
+    }
+    if (Object.keys(allowedHeaders).length > 0) metadata.headers = allowedHeaders;
+  }
+
+  return metadata.status !== undefined || metadata.headers ? metadata : undefined;
+}
+
 function getPreviousSessionFile(event: unknown): string | undefined {
   if (!isPlainObject(event)) return undefined;
   return typeof event.previousSessionFile === "string" ? event.previousSessionFile : undefined;
@@ -666,6 +720,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       metadata: {
         turn_number: session.totalTurns,
         active_model: safeModelName(ctx.model),
+        thinking_level: session.thinkingLevel,
       },
       name: `Turn ${session.totalTurns}`,
       type: "task",
@@ -683,6 +738,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       lastAssistantMessage: undefined,
       lastOutput: undefined,
       error: undefined,
+      thinkingLevel: session.thinkingLevel,
     };
 
     store.patch(session.sessionKey, {
@@ -699,6 +755,22 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
     });
   });
 
+  pi.on("after_provider_response", async (event) => {
+    if (!activeSession?.currentTurn) return;
+    const metadata = providerResponseMetadata(event);
+    if (!metadata) return;
+    const pending = [...activeSession.currentTurn.llmCalls]
+      .reverse()
+      .find((call) => !call.providerResponse);
+    if (pending) pending.providerResponse = metadata;
+  });
+
+  pi.on("thinking_level_select", async (event) => {
+    if (!isPlainObject(event) || typeof event.level !== "string") return;
+    if (activeSession) activeSession.thinkingLevel = event.level;
+    if (activeSession?.currentTurn) activeSession.currentTurn.thinkingLevel = event.level;
+  });
+
   pi.on("message_end", async (event) => {
     const session = activeSession;
     if (
@@ -716,7 +788,9 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       input: [{ role: "user", content: session.currentTurn.prompt }],
     };
 
-    const modelName = safeModelName(message) ?? message.model;
+    const requestedModelName = safeModelName(message) ?? message.model;
+    const responseModel = responseModelName(message);
+    const modelName = responseModel ?? requestedModelName;
     const endedAt = message.timestamp ?? Date.now();
     const normalizedOutput = normalizeAssistantMessage(message);
     const error =
@@ -740,7 +814,12 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
         api: message.api,
         provider: message.provider,
         model: modelName,
+        requested_model: requestedModelName,
+        response_model: responseModel,
         stop_reason: message.stopReason,
+        thinking_level: session.currentTurn.thinkingLevel ?? session.thinkingLevel,
+        provider_response_status: pending.providerResponse?.status,
+        provider_response_headers: pending.providerResponse?.headers,
         cache_read_tokens: message.usage?.cacheRead,
         cache_write_tokens: message.usage?.cacheWrite,
       },

package/src/types.ts

@@ -121,6 +121,7 @@ export interface AssistantMessageLike {
   api?: string;
   provider?: string;
   model?: string;
+  responseModel?: string;
   usage?: UsageLike;
   stopReason?: string;
   errorMessage?: string;