diff --git a/.changeset/fix-pi-token-rollups.md b/.changeset/fix-pi-token-rollups.md new file mode 100644 index 000000000..06b20e428 --- /dev/null +++ b/.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. diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.json index 141522a06..de10be740 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 80, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 4596, - "tokens": 4676 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.txt index e0f0873c9..2f90f6e73 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.json index cfe961420..93649800d 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.txt index ba235235a..ed734eb1d 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.json index cfe961420..93649800d 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.txt index ba235235a..ed734eb1d 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.json index 141522a06..de10be740 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 80, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 4596, - "tokens": 4676 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.txt index e0f0873c9..2f90f6e73 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.json index c94c8b853..92645a413 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.txt index 467bb2708..ba039908e 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.json index 160cc7985..10c8e6282 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.txt index e02b5f8d0..516cca194 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.json index 160cc7985..10c8e6282 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.txt index e02b5f8d0..516cca194 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.json b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.json index c94c8b853..92645a413 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.json +++ b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.json @@ -129,12 +129,7 @@ "provider": "anthropic" }, "metrics": { - "completion_tokens": 6, - "duration": 0, - "prompt_cache_creation_tokens": 0, - "prompt_cached_tokens": 0, - "prompt_tokens": 10, - "tokens": 16 + "duration": 0 } } ], diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.txt b/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.txt index 467bb2708..ba039908e 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.txt +++ b/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: [ diff --git a/e2e/scenarios/pi-coding-agent-instrumentation/assertions.ts b/e2e/scenarios/pi-coding-agent-instrumentation/assertions.ts index 1b13c19a7..6ad1ae742 100644 --- a/e2e/scenarios/pi-coding-agent-instrumentation/assertions.ts +++ b/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"); diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts index 2658b151b..462e7f04a 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts +++ b/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 () => { diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts index 9f71ed67d..9cb38ae16 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts @@ -41,10 +41,8 @@ type PiPromptState = { activeLlmSpans: Set; activeToolSpans: Map; agent: PiAgent; - collectedLlmUsageMetrics: boolean; deferCompletionUntilTurnEnd: boolean; finalized: boolean; - metrics: Record; onFinalize?: (state: PiPromptState) => void; metadata: Record; output?: unknown; @@ -202,13 +200,11 @@ function startPiPromptRun( activeLlmSpans: new Set(), activeToolSpans: new Map(), agent, - collectedLlmUsageMetrics: false, deferCompletionUntilTurnEnd: options?.streamingBehavior === "followUp" || options?.streamingBehavior === "steer", finalized: false, metadata, - metrics: {}, onFinalize, promptCallEnded: false, span, @@ -568,9 +564,6 @@ async function handlePiAgentEvent( state.turnEnded = true; if (isPiAssistantMessage(event.message)) { state.output = extractAssistantOutput(event.message); - if (!state.collectedLlmUsageMetrics) { - addMetrics(state.metrics, extractUsageMetrics(event.message.usage)); - } } if (state.promptCallEnded && state.deferCompletionUntilTurnEnd) { finishPiPromptRun(state); @@ -661,10 +654,8 @@ function finishPiPromptRun(state: PiPromptState, error?: unknown): void { safeLog(state.span, { ...(error ? { error: toLoggedError(error) } : {}), metadata, - metrics: { - ...cleanMetrics(state.metrics), - ...buildDurationMetrics(state.startTime), - }, + // Usage belongs to LLM spans; repeating it here doubles trace rollups and cost. + metrics: buildDurationMetrics(state.startTime), output: state.output, }); } finally { @@ -700,11 +691,6 @@ function finishPiLlmSpan( ...cleanMetrics(llmState.metrics), ...buildDurationMetrics(llmState.startTime), }; - const usageMetrics = extractUsageMetrics(message?.usage); - if (Object.keys(usageMetrics).length > 0) { - promptState.collectedLlmUsageMetrics = true; - addMetrics(promptState.metrics, usageMetrics); - } try { safeLog(llmState.span, { @@ -1073,15 +1059,6 @@ function extractUsageMetrics( }); } -function addMetrics( - target: Record, - source: Record, -): void { - for (const [key, value] of Object.entries(source)) { - target[key] = (target[key] ?? 0) + value; - } -} - function buildDurationMetrics(startTime: number): Record { const end = getCurrentUnixTimestamp(); return {