From c8db4d41146e976a1d0981ee330002d486dbe356 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 23 Jun 2026 12:36:53 -0400 Subject: [PATCH] feat: Track cache tokens --- src/client.ts | 2 ++ src/event-processor.ts | 14 +++++++-- src/span-sink.ts | 3 ++ src/test-helpers.ts | 68 ++++++++++++++++++++++++++++++++++++++---- src/tracing.test.ts | 38 +++++++++++++++++++++++ src/tracing.ts | 14 +++++++-- 6 files changed, 128 insertions(+), 11 deletions(-) diff --git a/src/client.ts b/src/client.ts index 42dcf72..f9baed1 100644 --- a/src/client.ts +++ b/src/client.ts @@ -50,6 +50,8 @@ export interface SpanData { prompt_tokens?: number completion_tokens?: number tokens?: number + prompt_cached_tokens?: number + prompt_cache_creation_tokens?: number reasoning_tokens?: number } context?: { diff --git a/src/event-processor.ts b/src/event-processor.ts index a94b9e3..ae66997 100644 --- a/src/event-processor.ts +++ b/src/event-processor.ts @@ -609,12 +609,18 @@ export class EventProcessor { // Mark as processed state.processedLlmMessages.add(messageId) - // Extract token info + // Extract token info. OpenCode reports input tokens exclusive of cache + // buckets, so normalize prompt_tokens to Braintrust's inclusive convention + // by adding cache read/write tokens back in. const tokens = messageInfo.tokens as Record | undefined + const cacheTokens = tokens?.cache as Record | undefined const inputTokens = (tokens?.input as number) || 0 const outputTokens = (tokens?.output as number) || 0 const reasoningTokens = (tokens?.reasoning as number) || 0 - const totalTokens = inputTokens + outputTokens + reasoningTokens + const cacheReadTokens = (cacheTokens?.read as number) || 0 + const cacheWriteTokens = (cacheTokens?.write as number) || 0 + const promptTokens = inputTokens + cacheReadTokens + cacheWriteTokens + const totalTokens = promptTokens + outputTokens + reasoningTokens // Extract model info const providerID = (messageInfo.providerID as string) || "unknown" @@ -662,9 +668,11 @@ export class EventProcessor { metrics: { start: msToSeconds(time.created as number), end: msToSeconds(time.completed as number), - prompt_tokens: inputTokens, + prompt_tokens: promptTokens, completion_tokens: outputTokens, tokens: totalTokens, + prompt_cached_tokens: cacheReadTokens || undefined, + prompt_cache_creation_tokens: cacheWriteTokens || undefined, reasoning_tokens: reasoningTokens || undefined, }, metadata: { diff --git a/src/span-sink.ts b/src/span-sink.ts index eed3523..a9df44b 100644 --- a/src/span-sink.ts +++ b/src/span-sink.ts @@ -65,6 +65,9 @@ export interface SpanTree { prompt_tokens?: number completion_tokens?: number tokens?: number + prompt_cached_tokens?: number + prompt_cache_creation_tokens?: number + reasoning_tokens?: number } metadata?: Record children: SpanTree[] diff --git a/src/test-helpers.ts b/src/test-helpers.ts index d367149..3ac1a0a 100644 --- a/src/test-helpers.ts +++ b/src/test-helpers.ts @@ -297,7 +297,12 @@ export function messageCompleted( sessionID: string, messageID: string, options?: { - tokens?: { input: number; output: number } + tokens?: { + input: number + output: number + reasoning?: number + cache?: { read?: number; write?: number } + } model?: { providerID: string; modelID: string } time?: { created: number; completed: number } }, @@ -317,10 +322,13 @@ export function messageCompleted( path: { cwd: "/test", root: "/test" }, cost: 0.001, tokens: { - input: options?.tokens?.input || 10, - output: options?.tokens?.output || 5, - reasoning: 0, - cache: { read: 0, write: 0 }, + input: options?.tokens?.input ?? 10, + output: options?.tokens?.output ?? 5, + reasoning: options?.tokens?.reasoning ?? 0, + cache: { + read: options?.tokens?.cache?.read ?? 0, + write: options?.tokens?.cache?.write ?? 0, + }, }, } return { @@ -379,6 +387,9 @@ export interface ExpectedSpan { prompt_tokens?: number completion_tokens?: number tokens?: number + prompt_cached_tokens?: number + prompt_cache_creation_tokens?: number + reasoning_tokens?: number } metadata?: Record children?: ExpectedSpan[] @@ -549,6 +560,21 @@ function spanMatchesSingle(actual: SpanTree, expected: ExpectedSpan): boolean { return false if (expected.metrics.tokens !== undefined && actual.metrics?.tokens !== expected.metrics.tokens) return false + if ( + expected.metrics.prompt_cached_tokens !== undefined && + actual.metrics?.prompt_cached_tokens !== expected.metrics.prompt_cached_tokens + ) + return false + if ( + expected.metrics.prompt_cache_creation_tokens !== undefined && + actual.metrics?.prompt_cache_creation_tokens !== expected.metrics.prompt_cache_creation_tokens + ) + return false + if ( + expected.metrics.reasoning_tokens !== undefined && + actual.metrics?.reasoning_tokens !== expected.metrics.reasoning_tokens + ) + return false } if (expected.input !== undefined) { @@ -630,6 +656,38 @@ export function getDiff(actual: SpanTree | null, expected: ExpectedSpan, path = `${path}.metrics.completion_tokens: expected ${expected.metrics.completion_tokens}, got ${actual.metrics?.completion_tokens}`, ) } + if ( + expected.metrics.tokens !== undefined && + actual.metrics?.tokens !== expected.metrics.tokens + ) { + diffs.push( + `${path}.metrics.tokens: expected ${expected.metrics.tokens}, got ${actual.metrics?.tokens}`, + ) + } + if ( + expected.metrics.prompt_cached_tokens !== undefined && + actual.metrics?.prompt_cached_tokens !== expected.metrics.prompt_cached_tokens + ) { + diffs.push( + `${path}.metrics.prompt_cached_tokens: expected ${expected.metrics.prompt_cached_tokens}, got ${actual.metrics?.prompt_cached_tokens}`, + ) + } + if ( + expected.metrics.prompt_cache_creation_tokens !== undefined && + actual.metrics?.prompt_cache_creation_tokens !== expected.metrics.prompt_cache_creation_tokens + ) { + diffs.push( + `${path}.metrics.prompt_cache_creation_tokens: expected ${expected.metrics.prompt_cache_creation_tokens}, got ${actual.metrics?.prompt_cache_creation_tokens}`, + ) + } + if ( + expected.metrics.reasoning_tokens !== undefined && + actual.metrics?.reasoning_tokens !== expected.metrics.reasoning_tokens + ) { + diffs.push( + `${path}.metrics.reasoning_tokens: expected ${expected.metrics.reasoning_tokens}, got ${actual.metrics?.reasoning_tokens}`, + ) + } } if (expected.error !== undefined && !nameMatches(actual.error, expected.error)) { diff --git a/src/tracing.test.ts b/src/tracing.test.ts index d75a759..8948f4a 100644 --- a/src/tracing.test.ts +++ b/src/tracing.test.ts @@ -70,6 +70,44 @@ describe("Event to Span Transformation", () => { ) }) + it("tracks OpenCode cache read/write tokens on LLM spans", async () => { + const sessionId = "ses_cache" + const messageId = "msg_cache" + + await assertEventsProduceTree( + session( + sessionId, + sessionCreated(sessionId), + chatMessage("Use the cached context"), + textPart(sessionId, messageId, "Done."), + messageCompleted(sessionId, messageId, { + tokens: { input: 10, output: 5, cache: { read: 3, write: 2 } }, + }), + sessionIdle(sessionId), + ), + { + span_attributes: { name: "OpenCode: test-project", type: "task" }, + children: [ + { + span_attributes: { name: "Turn 1", type: "task" }, + children: [ + { + span_attributes: { name: "anthropic/claude-3-haiku", type: "llm" }, + metrics: { + prompt_tokens: 15, + completion_tokens: 5, + tokens: 20, + prompt_cached_tokens: 3, + prompt_cache_creation_tokens: 2, + }, + }, + ], + }, + ], + }, + ) + }) + it("stores modelID on turn metadata from chat.message hook", async () => { const sessionId = "ses_turn_model" const messageId = "msg_turn_model" diff --git a/src/tracing.ts b/src/tracing.ts index 7d33c40..ef50a29 100644 --- a/src/tracing.ts +++ b/src/tracing.ts @@ -457,12 +457,18 @@ export function createTracingHooks( // Mark as processed state.processedLlmMessages.add(messageId) - // Extract token info + // Extract token info. OpenCode reports input tokens exclusive of cache + // buckets, so normalize prompt_tokens to Braintrust's inclusive + // convention by adding cache read/write tokens back in. const tokens = messageInfo.tokens as Record | undefined + const cacheTokens = tokens?.cache as Record | undefined const inputTokens = (tokens?.input as number) || 0 const outputTokens = (tokens?.output as number) || 0 const reasoningTokens = (tokens?.reasoning as number) || 0 - const totalTokens = inputTokens + outputTokens + reasoningTokens + const cacheReadTokens = (cacheTokens?.read as number) || 0 + const cacheWriteTokens = (cacheTokens?.write as number) || 0 + const promptTokens = inputTokens + cacheReadTokens + cacheWriteTokens + const totalTokens = promptTokens + outputTokens + reasoningTokens // Extract model info const providerID = (messageInfo.providerID as string) || "unknown" @@ -524,9 +530,11 @@ export function createTracingHooks( metrics: { start: msToSeconds(time.created as number), end: msToSeconds(time.completed as number), - prompt_tokens: inputTokens, + prompt_tokens: promptTokens, completion_tokens: outputTokens, tokens: totalTokens, + prompt_cached_tokens: cacheReadTokens || undefined, + prompt_cache_creation_tokens: cacheWriteTokens || undefined, reasoning_tokens: reasoningTokens || undefined, }, metadata: {