From 2cdacaf2f5af92ba827ef6422b69ce5027e66bd5 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 4 Sep 2026 10:13:21 +0200 Subject: [PATCH 1/2] fix(server-utils): Include Gemini reasoning tokens in Vercel AI token usage Forward port of #23433, which landed on `v10` because the Vercel AI OTel span processing it originally targeted was removed here in #23384. Gemini reports its reasoning ("thoughts") tokens separately from the candidate output count, so on `ai` v4/v5 the SDK's `outputTokens` covers only the visible answer and the reasoning count reaches us solely through `providerMetadata.google.usageMetadata`. The conventions define `gen_ai.usage.output_tokens` as reasoning-inclusive, so those spans were under-reporting rather than merely missing a breakdown. Read the `google`/`vertex` `usageMetadata` and derive output from `candidatesTokenCount + thoughtsTokenCount`, the total from `totalTokenCount`, and the reasoning breakdown alongside them. An absent `candidatesTokenCount` counts as zero: Gemini omits it when the response is truncated during thinking, which means no candidate tokens were produced, and skipping the recompute there would report zero output for a call that spent its whole budget reasoning. None of the three land on `gen_ai.invoke_agent` spans. Those carry usage summed across every step while `providerMetadata` describes the last step alone, so writing from it would replace an aggregate with one step's figures. Reasoning is gated with them because it is a subset of an output that span never recomputes, and nothing sums it across steps. Unlike `v10` this branch has a single caller, so the rule lives directly in `enrichSpanOnEnd` and there is no `addProviderMetadataToAttributes` equivalent. Ref #23993 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_018TR1cvQA7t6T2saCrHwwUh --- .../server-utils/src/ai/vercel-ai/index.ts | 27 ++++ .../src/ai/vercel-ai/vercel-ai-attributes.ts | 14 +++ .../vercel-ai/vercel-ai-dc-subscriber.ts | 16 ++- .../vercel-ai/gemini-reasoning-tokens.test.ts | 118 ++++++++++++++++++ 4 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts diff --git a/packages/server-utils/src/ai/vercel-ai/index.ts b/packages/server-utils/src/ai/vercel-ai/index.ts index 758defb019b5..9f047dcba7f8 100644 --- a/packages/server-utils/src/ai/vercel-ai/index.ts +++ b/packages/server-utils/src/ai/vercel-ai/index.ts @@ -3,7 +3,9 @@ import { GEN_AI_CONVERSATION_ID, GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, } from '@sentry/conventions/attributes'; import type { OpenAiProviderMetadata, ProviderMetadata } from './vercel-ai-attributes'; @@ -68,9 +70,34 @@ export function getProviderMetadataAttributes(providerMetadata: unknown): Record setAttributeIfDefined(attributes, 'gen_ai.usage.input_tokens.cache_miss', metadata.deepseek.promptCacheMissTokens); } + // Google (v5 uses 'google', v6 Vertex AI uses 'vertex'). Gemini reports its reasoning ("thoughts") + // tokens separately from the candidate output count, so the SDK's `outputTokens` covers only the + // visible answer. Recompute output from `candidatesTokenCount + thoughtsTokenCount` rather than + // adding onto the existing value, which stays correct on `ai` v6+ where `outputTokens` already + // folds reasoning in. `candidatesTokenCount` is omitted when the response is truncated during + // thinking, which means no candidate tokens, so it counts as zero. Reasoning is a subset of + // output per the conventions, so it is only written alongside the output it belongs to. + const googleUsage = (metadata.google ?? metadata.vertex)?.usageMetadata; + if (googleUsage && typeof googleUsage.thoughtsTokenCount === 'number' && googleUsage.thoughtsTokenCount > 0) { + attributes[GEN_AI_USAGE_OUTPUT_TOKENS] = (googleUsage.candidatesTokenCount ?? 0) + googleUsage.thoughtsTokenCount; + setAttributeIfDefined(attributes, GEN_AI_USAGE_TOTAL_TOKENS, googleUsage.totalTokenCount); + setAttributeIfDefined(attributes, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, googleUsage.thoughtsTokenCount); + } + return attributes; } +/** + * Usage attributes derived from `providerMetadata`, which describes only the last step of a call. + * 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, + GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, +]); + /** * Sets an attribute only if the value is not null or undefined. */ diff --git a/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts b/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts index 5370a62261f9..4c6538742c2a 100644 --- a/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts +++ b/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts @@ -181,6 +181,20 @@ export interface GoogleGenerativeAIProviderMetadata { * @see https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters */ safetyRatings?: null | unknown; + + /** + * Raw token usage returned by the Gemini API. Reasoning ("thoughts") tokens are reported here in + * `thoughtsTokenCount`, separately from the candidate output count, so they have to be added back + * into `gen_ai.usage.output_tokens`. + * @see https://ai.google.dev/api/generate-content#UsageMetadata + * @see https://github.com/vercel/ai/blob/main/packages/google/src/google-language-model.ts + */ + usageMetadata?: null | { + promptTokenCount?: number; + candidatesTokenCount?: number; + thoughtsTokenCount?: number; + totalTokenCount?: number; + }; } /** 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 b90fd705168e..3cfebd61d2b3 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 @@ -47,7 +47,7 @@ import { import type { TracingChannel } from 'node:diagnostics_channel'; import { GEN_AI_TOOL_CALL_ID_ATTRIBUTE } from '../../ai/core/gen-ai-attributes'; import type { GenAiOptions } from '../../ai/core/utils'; -import { getProviderMetadataAttributes } from '../../ai/vercel-ai'; +import { getProviderMetadataAttributes, LAST_STEP_ONLY_USAGE_KEYS } from '../../ai/vercel-ai'; import { WORKERS_AI_INTEGRATION_NAME } from '../../ai/workers-ai/constants'; import { bindTracingChannelToSpan } from '../../tracing-channel'; import { asNumber, asString, isReadableStream, type StreamedModelCallResult, sum, tapModelCallStream } from './util'; @@ -134,6 +134,19 @@ export function clearOperationCallId(callId: string): void { invokeAgentSpanByCallId.delete(callId); } +/** + * `providerMetadata` is last-step only; drop derived usage on spans that report an aggregate. + */ +function dropLastStepOnlyUsage(providerAttributes: Record, type: ChannelEventType): void { + if (!ROOT_OPERATION_TYPES.has(type)) { + return; + } + for (const key of LAST_STEP_ONLY_USAGE_KEYS) { + // oxlint-disable-next-line typescript/no-dynamic-delete + delete providerAttributes[key]; + } +} + /** Record tool name → description from an event's `tools`, so tool spans can backfill the description. */ function recordToolDescriptions(callId: string | undefined, tools: unknown): void { if (!callId || !Array.isArray(tools)) { @@ -573,6 +586,7 @@ export function enrichSpanOnEnd( // oxlint-disable-next-line typescript/no-dynamic-delete delete providerAttributes[GEN_AI_CONVERSATION_ID]; } + dropLastStepOnlyUsage(providerAttributes, type); span.setAttributes(providerAttributes); if (recordOutputs) { diff --git a/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts b/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts new file mode 100644 index 000000000000..df3cf4ce606a --- /dev/null +++ b/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts @@ -0,0 +1,118 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_REASONING_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 } from '../../../src/integrations/vercel-ai/vercel-ai-dc-subscriber'; +import { getDefaultTestClientOptions, TestClient } from '../../mocks/client'; + +/** + * Real usage from a Gemini reasoning model: one candidate token, the rest of the budget spent on + * hidden reasoning. `@ai-sdk/google` maps `outputTokens` from `candidatesTokenCount` alone, so the + * reasoning count reaches us only through `providerMetadata`. + */ +const REASONING_METADATA = { + google: { + usageMetadata: { promptTokenCount: 14, candidatesTokenCount: 1, thoughtsTokenCount: 100, totalTokenCount: 115 }, + }, +}; + +describe('Vercel AI Gemini reasoning 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): Record { + const endedSpans = setupClient(); + const message = { type, event: {}, result } as Parameters[0]; + const span = createSpanFromMessage(message, {} as Parameters[1]); + enrichSpanOnEnd(span!, message, {} as Parameters[2]); + span?.end(); + return spanToStaticSpanJSON(endedSpans[0]!).data ?? {}; + } + + it('adds reasoning tokens into output and total on a model-call span', () => { + const data = runSpan('languageModelCall', { + usage: { inputTokens: 14, outputTokens: 1, totalTokens: 115 }, + providerMetadata: REASONING_METADATA, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(101); + expect(data[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(115); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBe(100); + }); + + it('reads the v6 vertex key the same way', () => { + const data = runSpan('languageModelCall', { + usage: { inputTokens: 14, outputTokens: 1, totalTokens: 115 }, + providerMetadata: { vertex: REASONING_METADATA.google }, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(101); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBe(100); + }); + + // Gemini omits `candidatesTokenCount` when the response is truncated during thinking, which means + // no candidate tokens were produced. Dropping the recompute there would report zero output for a + // call that spent its whole budget reasoning. + it('counts an absent candidatesTokenCount as zero', () => { + const data = runSpan('languageModelCall', { + usage: { inputTokens: 14, outputTokens: 0, totalTokens: 514 }, + providerMetadata: { + google: { usageMetadata: { promptTokenCount: 14, thoughtsTokenCount: 500, totalTokenCount: 514 } }, + }, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(500); + expect(data[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(514); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBe(500); + }); + + it('leaves a non-reasoning response alone', () => { + const data = runSpan('languageModelCall', { + usage: { inputTokens: 14, outputTokens: 12, totalTokens: 26 }, + providerMetadata: { + google: { usageMetadata: { promptTokenCount: 14, candidatesTokenCount: 12, totalTokenCount: 26 } }, + }, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(12); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBeUndefined(); + }); + + // A top-level operation's span reports usage summed across every step, while `providerMetadata` + // describes the last step alone. Reasoning is gated with output and total: it is a subset of an + // output this span never recomputes, and nothing sums it across steps. + it('does not write last-step usage onto an invoke_agent span', () => { + const data = runSpan('generateText', { + usage: { inputTokens: 900, outputTokens: 350, totalTokens: 1250 }, + providerMetadata: REASONING_METADATA, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(350); + expect(data[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(1250); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBeUndefined(); + }); +}); From fe33688794a24e7c598f3e3bce6338162a2307b6 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 4 Sep 2026 11:00:19 +0200 Subject: [PATCH 2/2] fix(server-utils): Keep OpenAI reasoning tokens on aggregate spans Gating `gen_ai.usage.reasoning.output_tokens` unconditionally dropped it from `invoke_agent` spans for every provider, not just Google. OpenAI writes the same attribute from `providerMetadata.openai.reasoningTokens`, so a `generateText` call against a reasoning model silently lost a count it used to report. The reason to drop it only applies where a recomputed output is dropped with it: that leaves the span's own output reasoning-exclusive, so the documented subset relationship would not hold. A provider reporting reasoning against an already-inclusive output keeps it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_018TR1cvQA7t6T2saCrHwwUh --- packages/server-utils/src/ai/vercel-ai/index.ts | 6 +----- .../vercel-ai/vercel-ai-dc-subscriber.ts | 11 +++++++++++ .../vercel-ai/gemini-reasoning-tokens.test.ts | 12 ++++++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/packages/server-utils/src/ai/vercel-ai/index.ts b/packages/server-utils/src/ai/vercel-ai/index.ts index 9f047dcba7f8..ec491bd79e6c 100644 --- a/packages/server-utils/src/ai/vercel-ai/index.ts +++ b/packages/server-utils/src/ai/vercel-ai/index.ts @@ -92,11 +92,7 @@ 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, - GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, -]); +export const LAST_STEP_ONLY_USAGE_KEYS = new Set([GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_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 3cfebd61d2b3..15f526c1b86d 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 @@ -20,6 +20,7 @@ import { GEN_AI_TOOL_NAME, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS, SENTRY_OP, } from '@sentry/conventions/attributes'; @@ -136,15 +137,25 @@ export function clearOperationCallId(callId: string): void { /** * `providerMetadata` is last-step only; drop derived usage on spans that report an aggregate. + * + * Reasoning goes only when it accompanies a recomputed output (Google), because dropping that + * output leaves the span's own count reasoning-exclusive and the documented subset relationship + * would not hold. A provider that reports reasoning against an already-inclusive output (OpenAI) + * keeps it. */ function dropLastStepOnlyUsage(providerAttributes: Record, type: ChannelEventType): void { if (!ROOT_OPERATION_TYPES.has(type)) { return; } + const hasRecomputedOutput = GEN_AI_USAGE_OUTPUT_TOKENS in providerAttributes; for (const key of LAST_STEP_ONLY_USAGE_KEYS) { // oxlint-disable-next-line typescript/no-dynamic-delete delete providerAttributes[key]; } + if (hasRecomputedOutput) { + // oxlint-disable-next-line typescript/no-dynamic-delete + delete providerAttributes[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]; + } } /** Record tool name → description from an event's `tools`, so tool spans can backfill the description. */ diff --git a/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts b/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts index df3cf4ce606a..51040686b90e 100644 --- a/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts +++ b/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts @@ -115,4 +115,16 @@ describe('Vercel AI Gemini reasoning tokens', () => { expect(data[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(1250); expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBeUndefined(); }); + + // OpenAI reports reasoning against an `outputTokens` that already includes it, so there is no + // recomputed output to suppress and the subset relationship still holds on the aggregate span. + it('keeps OpenAI reasoning tokens on an invoke_agent span', () => { + const data = runSpan('generateText', { + usage: { inputTokens: 900, outputTokens: 350, totalTokens: 1250 }, + providerMetadata: { openai: { reasoningTokens: 120 } }, + }); + + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(350); + expect(data[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]).toBe(120); + }); });