diff --git a/packages/server-utils/src/ai/vercel-ai/index.ts b/packages/server-utils/src/ai/vercel-ai/index.ts index 758defb019b5..ec491bd79e6c 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,30 @@ 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]); + /** * 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..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'; @@ -47,7 +48,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 +135,29 @@ export function clearOperationCallId(callId: string): void { invokeAgentSpanByCallId.delete(callId); } +/** + * `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. */ function recordToolDescriptions(callId: string | undefined, tools: unknown): void { if (!callId || !Array.isArray(tools)) { @@ -573,6 +597,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..51040686b90e --- /dev/null +++ b/packages/server-utils/test/integrations/vercel-ai/gemini-reasoning-tokens.test.ts @@ -0,0 +1,130 @@ +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(); + }); + + // 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); + }); +});