@braintrust/pi-extension 0.3.1 → 0.4.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.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@braintrust/pi-extension",
-  "version": "0.3.1",
+  "version": "0.4.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.68.0",
+    "@mariozechner/pi-coding-agent": "^0.68.0",
     "@types/node": "^25.6.0",
     "typescript": "^6.0.2",
     "vite-plus": "^0.1.16",

package/src/index.test.ts

@@ -411,6 +411,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

@@ -132,7 +132,7 @@ function getPreviousSessionFile(event: unknown): string | 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 +606,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,
@@ -837,13 +837,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;