diff --git a/packages/server-utils/src/ai/vercel-ai/index.ts b/packages/server-utils/src/ai/vercel-ai/index.ts index ec491bd79e6c..76a355fdc588 100644 --- a/packages/server-utils/src/ai/vercel-ai/index.ts +++ b/packages/server-utils/src/ai/vercel-ai/index.ts @@ -92,7 +92,12 @@ export function getProviderMetadataAttributes(providerMetadata: unknown): Record * They must not be written onto a span that reports usage aggregated across steps * (`gen_ai.invoke_agent`), where they would replace the aggregate with one step's figures. */ -export const LAST_STEP_ONLY_USAGE_KEYS = new Set([GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS]); +export const LAST_STEP_ONLY_USAGE_KEYS = new Set([ + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, + GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, +]); /** * Sets an attribute only if the value is not null or undefined. diff --git a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts index 15f526c1b86d..924cde857b6b 100644 --- a/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts +++ b/packages/server-utils/src/integrations/vercel-ai/vercel-ai-dc-subscriber.ts @@ -18,6 +18,8 @@ import { GEN_AI_TOOL_DEFINITIONS, GEN_AI_TOOL_DESCRIPTION, GEN_AI_TOOL_NAME, + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, + GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, @@ -356,6 +358,9 @@ function enrichInvokeAgentFromStream( addTokensToSpan(span, GEN_AI_USAGE_INPUT_TOKENS, input); addTokensToSpan(span, GEN_AI_USAGE_OUTPUT_TOKENS, output); addTokensToSpan(span, GEN_AI_USAGE_TOTAL_TOKENS, tokenCount(usage.totalTokens) ?? sum(input, output)); + const { cacheRead, cacheWrite } = cacheTokens(usage); + addTokensToSpan(span, GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cacheRead); + addTokensToSpan(span, GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cacheWrite); } if (recordOutputs) { @@ -567,6 +572,7 @@ export function enrichSpanOnEnd( if (totalTokens !== undefined) { span.setAttribute(GEN_AI_USAGE_TOTAL_TOKENS, totalTokens); } + setCacheTokens(span, usage); } // Match the OTel integration: finish reasons live on the model-call (`generate_content`) span, not @@ -633,6 +639,32 @@ function tokenCount(value: unknown): number | undefined { return asNumber(value) ?? (isObjectLike(value) ? asNumber(value.total) : undefined); } +/** + * Cache token counts as the AI SDK normalizes them: v5 `cachedInputTokens`, v6 `inputTokenDetails`, + * v7 `inputTokens.{cacheRead,cacheWrite}`. + */ +function setCacheTokens(span: Span, usage: Record): void { + const { cacheRead, cacheWrite } = cacheTokens(usage); + if (cacheRead !== undefined) { + span.setAttribute(GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cacheRead); + } + if (cacheWrite !== undefined) { + span.setAttribute(GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cacheWrite); + } +} + +function cacheTokens(usage: Record): { cacheRead?: number; cacheWrite?: number } { + const inputTokens = isObjectLike(usage.inputTokens) ? usage.inputTokens : undefined; + const inputTokenDetails = isObjectLike(usage.inputTokenDetails) ? usage.inputTokenDetails : undefined; + return { + cacheRead: + asNumber(inputTokens?.cacheRead) ?? + asNumber(inputTokenDetails?.cacheReadTokens) ?? + asNumber(usage.cachedInputTokens), + cacheWrite: asNumber(inputTokens?.cacheWrite) ?? asNumber(inputTokenDetails?.cacheWriteTokens), + }; +} + function buildOutputMessages( parts: Array>, finishReason: string | undefined, diff --git a/packages/server-utils/test/integrations/vercel-ai/cache-tokens.test.ts b/packages/server-utils/test/integrations/vercel-ai/cache-tokens.test.ts new file mode 100644 index 000000000000..54f249be0591 --- /dev/null +++ b/packages/server-utils/test/integrations/vercel-ai/cache-tokens.test.ts @@ -0,0 +1,188 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, + GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, +} from '@sentry/conventions/attributes'; +import { getMainCarrier, setCurrentClient, spanToStaticSpanJSON } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { + createSpanFromMessage, + enrichSpanOnEnd, + streamedResultToChannelResult, +} from '../../../src/integrations/vercel-ai/vercel-ai-dc-subscriber'; +import { getDefaultTestClientOptions, TestClient } from '../../mocks/client'; + +describe('Vercel AI SDK cache tokens', () => { + beforeEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + afterEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + function setupClient(): Span[] { + const client = new TestClient( + getDefaultTestClientOptions({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + tracesSampleRate: 1, + }), + ); + setCurrentClient(client); + client.init(); + + const endedSpans: Span[] = []; + client.on('spanEnd', span => endedSpans.push(span)); + return endedSpans; + } + + function runSpan( + type: string, + result: Record, + existingAttributes: Record = {}, + ): Record { + const endedSpans = setupClient(); + const message = { type, event: {}, result } as Parameters[0]; + const span = createSpanFromMessage(message, {} as Parameters[1]); + span!.setAttributes(existingAttributes); + enrichSpanOnEnd(span!, message, {} as Parameters[2]); + span?.end(); + return spanToStaticSpanJSON(endedSpans[0]!).data ?? {}; + } + + it('reads v5 `cachedInputTokens` when providerMetadata has no provider key', () => { + const data = runSpan('languageModelCall', { + usage: { + inputTokens: 120, + outputTokens: 10, + totalTokens: 130, + cachedInputTokens: 100, + }, + providerMetadata: { gateway: { routing: {} } }, + }); + + expect(data[GEN_AI_USAGE_INPUT_TOKENS]).toBe(120); + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(100); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBeUndefined(); + }); + + it('reads v6 `inputTokenDetails` cache read and write counts', () => { + const data = runSpan('languageModelCall', { + usage: { + inputTokens: 120, + inputTokenDetails: { + noCacheTokens: 20, + cacheReadTokens: 80, + cacheWriteTokens: 20, + }, + outputTokens: 10, + totalTokens: 130, + }, + }); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(80); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(20); + }); + + it('prefers `inputTokenDetails` over the deprecated `cachedInputTokens`', () => { + const data = runSpan('languageModelCall', { + usage: { + inputTokens: 120, + inputTokenDetails: { cacheReadTokens: 80 }, + cachedInputTokens: 5, + outputTokens: 10, + }, + }); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(80); + }); + + it('reads v7 `inputTokens` / `outputTokens` objects', () => { + const data = runSpan('languageModelCall', { + usage: { + inputTokens: { total: 120, noCache: 20, cacheRead: 80, cacheWrite: 20 }, + outputTokens: { total: 10, text: 10, reasoning: 0 }, + }, + }); + + expect(data[GEN_AI_USAGE_INPUT_TOKENS]).toBe(120); + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(10); + expect(data[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(130); + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(80); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(20); + }); + + it('sets nothing when the usage object carries no cache counts', () => { + const data = runSpan('languageModelCall', { + usage: { inputTokens: 120, outputTokens: 10, totalTokens: 130 }, + }); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBeUndefined(); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBeUndefined(); + }); + + it('leaves an existing count in place when the SDK usage does not report it', () => { + const data = runSpan( + 'languageModelCall', + { usage: { inputTokens: 120, cachedInputTokens: 80, outputTokens: 10 } }, + { [GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]: 20 }, + ); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(80); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(20); + }); + + it('keeps providerMetadata-derived counts over the SDK usage counts', () => { + const data = runSpan('languageModelCall', { + usage: { + inputTokens: 120, + inputTokenDetails: { cacheReadTokens: 80, cacheWriteTokens: 20 }, + cachedInputTokens: 80, + outputTokens: 10, + }, + providerMetadata: { + anthropic: { cacheReadInputTokens: 81, cacheCreationInputTokens: 21 }, + }, + }); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(81); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(21); + }); + + it('keeps the aggregated SDK counts on a root operation over last-step providerMetadata', () => { + const data = runSpan('generateText', { + usage: { + inputTokens: 9500, + inputTokenDetails: { cacheReadTokens: 8000, cacheWriteTokens: 500 }, + outputTokens: 40, + }, + providerMetadata: { + anthropic: { cacheReadInputTokens: 3000, cacheCreationInputTokens: 0 }, + }, + }); + + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(8000); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(500); + }); + + it('reads the counts from a streamed model call result', () => { + const data = runSpan( + 'languageModelCall', + streamedResultToChannelResult({ + text: 'hi', + toolCalls: [], + usage: { + inputTokens: { total: 120, cacheRead: 80, cacheWrite: 20 }, + outputTokens: { total: 10 }, + }, + }), + ); + + expect(data[GEN_AI_USAGE_INPUT_TOKENS]).toBe(120); + expect(data[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]).toBe(80); + expect(data[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]).toBe(20); + }); +});