fix(ai): declare gemini cache reporting as inclusive on generations

.changeset/gemini-cache-reporting-exclusive.md

@@ -0,0 +1,5 @@
+---
+'@posthog/ai': patch
+---
+
+fix(gemini): declare Gemini's cache accounting model on generations with cache reads, so ingestion prices cached tokens from `$ai_cache_reporting_exclusive` instead of inferring it from the token counts

packages/ai/src/captureAiGeneration.ts

@@ -4,7 +4,7 @@ import { uuidv7, ErrorTracking as CoreErrorTracking, toJsonSafeValue } from '@po
 import { version } from '../package.json'
 import type { TokenUsage } from './types'
 import { stringifyError } from './serializeError'
-import { AIEvent, CostOverride, getTokensSource, withPrivacyMode } from './utils'
+import { AIEvent, CostOverride, getTokensSource, hasTokenOverrides, withPrivacyMode } from './utils'
 import { warnIfPostHogAiGateway } from './gatewayWarning'
 
 /**
@@ -161,10 +161,26 @@ export const captureAiGeneration = async (client: PostHog, options: CaptureAiGen
       }
     }
 
+    // The caller's own token counts override the SDK-derived ones further down, via the
+    // `options.properties` spread.
+    const tokensOverridden = hasTokenOverrides(options.properties)
+
     const additionalTokenValues = {
       ...(usage.reasoningTokens ? { $ai_reasoning_tokens: usage.reasoningTokens } : {}),
       ...(usage.cacheReadInputTokens ? { $ai_cache_read_input_tokens: usage.cacheReadInputTokens } : {}),
       ...(usage.cacheCreationInputTokens ? { $ai_cache_creation_input_tokens: usage.cacheCreationInputTokens } : {}),
+      // Checked against undefined rather than truthiness, because false is the meaningful
+      // value here and a truthiness guard would drop it.
+      //
+      // Dropped entirely when the caller overrides the token counts: the flag describes how
+      // the SDK-derived counts relate to each other, so against passthrough counts it can be
+      // wrong in the expensive direction. Declaring inclusive over counts that are actually
+      // exclusive makes ingestion subtract the cache pool that was never in the input. A
+      // caller who knows their own accounting model can still pass
+      // `$ai_cache_reporting_exclusive` themselves, and that value wins.
+      ...(usage.cacheReportingExclusive !== undefined && !tokensOverridden
+        ? { $ai_cache_reporting_exclusive: usage.cacheReportingExclusive }
+        : {}),
       ...(usage.webSearchCount ? { $ai_web_search_count: usage.webSearchCount } : {}),
       ...(usage.rawUsage ? { $ai_usage: usage.rawUsage } : {}),
     }

packages/ai/src/gemini/index.ts

@@ -80,6 +80,11 @@ export class WrappedModels {
             (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number })?.thoughtsTokenCount ??
             0,
           cacheReadInputTokens: metadata?.cachedContentTokenCount ?? 0,
+          // Gemini counts cachedContentTokenCount inside promptTokenCount, so declare the
+          // accounting model rather than leaving ingestion to infer it. Under explicit
+          // context caching the two counts come from separate measurements and can disagree
+          // by a few percent, which makes inference from the counts alone unreliable.
+          ...(metadata?.cachedContentTokenCount ? { cacheReportingExclusive: false } : {}),
           webSearchCount: calculateGoogleWebSearchCount(response),
           rawUsage: metadata,
         },
@@ -196,6 +201,9 @@ export class WrappedModels {
               (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number }).thoughtsTokenCount ??
               0,
             cacheReadInputTokens: metadata.cachedContentTokenCount ?? 0,
+            // See the non-streaming path: Gemini counts cachedContentTokenCount inside
+            // promptTokenCount, so the accounting model is declared rather than inferred.
+            ...(metadata.cachedContentTokenCount ? { cacheReportingExclusive: false } : {}),
             webSearchCount: usage.webSearchCount,
             rawUsage: metadata,
           }

packages/ai/src/types.ts

@@ -95,6 +95,10 @@ export interface TokenUsage {
   reasoningTokens?: unknown // Use unknown since various providers return different types
   cacheReadInputTokens?: unknown // Use unknown for provider flexibility
   cacheCreationInputTokens?: unknown // Use unknown for provider flexibility
+  // Whether cache tokens are counted separately from inputTokens. Providers that report
+  // them as a subset of inputTokens set this false. Left undefined when the provider's
+  // accounting model is not known, in which case ingestion infers it from the counts.
+  cacheReportingExclusive?: boolean
   webSearchCount?: number // Count of web search queries/calls used
   rawUsage?: unknown // Raw provider usage metadata for backend processing
 }

packages/ai/src/utils.ts

@@ -28,11 +28,16 @@ const TOKEN_PROPERTY_KEYS = new Set([
   '$ai_reasoning_tokens',
 ])
 
+/**
+ * Whether the caller supplied their own token counts, which override the ones the SDK
+ * derived from the provider response.
+ */
+export function hasTokenOverrides(posthogProperties?: Record<string, unknown>): boolean {
+  return !!posthogProperties && Object.keys(posthogProperties).some((key) => TOKEN_PROPERTY_KEYS.has(key))
+}
+
 export function getTokensSource(posthogProperties?: Record<string, unknown>): string {
-  if (posthogProperties && Object.keys(posthogProperties).some((key) => TOKEN_PROPERTY_KEYS.has(key))) {
-    return 'passthrough'
-  }
-  return 'sdk'
+  return hasTokenOverrides(posthogProperties) ? 'passthrough' : 'sdk'
 }
 
 // limit large outputs by truncating to 200kb (approx 200k bytes)

packages/ai/tests/gemini.test.ts

@@ -675,6 +675,96 @@ describe('PostHogGemini - Jest test suite', () => {
     ])
   })
 
+  describe('Cache reporting', () => {
+    test('declares inclusive cache reporting when cached tokens are present', async () => {
+      mockGeminiResponse = {
+        text: 'Cached answer',
+        candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }],
+        usageMetadata: {
+          promptTokenCount: 23000,
+          candidatesTokenCount: 8,
+          cachedContentTokenCount: 25000,
+        },
+      }
+      ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse)
+
+      await client.models.generateContent({
+        model: 'gemini-2.0-flash-001',
+        contents: 'Test',
+        posthogDistinctId: 'test-id',
+      })
+
+      const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0]
+      expect(properties['$ai_cache_read_input_tokens']).toBe(25000)
+      expect(properties['$ai_cache_reporting_exclusive']).toBe(false)
+    })
+
+    test('drops the cache reporting flag when the caller overrides token counts', async () => {
+      mockGeminiResponse = {
+        text: 'Cached answer',
+        candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }],
+        usageMetadata: {
+          promptTokenCount: 23000,
+          candidatesTokenCount: 8,
+          cachedContentTokenCount: 25000,
+        },
+      }
+      ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse)
+
+      await client.models.generateContent({
+        model: 'gemini-2.0-flash-001',
+        contents: 'Test',
+        posthogDistinctId: 'test-id',
+        posthogProperties: { $ai_input_tokens: 400, $ai_cache_read_input_tokens: 25000 },
+      })
+
+      const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0]
+      expect(properties['$ai_tokens_source']).toBe('passthrough')
+      expect(properties['$ai_input_tokens']).toBe(400)
+      expect(properties).not.toHaveProperty('$ai_cache_reporting_exclusive')
+    })
+
+    test('keeps an explicit cache reporting flag from the caller', async () => {
+      mockGeminiResponse = {
+        text: 'Cached answer',
+        candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }],
+        usageMetadata: {
+          promptTokenCount: 23000,
+          candidatesTokenCount: 8,
+          cachedContentTokenCount: 25000,
+        },
+      }
+      ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse)
+
+      await client.models.generateContent({
+        model: 'gemini-2.0-flash-001',
+        contents: 'Test',
+        posthogDistinctId: 'test-id',
+        posthogProperties: {
+          $ai_input_tokens: 400,
+          $ai_cache_read_input_tokens: 25000,
+          $ai_cache_reporting_exclusive: true,
+        },
+      })
+
+      const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0]
+      expect(properties['$ai_cache_reporting_exclusive']).toBe(true)
+    })
+
+    test('omits the cache reporting flag when no tokens were cached', async () => {
+      ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse)
+
+      await client.models.generateContent({
+        model: 'gemini-2.0-flash-001',
+        contents: 'Test',
+        posthogDistinctId: 'test-id',
+      })
+
+      const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0]
+      expect(properties).not.toHaveProperty('$ai_cache_reporting_exclusive')
+    })
+  })
+
   describe('Web Search Tracking', () => {
     test('should detect grounding metadata (binary detection)', async () => {
       mockGeminiResponse = {