feat(tracing): add session trace metadata (#92)

Record the pi version, run mode, and session display name on root spans so Braintrust traces are easier to identify across pi releases and execution modes.

Use the session display name as the root span name when one is available, while preserving the existing workspace-based fallback.

Author: AbhiPrasad

src/index.test.ts

@@ -159,7 +159,9 @@ async function createHarness() {
     sessionManager: {
       getSessionFile: () => "/tmp/session.json",
       getSessionId: () => "session-1",
+      getSessionName: () => "Checkout fix",
     },
+    mode: "json",
   };
 
   async function emit(eventName: string, event: Record<string, unknown> = {}): Promise<void> {
@@ -244,6 +246,27 @@ describe("braintrustPiExtension", () => {
     expect(mockState.updateSpans).toEqual([]);
   });
 
+  it("annotates root spans with pi version, mode, and session name metadata", async () => {
+    const { emit } = await createHarness();
+
+    await emit("before_agent_start", {
+      prompt: "Inspect the package",
+      images: [],
+    });
+
+    expect(mockState.startSpans[0]).toMatchObject({
+      name: "Checkout fix",
+      type: "task",
+      metadata: {
+        source: "pi",
+        extension_version: packageVersion,
+        pi_version: expect.any(String),
+        pi_mode: "json",
+        session_name: "Checkout fix",
+      },
+    });
+  });
+
   it("annotates turn spans with idle input event metadata", async () => {
     const { emit } = await createHarness();
 
@@ -370,7 +393,7 @@ describe("braintrustPiExtension", () => {
     });
 
     const rootSpan = mockState.startSpans.find(
-      (span) => span.type === "task" && String(span.name).startsWith("pi:"),
+      (span) => span.type === "task" && span.parentSpanId === undefined,
     );
     const compactionSpan = mockState.startSpans.find((span) => span.name === "Compaction");
 
@@ -513,7 +536,9 @@ describe("braintrustPiExtension", () => {
       sessionManager: {
         getSessionFile: () => "/tmp/session.json",
         getSessionId: () => "session-1",
+        getSessionName: () => undefined,
       },
+      mode: "tui",
     };
 
     const emit = async (eventName: string, event: Record<string, unknown> = {}): Promise<void> => {

src/index.ts

@@ -1,9 +1,10 @@
 import { hostname, userInfo } from "node:os";
 import { basename, resolve } from "node:path";
-import type {
-  AgentEndEvent,
-  ExtensionAPI,
-  ExtensionContext,
+import {
+  VERSION as PI_VERSION,
+  type AgentEndEvent,
+  type ExtensionAPI,
+  type ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
 import { BraintrustClient, type BraintrustSpanHandle } from "./client.ts";
 import { createLogger, loadConfig } from "./config.ts";
@@ -241,12 +242,25 @@ function findLastAssistant(
   return undefined;
 }
 
+function getSessionName(ctx: ExtensionContext): string | undefined {
+  try {
+    const name = ctx.sessionManager.getSessionName?.();
+    return typeof name === "string" && name.trim() ? name : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
 function standardRootMetadata(ctx: ExtensionContext, config: TraceConfig): Record<string, unknown> {
   const descriptor = getSessionDescriptor(ctx);
+  const sessionName = getSessionName(ctx);
   return {
     ...config.additionalMetadata,
     source: "pi",
     extension_version: EXTENSION_VERSION,
+    pi_version: PI_VERSION,
+    pi_mode: ctx.mode,
+    session_name: sessionName,
     session_id: descriptor.sessionId,
     session_key: descriptor.sessionKey,
     session_file: descriptor.sessionFile,
@@ -487,7 +501,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void {
       rootSpanId: traceRootSpanId,
       parentSpanId,
       startedAt,
-      name: rootSpanName(ctx.cwd),
+      name: getSessionName(ctx) ?? rootSpanName(ctx.cwd),
       type: "task",
       metadata: {
         ...standardRootMetadata(ctx, config),