fix(ai): preserve Anthropic cache-write TTL breakdowns

.changeset/pretty-corners-peel.md

@@ -0,0 +1,5 @@
+---
+'@posthog/ai': patch
+---
+
+Preserve Anthropic cache creation TTL breakdowns in streaming and LangChain generation events.

packages/ai/src/anthropic/index.ts

@@ -95,7 +95,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
           cacheReadInputTokens: 0,
           webSearchCount: 0,
         }
-        let lastRawUsage: unknown
+        let rawUsage: Record<string, unknown> = {}
         if (Symbol.asyncIterator in value) {
           const [stream1, stream2] = monitoredStreamTee<RawMessageStreamEvent, Stream<RawMessageStreamEvent>>(
             value as Stream<RawMessageStreamEvent>,
@@ -192,14 +192,17 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                 }
 
                 if (chunk.type == 'message_start') {
-                  lastRawUsage = chunk.message.usage
+                  rawUsage = { ...chunk.message.usage }
                   usage.inputTokens = chunk.message.usage.input_tokens ?? 0
                   usage.cacheCreationInputTokens = chunk.message.usage.cache_creation_input_tokens ?? 0
                   usage.cacheReadInputTokens = chunk.message.usage.cache_read_input_tokens ?? 0
                   usage.webSearchCount = chunk.message.usage.server_tool_use?.web_search_requests ?? 0
                 }
                 if ('usage' in chunk) {
-                  lastRawUsage = chunk.usage
+                  rawUsage = {
+                    ...rawUsage,
+                    ...Object.fromEntries(Object.entries(chunk.usage).filter(([, value]) => value != null)),
+                  }
                   usage.outputTokens = chunk.usage.output_tokens ?? 0
                   // Update web search count if present in delta
                   if (chunk.usage.server_tool_use?.web_search_requests !== undefined) {
@@ -214,7 +217,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                   }
                 }
               }
-              usage.rawUsage = lastRawUsage
+              usage.rawUsage = rawUsage
 
               const latency = (Date.now() - startTime) / 1000
               const timeToFirstToken = firstTokenTime !== undefined ? (firstTokenTime - startTime) / 1000 : undefined

packages/ai/src/langchain/callbacks.ts

@@ -12,6 +12,7 @@ import { BaseMessage } from '@langchain/core/messages'
 import { sanitizeLangChain } from '../sanitization'
 import { stringifyError } from '../serializeError'
 import { warnIfPostHogAiGateway } from '../gatewayWarning'
+import { isObject } from '../typeGuards'
 
 // Mirror LangGraph's isGraphBubbleUp guard without adding LangGraph as a dependency. Every
 // LangGraph control-flow exception (GraphInterrupt, NodeInterrupt, ParentCommand, GraphDrained,
@@ -530,6 +531,13 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
       if (additionalTokenData.cacheWriteInputTokens) {
         eventProperties['$ai_cache_creation_input_tokens'] = additionalTokenData.cacheWriteInputTokens
       }
+      if (
+        additionalTokenData.cacheWrite5mInputTokens !== undefined &&
+        additionalTokenData.cacheWrite1hInputTokens !== undefined
+      ) {
+        eventProperties['$ai_cache_creation_5m_input_tokens'] = additionalTokenData.cacheWrite5mInputTokens
+        eventProperties['$ai_cache_creation_1h_input_tokens'] = additionalTokenData.cacheWrite1hInputTokens
+      }
       if (additionalTokenData.reasoningTokens) {
         eventProperties['$ai_reasoning_tokens'] = additionalTokenData.reasoningTokens
       }
@@ -702,7 +710,86 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
     return stopReason != null ? String(stopReason) : undefined
   }
 
-  private _parseUsageModel(usage: any, provider?: string, model?: string): [number, number, Record<string, any>] {
+  private _extractCacheCreationTtlBreakdown(
+    cacheCreation: unknown,
+    aggregateValues: unknown[]
+  ): [number, number] | undefined {
+    if (!isObject(cacheCreation)) {
+      return undefined
+    }
+
+    const { ephemeral_5m_input_tokens: cache5m, ephemeral_1h_input_tokens: cache1h } = cacheCreation
+    const providedValues = [cache5m, cache1h].filter((value) => value != null)
+    if (
+      providedValues.length === 0 ||
+      !providedValues.every((value) => typeof value === 'number' && Number.isFinite(value) && value >= 0)
+    ) {
+      return undefined
+    }
+
+    const breakdown: [number, number] = [
+      typeof cache5m === 'number' ? cache5m : 0,
+      typeof cache1h === 'number' ? cache1h : 0,
+    ]
+    const total = breakdown[0] + breakdown[1]
+    const validAggregates = aggregateValues.filter(
+      (value): value is number => typeof value === 'number' && Number.isFinite(value) && value >= 0
+    )
+    return total > 0 && !validAggregates.some((aggregate) => aggregate !== total) ? breakdown : undefined
+  }
+
+  private _extractBedrockCacheCreationTtlBreakdown(
+    cacheDetails: unknown,
+    aggregateValues: unknown[]
+  ): [number, number] | undefined {
+    if (!Array.isArray(cacheDetails)) {
+      return undefined
+    }
+
+    let cache5m = 0
+    let cache1h = 0
+
+    for (const detail of cacheDetails) {
+      if (!isObject(detail)) {
+        continue
+      }
+
+      const ttl = typeof detail.ttl === 'string' ? detail.ttl.toLowerCase() : undefined
+      const inputTokens = detail.inputTokens
+      if (
+        (ttl !== '5m' && ttl !== 't5m' && ttl !== '1h' && ttl !== 't1h') ||
+        typeof inputTokens !== 'number' ||
+        !Number.isFinite(inputTokens) ||
+        inputTokens < 0
+      ) {
+        continue
+      }
+
+      if (ttl === '5m' || ttl === 't5m') {
+        cache5m += inputTokens
+      } else {
+        cache1h += inputTokens
+      }
+    }
+
+    const total = cache5m + cache1h
+    const validAggregates = aggregateValues.filter(
+      (value): value is number => typeof value === 'number' && Number.isFinite(value) && value >= 0
+    )
+    if (total === 0 || validAggregates.some((aggregate) => aggregate !== total)) {
+      return undefined
+    }
+
+    return [cache5m, cache1h]
+  }
+
+  private _parseUsageModel(
+    usage: any,
+    provider?: string,
+    model?: string,
+    inputIncludesCacheTokens = true,
+    rawUsage?: any
+  ): [number, number, Record<string, any>] {
     const conversionList: Array<[string, 'input' | 'output']> = [
       ['promptTokens', 'input'],
       ['completionTokens', 'output'],
@@ -751,6 +838,33 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
       additionalTokenData.cacheWriteInputTokens = usage.input_token_details.cache_creation
     }
 
+    const directCacheCreationAggregates = [
+      usage.cache_creation_input_tokens,
+      usage.input_token_details?.cache_creation,
+      usage.cacheWriteInputTokens,
+      rawUsage?.cache_creation_input_tokens,
+      rawUsage?.input_token_details?.cache_creation,
+      rawUsage?.cacheWriteInputTokens,
+      additionalTokenData.cacheWriteInputTokens,
+    ]
+    const cacheCreationTtl =
+      this._extractCacheCreationTtlBreakdown(usage.cache_creation, directCacheCreationAggregates) ??
+      this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation, directCacheCreationAggregates) ??
+      this._extractBedrockCacheCreationTtlBreakdown(usage.cacheDetails, [
+        usage.cacheWriteInputTokens,
+        additionalTokenData.cacheWriteInputTokens,
+      ]) ??
+      this._extractBedrockCacheCreationTtlBreakdown(rawUsage?.cacheDetails, [
+        rawUsage?.cacheWriteInputTokens,
+        additionalTokenData.cacheWriteInputTokens,
+      ])
+    if (cacheCreationTtl) {
+      const [cacheWrite5mInputTokens, cacheWrite1hInputTokens] = cacheCreationTtl
+      additionalTokenData.cacheWrite5mInputTokens = cacheWrite5mInputTokens
+      additionalTokenData.cacheWrite1hInputTokens = cacheWrite1hInputTokens
+      additionalTokenData.cacheWriteInputTokens = cacheWrite5mInputTokens + cacheWrite1hInputTokens
+    }
+
     // Check for reasoning tokens in various formats
     if (usage.completion_tokens_details?.reasoning_tokens != null) {
       additionalTokenData.reasoningTokens = usage.completion_tokens_details.reasoning_tokens
@@ -816,7 +930,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
       isAnthropic = true
     }
 
-    if (isAnthropic && parsedUsage.input) {
+    if (isAnthropic && inputIncludesCacheTokens && parsedUsage.input) {
       const cacheTokens =
         (additionalTokenData.cacheReadInputTokens || 0) + (additionalTokenData.cacheWriteInputTokens || 0)
       if (cacheTokens > 0) {
@@ -828,39 +942,86 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
   }
 
   private parseUsage(response: LLMResult, provider?: string, model?: string): [number, number, Record<string, any>] {
-    let llmUsage: [number, number, Record<string, any>] = [0, 0, {}]
+    const isNonEmptyUsage = (usage: unknown): usage is Record<string, any> =>
+      isObject(usage) && Object.keys(usage).length > 0
+    const firstNonEmptyUsage = (...candidates: unknown[]): Record<string, any> | undefined =>
+      candidates.find(isNonEmptyUsage)
+
+    let normalizedGenerationUsage: any
+    let rawGenerationUsage: any
+    let fallbackGenerationUsage: any
+
+    for (const generation of response.generations ?? []) {
+      for (const genChunk of generation) {
+        const generationInfo = genChunk.generationInfo ?? {}
+        const message = 'message' in genChunk ? genChunk.message : undefined
+        const messageUsage =
+          message && typeof message === 'object' && 'usage_metadata' in message ? message.usage_metadata : undefined
+        normalizedGenerationUsage = firstNonEmptyUsage(
+          normalizedGenerationUsage,
+          messageUsage,
+          generationInfo.usage_metadata
+        )
+
+        const messageResponseMetadata =
+          message &&
+          typeof message === 'object' &&
+          'response_metadata' in message &&
+          isObject(message.response_metadata)
+            ? message.response_metadata
+            : undefined
+        const generationResponseMetadata = isObject(generationInfo.response_metadata)
+          ? generationInfo.response_metadata
+          : undefined
+        const messageStreamMetadata = isObject(messageResponseMetadata?.metadata)
+          ? messageResponseMetadata.metadata
+          : undefined
+        const generationStreamMetadata = isObject(generationResponseMetadata?.metadata)
+          ? generationResponseMetadata.metadata
+          : undefined
+        rawGenerationUsage = firstNonEmptyUsage(
+          rawGenerationUsage,
+          messageResponseMetadata?.usage,
+          messageStreamMetadata?.usage,
+          generationResponseMetadata?.usage,
+          generationStreamMetadata?.usage
+        )
+        fallbackGenerationUsage = firstNonEmptyUsage(
+          fallbackGenerationUsage,
+          messageResponseMetadata?.['amazon-bedrock-invocationMetrics'],
+          generationResponseMetadata?.['amazon-bedrock-invocationMetrics'],
+          generationInfo.usage_metadata
+        )
+      }
+    }
+
+    const isAnthropic = provider?.toLowerCase() === 'anthropic' || model?.toLowerCase().includes('anthropic') === true
+    if (isAnthropic && isNonEmptyUsage(normalizedGenerationUsage)) {
+      return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage)
+    }
+
     const llmUsageKeys = ['token_usage', 'usage', 'tokenUsage']
 
     if (response.llmOutput != null) {
-      const key = llmUsageKeys.find((k) => response.llmOutput?.[k] != null)
-      if (key) {
-        llmUsage = this._parseUsageModel(response.llmOutput[key], provider, model)
+      for (const key of llmUsageKeys) {
+        const llmUsage = response.llmOutput[key]
+        if (!isNonEmptyUsage(llmUsage)) {
+          continue
+        }
+        return this._parseUsageModel(llmUsage, provider, model, key !== 'usage', llmUsage)
       }
     }
 
-    // If top-level usage info was not found, try checking the generations.
-    if (llmUsage[0] === 0 && llmUsage[1] === 0 && response.generations) {
-      for (const generation of response.generations) {
-        for (const genChunk of generation) {
-          const message = (genChunk as any).message ?? {}
-          const messageResponseMetadata = message.response_metadata ?? {}
-          const generationInfo = genChunk.generationInfo ?? {}
-          const generationResponseMetadata = generationInfo.response_metadata ?? {}
-          const chunkUsage =
-            message.usage_metadata ??
-            messageResponseMetadata['usage'] ??
-            messageResponseMetadata['amazon-bedrock-invocationMetrics'] ??
-            generationInfo.usage_metadata ??
-            generationResponseMetadata['usage'] ??
-            generationResponseMetadata['amazon-bedrock-invocationMetrics']
-          if (chunkUsage) {
-            llmUsage = this._parseUsageModel(chunkUsage, provider, model)
-            return llmUsage
-          }
-        }
-      }
+    if (isNonEmptyUsage(normalizedGenerationUsage)) {
+      return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage)
+    }
+    if (isNonEmptyUsage(rawGenerationUsage)) {
+      return this._parseUsageModel(rawGenerationUsage, provider, model, false, rawGenerationUsage)
+    }
+    if (isNonEmptyUsage(fallbackGenerationUsage)) {
+      return this._parseUsageModel(fallbackGenerationUsage, provider, model)
     }
 
-    return llmUsage
+    return [0, 0, {}]
   }
 }

packages/ai/tests/anthropic.test.ts

@@ -18,6 +18,10 @@ interface MockAnthropicResponseOptions {
     output_tokens: number
     cache_creation_input_tokens?: number
     cache_read_input_tokens?: number
+    cache_creation?: {
+      ephemeral_5m_input_tokens: number
+      ephemeral_1h_input_tokens: number
+    }
     server_tool_use?: {
       web_search_requests?: number
     }
@@ -160,6 +164,7 @@ const createMockStreamChunks = (options: MockAnthropicResponseOptions = {}): Moc
         input_tokens: options.usage?.input_tokens || 20,
         cache_creation_input_tokens: options.usage?.cache_creation_input_tokens || 0,
         cache_read_input_tokens: options.usage?.cache_read_input_tokens || 0,
+        ...(options.usage?.cache_creation ? { cache_creation: options.usage.cache_creation } : {}),
         output_tokens: 0,
         ...(options.usage?.server_tool_use?.web_search_requests
           ? {
@@ -804,6 +809,51 @@ describe('PostHogAnthropic', () => {
         cacheReadInputTokens: 5,
       })
     })
+
+    it('should retain cache creation TTL usage after the final streaming delta', async () => {
+      mockStreamChunks = createMockStreamChunks({
+        content: 'Streaming response',
+        usage: {
+          input_tokens: 75,
+          output_tokens: 40,
+          cache_creation_input_tokens: 30,
+          cache_read_input_tokens: 5,
+          cache_creation: {
+            ephemeral_5m_input_tokens: 10,
+            ephemeral_1h_input_tokens: 20,
+          },
+        },
+      })
+
+      const stream = await client.messages.create({
+        model: 'claude-3-opus-20240229',
+        messages: [{ role: 'user', content: 'Hello' }],
+        max_tokens: 100,
+        stream: true,
+        posthogDistinctId: 'test-user-123',
+      })
+
+      for await (const _chunk of stream) {
+        // Consume the stream so the monitoring branch reaches message_delta.
+      }
+      await waitForAsyncCapture()
+
+      const captureMock = mockPostHogClient.capture as jest.Mock
+      const [captureArgs] = captureMock.mock.calls
+
+      expect(captureArgs[0].properties['$ai_usage']).toEqual({
+        input_tokens: 75,
+        output_tokens: 40,
+        cache_creation_input_tokens: 30,
+        cache_read_input_tokens: 5,
+        cache_creation: {
+          ephemeral_5m_input_tokens: 10,
+          ephemeral_1h_input_tokens: 20,
+        },
+      })
+      expect(captureArgs[0].properties['$ai_cache_creation_input_tokens']).toBe(30)
+      expect(captureArgs[0].properties['$ai_output_tokens']).toBe(40)
+    })
   })
 
   describe('Telemetry failure isolation', () => {