fix(pi-coding-agent): avoid double-counting token usage and cost

.changeset/fix-pi-token-rollups.md

@@ -0,0 +1,7 @@
+---
+"braintrust": patch
+---
+
+fix(pi-coding-agent): avoid double-counting token usage and cost
+
+Keep token usage on child LLM spans instead of repeating it on AgentSession.prompt task spans, so trace rollups count each provider call once.

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 80,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 4596,
-              "tokens": 4676
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 80,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 4596,
-              "tokens": 4676
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.txt

@@ -28,12 +28,7 @@ span_tree:
               "provider": "anthropic"
             }
             metrics: {
-              "completion_tokens": 6,
-              "duration": 0,
-              "prompt_cache_creation_tokens": 0,
-              "prompt_cached_tokens": 0,
-              "prompt_tokens": 10,
-              "tokens": 16
+              "duration": 0
             }
             ├── anthropic.messages.create [llm]
             │   input: [

e2e/scenarios/pi-coding-agent-instrumentation/assertions.ts

@@ -185,11 +185,7 @@ export function definePiCodingAgentInstrumentationAssertions(options: {
           "anthropic.messages.create",
           task?.span.id,
         );
-        const llm = findChildSpans(
-          events,
-          "anthropic.messages.create",
-          task?.span.id,
-        ).at(-1);
+        const llm = llmSpans.at(-1);
         const tool = findChildSpans(events, "bash", task?.span.id).at(-1);
         const toolChild = findChildSpans(
           events,
@@ -205,6 +201,17 @@ export function definePiCodingAgentInstrumentationAssertions(options: {
           "pi_coding_agent.operation": "AgentSession.prompt",
           provider: "anthropic",
         });
+        for (const metric of [
+          "prompt_tokens",
+          "completion_tokens",
+          "tokens",
+          "prompt_cached_tokens",
+          "prompt_cache_creation_tokens",
+          "estimated_cost",
+        ]) {
+          expect(task?.metrics).not.toHaveProperty(metric);
+        }
+        expect(task?.metrics).toMatchObject({ duration: expect.any(Number) });
 
         expect(llm).toBeDefined();
         expect(llmSpans).toHaveLength(2);
@@ -216,13 +223,13 @@ export function definePiCodingAgentInstrumentationAssertions(options: {
         expect(String(llm?.row.metadata?.model)).toContain("claude-haiku-4-5");
         expect(llm?.input).toEqual(expect.any(Array));
         expect(llm?.output).toBeDefined();
-        expect(llm?.metrics).toEqual(
-          expect.objectContaining({
+        for (const llmSpan of llmSpans) {
+          expect(llmSpan.metrics).toMatchObject({
             completion_tokens: expect.any(Number),
             prompt_tokens: expect.any(Number),
             tokens: expect.any(Number),
-          }),
-        );
+          });
+        }
 
         expect(tool).toBeDefined();
         expect(tool?.span.type).toBe("tool");

js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts

@@ -209,16 +209,16 @@ describe("PiCodingAgentPlugin", () => {
         expect.objectContaining({ output: { stdout: "pi_tool_ok" } }),
       );
       const finalTaskLog = taskSpan?.log.mock.calls.at(-1)?.[0];
-      expect(finalTaskLog?.metrics).toMatchObject({
-        completion_tokens: 3,
-        prompt_tokens: 5,
-        tokens: 8,
+      expect(finalTaskLog?.metrics).toEqual({
+        duration: expect.any(Number),
+        end: expect.any(Number),
+        start: expect.any(Number),
       });
       expect(taskSpan?.end).toHaveBeenCalledTimes(1);
     },
   );
 
-  it("counts cache reads and writes in prompt_tokens", async () => {
+  it("counts cached usage once across multiple LLM turns", async () => {
     const interceptor = promptInterceptor(enablePlugin(plugins));
     // Pi's `input` excludes cached tokens (its OpenAI client subtracts them from
     // `input_tokens`; Anthropic reports them separately), so 5 + 20 + 4 + 3 = 32.
@@ -237,10 +237,18 @@ describe("PiCodingAgentPlugin", () => {
 
     await interceptor(
       async function (this: typeof session) {
-        const stream = await this.agent.streamFn(anthropicModel(), {
-          messages: [{ role: "user", content: "hello" }],
-        });
-        await stream.result();
+        for (let turnIndex = 0; turnIndex < 2; turnIndex++) {
+          const stream = await this.agent.streamFn(anthropicModel(), {
+            messages: [{ role: "user", content: "hello" }],
+          });
+          await stream.result();
+          await this.agent.emit({
+            type: "turn_end",
+            message,
+            toolResults: [],
+            turnIndex,
+          });
+        }
       },
       session,
       ["hello", undefined],
@@ -254,16 +262,23 @@ describe("PiCodingAgentPlugin", () => {
       prompt_tokens: 29,
       tokens: 32,
     };
-    expect(
-      findSpan(spans, "anthropic.messages.create")?.log,
-    ).toHaveBeenCalledWith(
-      expect.objectContaining({ metrics: expect.objectContaining(expected) }),
+    const llmSpans = spans.filter(
+      (span) => span.args.spanAttributes.type === "llm",
     );
-    const finalTaskLog = findSpan(
-      spans,
-      "AgentSession.prompt",
-    )?.log.mock.calls.at(-1)?.[0];
-    expect(finalTaskLog?.metrics).toMatchObject(expected);
+    expect(llmSpans).toHaveLength(2);
+    for (const span of llmSpans) {
+      expect(span.log).toHaveBeenCalledWith(
+        expect.objectContaining({ metrics: expect.objectContaining(expected) }),
+      );
+    }
+    for (const [metric, value] of Object.entries(expected)) {
+      const total = spans.reduce(
+        (sum, span) =>
+          sum + (span.log.mock.calls.at(-1)?.[0].metrics?.[metric] ?? 0),
+        0,
+      );
+      expect(total).toBe(value * 2);
+    }
   });
 
   it("isolates overlapping prompts on the same agent without prompt matching", async () => {
@@ -360,6 +375,11 @@ describe("PiCodingAgentPlugin", () => {
       expect.objectContaining({ output: expect.any(Array) }),
     );
     expect(taskSpan?.end).toHaveBeenCalledTimes(1);
+    expect(taskSpan?.log.mock.calls.at(-1)?.[0].metrics).toEqual({
+      duration: expect.any(Number),
+      end: expect.any(Number),
+      start: expect.any(Number),
+    });
   });
 
   it("preserves full-iterator stream behavior", async () => {