From 3dfe7cd407a1e74ee68b758f8da1532debfc9338 Mon Sep 17 00:00:00 2001 From: coryslater <25396141+fivestarspicy@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:58:28 -0700 Subject: [PATCH 1/2] fix(ai): declare gemini cache reporting as inclusive on generations Gemini counts `cachedContentTokenCount` inside `promptTokenCount`, but the SDK never said so, leaving ingestion to infer the accounting model from the token counts alone. That inference is unreliable here. Under explicit context caching the two counts come from separate measurements, the cache object at creation time and the prompt per request, so they can disagree by a few percent and the cache pool can land just above the input total. Set `cacheReportingExclusive` to false on generations that report cache reads, on both the streaming and non-streaming paths, and map it onto `$ai_cache_reporting_exclusive`. The property mapping checks against undefined rather than truthiness, because false is the meaningful value here. Mirrors PostHog/posthog-python#860. Generated-By: PostHog Code Task-Id: 06160e48-feb9-4d39-8b7b-3dcfd1d9ca24 --- .../gemini-cache-reporting-exclusive.md | 5 +++ packages/ai/src/captureAiGeneration.ts | 5 +++ packages/ai/src/gemini/index.ts | 8 ++++ packages/ai/src/types.ts | 4 ++ packages/ai/tests/gemini.test.ts | 38 +++++++++++++++++++ 5 files changed, 60 insertions(+) create mode 100644 .changeset/gemini-cache-reporting-exclusive.md diff --git a/.changeset/gemini-cache-reporting-exclusive.md b/.changeset/gemini-cache-reporting-exclusive.md new file mode 100644 index 0000000000..231ae16c4d --- /dev/null +++ b/.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 diff --git a/packages/ai/src/captureAiGeneration.ts b/packages/ai/src/captureAiGeneration.ts index 57e0bca0da..b33d2b25b3 100644 --- a/packages/ai/src/captureAiGeneration.ts +++ b/packages/ai/src/captureAiGeneration.ts @@ -165,6 +165,11 @@ export const captureAiGeneration = async (client: PostHog, options: CaptureAiGen ...(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. + ...(usage.cacheReportingExclusive !== undefined + ? { $ai_cache_reporting_exclusive: usage.cacheReportingExclusive } + : {}), ...(usage.webSearchCount ? { $ai_web_search_count: usage.webSearchCount } : {}), ...(usage.rawUsage ? { $ai_usage: usage.rawUsage } : {}), } diff --git a/packages/ai/src/gemini/index.ts b/packages/ai/src/gemini/index.ts index 12ef5783fc..06877a9aca 100644 --- a/packages/ai/src/gemini/index.ts +++ b/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, } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 33b8bd98bf..a6ef560c47 100644 --- a/packages/ai/src/types.ts +++ b/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 } diff --git a/packages/ai/tests/gemini.test.ts b/packages/ai/tests/gemini.test.ts index b028050552..62d27f69b4 100644 --- a/packages/ai/tests/gemini.test.ts +++ b/packages/ai/tests/gemini.test.ts @@ -675,6 +675,44 @@ 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('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 = { From cf73d3e663cb928007786c6e9a12ac1d7d06e095 Mon Sep 17 00:00:00 2001 From: coryslater <25396141+fivestarspicy@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:08:20 -0700 Subject: [PATCH 2/2] fix(ai): drop derived cache reporting flag when token counts are overridden The derived flag describes how the SDK's own input and cache counts relate to each other. When a caller passes their own token counts through `posthogProperties`, those override the derived counts but the flag survived alongside them, so it could describe numbers that are no longer on the event. That is wrong in the expensive direction: declaring inclusive over counts that are actually exclusive makes ingestion subtract a cache pool that was never part of the input. Suppress the derived flag whenever the caller overrides any token count. A caller who knows their own accounting model can still pass `$ai_cache_reporting_exclusive` explicitly, and that value wins. Extracts the passthrough check behind `hasTokenOverrides`, which `getTokensSource` now shares. Generated-By: PostHog Code Task-Id: 06160e48-feb9-4d39-8b7b-3dcfd1d9ca24 --- packages/ai/src/captureAiGeneration.ts | 15 +++++++- packages/ai/src/utils.ts | 13 +++++-- packages/ai/tests/gemini.test.ts | 52 ++++++++++++++++++++++++++ 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/packages/ai/src/captureAiGeneration.ts b/packages/ai/src/captureAiGeneration.ts index b33d2b25b3..e11d88c8e1 100644 --- a/packages/ai/src/captureAiGeneration.ts +++ b/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,13 +161,24 @@ 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. - ...(usage.cacheReportingExclusive !== undefined + // + // 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 } : {}), diff --git a/packages/ai/src/utils.ts b/packages/ai/src/utils.ts index 9c4a54f53d..cfbe4a6b3a 100644 --- a/packages/ai/src/utils.ts +++ b/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): boolean { + return !!posthogProperties && Object.keys(posthogProperties).some((key) => TOKEN_PROPERTY_KEYS.has(key)) +} + export function getTokensSource(posthogProperties?: Record): 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) diff --git a/packages/ai/tests/gemini.test.ts b/packages/ai/tests/gemini.test.ts index 62d27f69b4..f5484e46c1 100644 --- a/packages/ai/tests/gemini.test.ts +++ b/packages/ai/tests/gemini.test.ts @@ -699,6 +699,58 @@ describe('PostHogGemini - Jest test suite', () => { 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)