fix(server-utils): Read AI SDK cache token counts in Vercel AI channel subscriber

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<string>([GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS]);
+export const LAST_STEP_ONLY_USAGE_KEYS = new Set<string>([
+  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.

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<string, unknown>): 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<string, unknown>): { 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<Record<string, unknown>>,
   finishReason: string | undefined,

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<string, unknown>,
+    existingAttributes: Record<string, number> = {},
+  ): Record<string, unknown> {
+    const endedSpans = setupClient();
+    const message = { type, event: {}, result } as Parameters<typeof createSpanFromMessage>[0];
+    const span = createSpanFromMessage(message, {} as Parameters<typeof createSpanFromMessage>[1]);
+    span!.setAttributes(existingAttributes);
+    enrichSpanOnEnd(span!, message, {} as Parameters<typeof enrichSpanOnEnd>[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);
+  });
+});