@braintrust/pi-extension 0.4.0 → 0.5.1

package/README.md

@@ -2,7 +2,7 @@
 
 [![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).
+Braintrust extension for [pi](https://github.com/earendil-works/pi-coding-agent).
 
 Today this extension automatically traces pi sessions, turns, model calls, and tool executions to Braintrust.
 
@@ -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 five 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.1",
   "description": "Braintrust extension for pi. Includes automatic tracing for pi sessions, turns, LLM calls, and tool executions to Braintrust.",
   "keywords": [
     "braintrust",
@@ -28,15 +28,15 @@
     "valibot": "^1.3.1"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.68.0",
-    "@mariozechner/pi-coding-agent": "^0.68.0",
+    "@earendil-works/pi-ai": "^0.74.0",
+    "@earendil-works/pi-coding-agent": "^0.74.0",
     "@types/node": "^25.6.0",
     "typescript": "^6.0.2",
     "vite-plus": "^0.1.16",
     "vitest": "^4.1.4"
   },
   "peerDependencies": {
-    "@mariozechner/pi-coding-agent": "*"
+    "@earendil-works/pi-coding-agent": "*"
   },
   "devEngines": {
     "packageManager": {
@@ -57,10 +57,11 @@
   },
   "scripts": {
     "preinstall": "node -e \"const userAgent = process.env.npm_config_user_agent || ''; if (process.env.INIT_CWD === process.cwd() && !userAgent.includes('pnpm/')) { console.error('Use pnpm in this repo.'); process.exit(1); }\"",
-    "check": "vp check",
+    "check": "node scripts/sync-version.mjs --check && vp check",
     "fmt": "vp fmt",
     "lint": "vp lint",
     "pack": "vp pack",
+    "sync:version": "node scripts/sync-version.mjs",
     "test": "vitest run",
     "test:integration": "vitest run src/index.integration.test.ts",
     "test:watch": "vitest",

package/src/config.test.ts

@@ -324,9 +324,8 @@ describe("loadConfig", () => {
 });
 
 describe("createLogger", () => {
-  it("writes json log lines to the default log file when debug is enabled", async () => {
-    const stateDir = makeTempDir("pi-extension-state-");
-    const config: TraceConfig = {
+  function makeLoggerConfig(overrides: Partial<TraceConfig> = {}): TraceConfig {
+    return {
       enabled: true,
       apiKey: "key",
       apiUrl: undefined,
@@ -335,14 +334,19 @@ describe("createLogger", () => {
       projectName: "pi",
       debug: true,
       logFile: undefined,
-      stateDir,
+      stateDir: makeTempDir("pi-extension-state-"),
       additionalMetadata: {},
       parentSpanId: undefined,
       rootSpanId: undefined,
       showUi: true,
       showTraceLink: true,
       configIssues: [],
+      ...overrides,
     };
+  }
+
+  it("writes json log lines to the default log file when debug is enabled", async () => {
+    const config = makeLoggerConfig();
 
     const logger = createLogger(config);
     logger.debug("debug message", { nested: { value: 1 } });
@@ -361,4 +365,27 @@ describe("createLogger", () => {
       message: "warn message",
     });
   });
+
+  it("writes warnings and errors to the default log file when debug is disabled", async () => {
+    const config = makeLoggerConfig({ debug: false, logFile: undefined });
+    const logger = createLogger(config);
+
+    logger.debug("debug message");
+    logger.info("info message");
+    logger.warn("warn message", { reason: "test" });
+    logger.error("error message");
+    await logger.flush();
+
+    const lines = readFileSync(logger.filePath, "utf8").trim().split("\n");
+    expect(lines).toHaveLength(2);
+    expect(JSON.parse(lines[0])).toMatchObject({
+      level: "warn",
+      message: "warn message",
+      data: { reason: "test" },
+    });
+    expect(JSON.parse(lines[1])).toMatchObject({
+      level: "error",
+      message: "error message",
+    });
+  });
 });

package/src/config.ts

@@ -450,14 +450,25 @@ export function createLogger(config: TraceConfig): Logger {
     config.logFile && config.logFile !== "true" && config.logFile !== "auto"
       ? config.logFile
       : join(config.stateDir, "braintrust-pi-extension.log");
-  const loggingEnabled = config.debug || Boolean(config.logFile);
+  const infoLoggingEnabled = config.debug || Boolean(config.logFile);
 
-  if (loggingEnabled) ensureDir(dirname(explicitLogFile));
+  let logDirEnsured = false;
+
+  function shouldLog(level: LogLevel): boolean {
+    return level === "warn" || level === "error" || infoLoggingEnabled;
+  }
+
+  function ensureLogDir(): void {
+    if (logDirEnsured) return;
+    ensureDir(dirname(explicitLogFile));
+    logDirEnsured = true;
+  }
 
   let pendingWrite = Promise.resolve();
 
   function emit(level: LogLevel, message: string, data?: unknown): void {
-    if (!loggingEnabled) return;
+    if (!shouldLog(level)) return;
+    ensureLogDir();
     pendingWrite = pendingWrite
       .catch(() => {})
       .then(async () => {

package/src/index.integration.test.ts

@@ -11,14 +11,14 @@ import {
   type Model,
   type SimpleStreamOptions,
   type ToolCall,
-} from "@mariozechner/pi-ai";
-import * as piCodingAgent from "@mariozechner/pi-coding-agent";
+} from "@earendil-works/pi-ai";
+import * as piCodingAgent from "@earendil-works/pi-coding-agent";
 import {
   createAgentSession,
   DefaultResourceLoader,
   SessionManager,
   type ExtensionAPI,
-} from "@mariozechner/pi-coding-agent";
+} from "@earendil-works/pi-coding-agent";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import braintrustPiExtension from "./index.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

@@ -1,5 +1,9 @@
+import { readFileSync } from "node:fs";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
+const packageVersion = (JSON.parse(readFileSync("package.json", "utf8")) as { version: string })
+  .version;
+
 const mockState = vi.hoisted(() => ({
   startSpans: [] as Array<Record<string, unknown>>,
   logSpans: [] as Array<Record<string, unknown>>,
@@ -193,6 +197,7 @@ describe("braintrustPiExtension", () => {
     expect(mockState.widgets.at(-1)?.content?.[1]).toBe(
       "braintrust.dev/app/test-org/p/pi/logs?oid=trace-row-1",
     );
+    expect(mockState.startSpans[0]?.metadata).toMatchObject({ extension_version: packageVersion });
 
     await emit("session_shutdown");
 
@@ -239,6 +244,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

@@ -1,6 +1,10 @@
-import { basename, resolve } from "node:path";
 import { hostname, userInfo } from "node:os";
-import type { AgentEndEvent, ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import { basename, resolve } from "node:path";
+import type {
+  AgentEndEvent,
+  ExtensionAPI,
+  ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
 import { BraintrustClient, type BraintrustSpanHandle } from "./client.ts";
 import { createLogger, loadConfig } from "./config.ts";
 import { createStateStore } from "./state.ts";
@@ -13,6 +17,7 @@ import type {
   NormalizedAssistantMessage,
   TraceConfig,
 } from "./types.ts";
+import { EXTENSION_VERSION } from "./version.ts";
 import {
   buildTurnInput,
   extractErrorText,
@@ -29,7 +34,6 @@ import {
   toUnixSeconds,
 } from "./utils.ts";
 
-const EXTENSION_VERSION = "0.1.0";
 const TRACING_STATUS_KEY = "braintrust-tracing";
 const TRACING_WIDGET_KEY = "braintrust-trace-link";
 
@@ -39,9 +43,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 +72,7 @@ interface ActiveTurn {
   lastAssistantMessage?: AssistantMessageLike;
   lastOutput?: NormalizedAssistantMessage;
   error?: string;
+  thinkingLevel?: string;
 }
 
 interface ActiveSession {
@@ -80,6 +91,7 @@ interface ActiveSession {
   startedAt?: number;
   totalTurns: number;
   totalToolCalls: number;
+  thinkingLevel?: string;
   currentTurn?: ActiveTurn;
 }
 
@@ -127,6 +139,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 +724,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 +742,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       lastAssistantMessage: undefined,
       lastOutput: undefined,
       error: undefined,
+      thinkingLevel: session.thinkingLevel,
     };
 
     store.patch(session.sessionKey, {
@@ -699,6 +759,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 +792,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 +818,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;

package/src/version.ts

@@ -0,0 +1,2 @@
+// This file is generated by scripts/sync-version.mjs. Do not edit by hand.
+export const EXTENSION_VERSION = "0.5.1";