@braintrust/pi-extension 0.3.1 → 0.5.0

package/README.md

@@ -1,5 +1,7 @@
 # @braintrust/pi-extension
 
+[![npm version](https://img.shields.io/npm/v/%40braintrust%2Fpi-extension)](https://www.npmjs.com/package/@braintrust/pi-extension)
+
 Braintrust extension for [pi](https://github.com/mariozechner/pi-coding-agent).
 
 Today this extension automatically traces pi sessions, turns, model calls, and tool executions to Braintrust.
@@ -45,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.3.1",
+  "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",
@@ -24,12 +24,12 @@
     "access": "public"
   },
   "dependencies": {
-    "braintrust": "^3.8.0",
+    "braintrust": "^3.9.0",
     "valibot": "^1.3.1"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.67.2",
-    "@mariozechner/pi-coding-agent": "^0.67.2",
+    "@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();
 
@@ -411,6 +464,55 @@ describe("braintrustPiExtension", () => {
     );
   });
 
+  it("records the structured shutdown reason on the finalized root span", async () => {
+    const { emit } = await createHarness();
+
+    await emit("session_start");
+    await emit("before_agent_start", {
+      prompt: "Inspect the package",
+      images: [],
+    });
+    await emit("session_shutdown", { reason: "quit" });
+
+    const rootFinalizeLog = mockState.logSpans
+      .map((entry) => entry.event as Record<string, unknown>)
+      .find(
+        (event) =>
+          (event.metadata as Record<string, unknown> | undefined)?.last_close_reason === "quit",
+      );
+    expect(rootFinalizeLog).toBeDefined();
+    expect(mockState.endSpans.length).toBeGreaterThan(0);
+    expect(mockState.flushCalls).toBeGreaterThan(0);
+  });
+
+  it("does not finalize the root span on reload shutdowns", async () => {
+    const { emit } = await createHarness();
+
+    await emit("session_start");
+    await emit("before_agent_start", {
+      prompt: "Inspect the package",
+      images: [],
+    });
+
+    const startsBefore = mockState.startSpans.length;
+    const endsBefore = mockState.endSpans.length;
+    const flushesBefore = mockState.flushCalls;
+
+    await emit("session_shutdown", { reason: "reload" });
+
+    // No additional span endings during reload, but pending writes are still flushed.
+    expect(mockState.startSpans.length).toBe(startsBefore);
+    expect(mockState.endSpans.length).toBe(endsBefore);
+    expect(mockState.flushCalls).toBeGreaterThan(flushesBefore);
+    const reloadClose = mockState.logSpans
+      .map((entry) => entry.event as Record<string, unknown>)
+      .some(
+        (event) =>
+          (event.metadata as Record<string, unknown> | undefined)?.last_close_reason === "reload",
+      );
+    expect(reloadClose).toBe(false);
+  });
+
   it("hides all UI when showUi is false", async () => {
     mockState.config.showUi = false;
 

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,12 +135,58 @@ 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;
 }
 
-function getSessionStartReason(event: unknown): string | undefined {
+function getEventReason(event: unknown): string | undefined {
   if (!isPlainObject(event)) return undefined;
   return typeof event.reason === "string" ? event.reason : undefined;
 }
@@ -606,7 +660,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
   pi.on("session_start", async (event, ctx) => {
     refreshTracingUi(ctx);
 
-    const reason = getSessionStartReason(event);
+    const reason = getEventReason(event);
     if (reason === "new" || reason === "resume" || reason === "fork") {
       await rolloverSession(
         ctx,
@@ -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,
       },
@@ -837,13 +916,27 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
     await finishTurn("agent_end", Date.now(), finalAssistant);
   });
 
-  pi.on("session_shutdown", async (_event, ctx) => {
+  pi.on("session_shutdown", async (event, ctx) => {
     if (ctx.hasUI) {
       ctx.ui.setStatus(TRACING_STATUS_KEY, undefined);
       ctx.ui.setWidget(TRACING_WIDGET_KEY, undefined);
     }
+
+    // pi 0.68.0+ exposes a structured reason ("quit" | "reload" | "new" | "resume"
+    // | "fork"). Older pi hosts pass no payload, so we fall back to the generic
+    // label to stay backwards-compatible and keep the existing metadata shape.
+    const reason = getEventReason(event) ?? "session_shutdown";
+    logger.debug("session_shutdown", { reason });
+
     if (client && !clientInitializationError) {
-      await finalizeSession("session_shutdown");
+      // On reload the same pi session is about to resume in a freshly imported
+      // extension instance, which restores its state from the persisted store and
+      // keeps writing to the existing root span. Finalizing here would close that
+      // root span out from under the reloaded instance, so we just flush pending
+      // writes and let the new instance continue the trace.
+      if (reason !== "reload") {
+        await finalizeSession(reason);
+      }
       await client.flush();
     }
     activeSession = undefined;

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;