Further cache debugging code to track usage

Author: jahooma

common/src/types/contracts/llm.ts

@@ -25,6 +25,13 @@ export type StreamChunk =
     >
   | { type: 'error'; message: string }
 
+export type CacheDebugUsageData = {
+  inputTokens: number
+  outputTokens: number
+  cachedInputTokens: number
+  totalTokens: number
+}
+
 export type PromptAiSdkStreamFn = (
   params: {
     apiKey: string
@@ -45,6 +52,7 @@ export type PromptAiSdkStreamFn = (
       rawBody: unknown
       normalizedBody?: unknown
     }) => void
+    onCacheDebugUsageReceived?: (usage: CacheDebugUsageData) => void
     includeCacheControl?: boolean
     cacheDebugCorrelation?: string
     agentProviderOptions?: OpenRouterProviderRoutingOptions
@@ -79,6 +87,7 @@ export type PromptAiSdkFn = (
       rawBody: unknown
       normalizedBody?: unknown
     }) => void
+    onCacheDebugUsageReceived?: (usage: CacheDebugUsageData) => void
     includeCacheControl?: boolean
     cacheDebugCorrelation?: string
     agentProviderOptions?: OpenRouterProviderRoutingOptions
@@ -114,6 +123,7 @@ export type PromptAiSdkStructuredInput<T> = {
     rawBody: unknown
     normalizedBody?: unknown
   }) => void
+  onCacheDebugUsageReceived?: (usage: CacheDebugUsageData) => void
   includeCacheControl?: boolean
   cacheDebugCorrelation?: string
   agentProviderOptions?: OpenRouterProviderRoutingOptions

packages/agent-runtime/src/prompt-agent-stream.ts

@@ -3,7 +3,7 @@ import { globalStopSequence } from './constants'
 import type { AgentTemplate } from './templates/types'
 import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
 import type { SendActionFn } from '@codebuff/common/types/contracts/client'
-import type { PromptAiSdkStreamFn } from '@codebuff/common/types/contracts/llm'
+import type { CacheDebugUsageData, PromptAiSdkStreamFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsOf } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
@@ -32,6 +32,7 @@ export const getAgentStreamFromTemplate = (params: {
     rawBody: unknown
     normalizedBody?: unknown
   }) => void
+  onCacheDebugUsageReceived?: (usage: CacheDebugUsageData) => void
 
   onCostCalculated?: (credits: number) => Promise<void>
   promptAiSdkStream: PromptAiSdkStreamFn
@@ -55,6 +56,7 @@ export const getAgentStreamFromTemplate = (params: {
     userInputId,
     cacheDebugCorrelation,
     onCacheDebugProviderRequestBuilt,
+    onCacheDebugUsageReceived,
 
     sendAction,
     onCostCalculated,
@@ -90,6 +92,7 @@ export const getAgentStreamFromTemplate = (params: {
     userInputId,
     cacheDebugCorrelation,
     onCacheDebugProviderRequestBuilt,
+    onCacheDebugUsageReceived,
 
     onCostCalculated,
     sendAction,

packages/agent-runtime/src/run-agent-step.ts

@@ -23,6 +23,7 @@ import { getAgentOutput } from './util/agent-output'
 import {
   createCacheDebugSnapshot,
   enrichCacheDebugSnapshotWithProviderRequest,
+  enrichCacheDebugSnapshotWithUsage,
 } from './util/cache-debug'
 import {
   withSystemInstructionTags,
@@ -39,7 +40,7 @@ import type {
   FinishAgentRunFn,
   StartAgentRunFn,
 } from '@codebuff/common/types/contracts/database'
-import type { PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
+import type { CacheDebugUsageData, PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
   ParamsExcluding,
@@ -312,6 +313,17 @@ export const runAgentStep = async (
         }
       : undefined
 
+  const onCacheDebugUsageReceived =
+    cacheDebugCorrelation
+      ? (usage: CacheDebugUsageData) => {
+          enrichCacheDebugSnapshotWithUsage({
+            correlation: cacheDebugCorrelation,
+            usage,
+            logger,
+          })
+        }
+      : undefined
+
   logger.debug(
     {
       iteration: iterationNum,
@@ -343,6 +355,7 @@ export const runAgentStep = async (
         ? serializeCacheDebugCorrelation(cacheDebugCorrelation)
         : undefined,
       onCacheDebugProviderRequestBuilt,
+      onCacheDebugUsageReceived,
     })
 
     if (result.aborted) {
@@ -399,6 +412,7 @@ export const runAgentStep = async (
     includeCacheControl: supportsCacheControl(agentTemplate.model),
     messages: [systemMessage(system), ...agentState.messageHistory],
     onCacheDebugProviderRequestBuilt,
+    onCacheDebugUsageReceived,
     template: agentTemplate,
     onCostCalculated,
   })

packages/agent-runtime/src/util/cache-debug.ts

@@ -5,6 +5,7 @@ import { dirname, join } from 'path'
 import {
   type CacheDebugCorrelation,
 } from '@codebuff/common/util/cache-debug'
+import type { CacheDebugUsageData } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProviderMetadata } from '@codebuff/common/types/messages/provider-metadata'
@@ -50,6 +51,7 @@ export type CacheDebugSnapshot = {
   toolsHash?: string
   preConversion: CacheDebugPreConversionSnapshot
   providerRequest?: CacheDebugProviderRequestSnapshot
+  usage?: CacheDebugUsageData
 }
 
 function getCacheDebugDir(projectRoot: string) {
@@ -241,6 +243,42 @@ export function createCacheDebugSnapshot(params: {
   return { snapshotId, filename, projectRoot }
 }
 
+export function enrichCacheDebugSnapshotWithUsage(params: {
+  correlation: CacheDebugCorrelation
+  usage: CacheDebugUsageData
+  logger: Logger
+}) {
+  const { correlation, usage, logger } = params
+  try {
+    const existing = loadSnapshot({
+      projectRoot: correlation.projectRoot,
+      filename: correlation.filename,
+    })
+    if (!existing) {
+      logger.warn(
+        `[Cache Debug] Could not find snapshot ${correlation.filename} to enrich with usage`,
+      )
+      return
+    }
+
+    if (existing.id !== correlation.snapshotId) {
+      logger.warn(
+        `[Cache Debug] Snapshot ID mismatch while enriching ${correlation.filename} with usage`,
+      )
+      return
+    }
+
+    const updated: CacheDebugSnapshot = {
+      ...existing,
+      usage,
+    }
+
+    writeSnapshot({ snapshot: updated, logger })
+  } catch (err) {
+    logger.warn({ error: err }, '[Cache Debug] Failed to enrich snapshot with usage')
+  }
+}
+
 export function enrichCacheDebugSnapshotWithProviderRequest(params: {
   correlation: CacheDebugCorrelation
   provider: string

scripts/compare-cache-debug.ts

@@ -54,6 +54,12 @@ interface Snapshot {
     rawBody: unknown
     normalized: unknown
   }
+  usage?: {
+    inputTokens: number
+    outputTokens: number
+    cachedInputTokens: number
+    totalTokens: number
+  }
 }
 
 function findFirstDifference(
@@ -261,6 +267,13 @@ function comparePair(prev: Snapshot, curr: Snapshot, prevFile: string, currFile:
   if (prev.systemHash || curr.systemHash) {
     console.log(`  Hashes: system=${prev.systemHash ?? '?'}→${curr.systemHash ?? '?'}  tools=${prev.toolsHash ?? '?'}→${curr.toolsHash ?? '?'}`)
   }
+  for (const snap of [{ label: 'A', data: prev }, { label: 'B', data: curr }]) {
+    if (snap.data.usage) {
+      const u = snap.data.usage
+      const hitRate = u.inputTokens > 0 ? ((u.cachedInputTokens / u.inputTokens) * 100).toFixed(1) : '0.0'
+      console.log(`  Usage ${snap.label}: ${u.inputTokens} in, ${u.outputTokens} out, ${u.cachedInputTokens} cached (${hitRate}% cache hit)`)
+    }
+  }
   if (prev.runId !== curr.runId) {
     console.log(`  ⚠️  Different runs: ${prev.runId ?? '?'} → ${curr.runId ?? '?'}`)
   }

sdk/src/impl/llm.ts

@@ -216,6 +216,30 @@ function emitCacheDebugProviderRequest(params: {
   })
 }
 
+function emitCacheDebugUsage(params: {
+  callback?: (usage: {
+    inputTokens: number
+    outputTokens: number
+    cachedInputTokens: number
+    totalTokens: number
+  }) => void
+  usage: {
+    inputTokens?: number
+    outputTokens?: number
+    totalTokens?: number
+    cachedInputTokens?: number
+  }
+}) {
+  if (!params.callback) return
+
+  params.callback({
+    inputTokens: params.usage.inputTokens ?? 0,
+    outputTokens: params.usage.outputTokens ?? 0,
+    cachedInputTokens: params.usage.cachedInputTokens ?? 0,
+    totalTokens: params.usage.totalTokens ?? 0,
+  })
+}
+
 export async function* promptAiSdkStream(
   params: ParamsOf<PromptAiSdkStreamFn> & {
     skipClaudeOAuth?: boolean
@@ -587,6 +611,12 @@ export async function* promptAiSdkStream(
     rawBody: requestMetadata.body,
   })
 
+  const usageResult = await response.usage
+  emitCacheDebugUsage({
+    callback: params.onCacheDebugUsageReceived,
+    usage: usageResult,
+  })
+
   // Skip cost tracking for Claude OAuth (user is on their own subscription)
   if (!isClaudeOAuth) {
     const providerMetadataResult = await response.providerMetadata
@@ -654,6 +684,10 @@ export async function promptAiSdk(
     provider: getModelProvider(aiSDKModel),
     rawBody: response.request?.body,
   })
+  emitCacheDebugUsage({
+    callback: params.onCacheDebugUsageReceived,
+    usage: response.usage,
+  })
   const content = response.text
 
   const providerMetadata = response.providerMetadata ?? {}
@@ -719,6 +753,10 @@ export async function promptAiSdkStructured<T>(
     provider: getModelProvider(aiSDKModel),
     rawBody: response.request?.body,
   })
+  emitCacheDebugUsage({
+    callback: params.onCacheDebugUsageReceived,
+    usage: response.usage,
+  })
 
   const content = response.object