diff --git a/packages/core/src/shared-exports.ts b/packages/core/src/shared-exports.ts index 4c1100142fec..3ed79e5347aa 100644 --- a/packages/core/src/shared-exports.ts +++ b/packages/core/src/shared-exports.ts @@ -178,7 +178,7 @@ export { export * as metrics from './metrics/public-api'; export type { MetricOptions } from './metrics/public-api'; export { createConsolaReporter } from './integrations/consola'; -export { addVercelAiProcessors, getProviderMetadataAttributes } from './tracing/vercel-ai'; +export { addVercelAiProcessors, getProviderMetadataAttributes, LAST_STEP_ONLY_USAGE_KEYS } from './tracing/vercel-ai'; export { getTruncatedJsonString, shouldEnableTruncation, resolveAIRecordingOptions } from './tracing/ai/utils'; export { GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, diff --git a/packages/core/src/tracing/vercel-ai/index.ts b/packages/core/src/tracing/vercel-ai/index.ts index f9abcb096dc8..c6f296b77819 100644 --- a/packages/core/src/tracing/vercel-ai/index.ts +++ b/packages/core/src/tracing/vercel-ai/index.ts @@ -601,14 +601,50 @@ 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 even if a future SDK version 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 reasoning-inclusive 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_ATTRIBUTE] = + (googleUsage.candidatesTokenCount ?? 0) + googleUsage.thoughtsTokenCount; + setAttributeIfDefined(attributes, GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, googleUsage.totalTokenCount); + setAttributeIfDefined(attributes, 'gen_ai.usage.reasoning.output_tokens', googleUsage.thoughtsTokenCount); + } + return attributes; } +/** + * Usage attributes that `getProviderMetadataAttributes` derives 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. Exported so the channel/orchestrion subscribers, which call + * `getProviderMetadataAttributes` directly rather than through `addProviderMetadataToAttributes`, + * apply the same rule. + */ +export const LAST_STEP_ONLY_USAGE_KEYS = new Set([ + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, + 'gen_ai.usage.reasoning.output_tokens', +]); + function addProviderMetadataToAttributes(attributes: Record): void { const providerMetadata = attributes[AI_RESPONSE_PROVIDER_METADATA_ATTRIBUTE] as string | undefined; if (!providerMetadata) { return; } + // An `invoke_agent` span carries the summed `ai.usage.*` of every step, while `providerMetadata` + // describes the last step alone. Writing output or total from it would replace the aggregate with + // one step's figures, so a two-step call reports output 50 against input 900. The event-processor + // path happens to overwrite it again in `applyAccumulatedTokens`; the streamed path ships it. + // Reasoning goes with them: it is a subset of an output the parent never recomputes, and the + // accumulator never sums it, so the last step's count would stand in for the whole call. + const lastStepOnly = attributes[GEN_AI_OPERATION_NAME_ATTRIBUTE] === 'invoke_agent'; try { const derived = getProviderMetadataAttributes(JSON.parse(providerMetadata) as ProviderMetadata); for (const [key, value] of Object.entries(derived)) { @@ -616,6 +652,9 @@ function addProviderMetadataToAttributes(attributes: Record): v if (key === GEN_AI_CONVERSATION_ID_ATTRIBUTE && attributes[key]) { continue; } + if (lastStepOnly && LAST_STEP_ONLY_USAGE_KEYS.has(key)) { + continue; + } attributes[key] = value; } } catch { diff --git a/packages/core/src/tracing/vercel-ai/vercel-ai-attributes.ts b/packages/core/src/tracing/vercel-ai/vercel-ai-attributes.ts index 62d89f50c17c..63547b1f058c 100644 --- a/packages/core/src/tracing/vercel-ai/vercel-ai-attributes.ts +++ b/packages/core/src/tracing/vercel-ai/vercel-ai-attributes.ts @@ -446,6 +446,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/core/test/lib/tracing/vercel-ai-reasoning-tokens.test.ts b/packages/core/test/lib/tracing/vercel-ai-reasoning-tokens.test.ts new file mode 100644 index 000000000000..90d8baa1582e --- /dev/null +++ b/packages/core/test/lib/tracing/vercel-ai-reasoning-tokens.test.ts @@ -0,0 +1,216 @@ +import { describe, expect, it } from 'vitest'; +import { addVercelAiProcessors } from '../../../src/tracing/vercel-ai'; +import type { SpanJSON } from '../../../src/types/span'; +import { getDefaultTestClientOptions, TestClient } from '../../mocks/client'; + +/** + * Real usage from a Gemini reasoning model: the candidate output is one token and the model spent + * the rest of its budget on hidden reasoning ("thoughts"). The AI SDK reports the candidate count + * as `outputTokens` and exposes the reasoning count only through + * `providerMetadata.google.usageMetadata`, so a span built from `ai.usage.*` alone undercounts + * output by 100 and undercounts the total by the same. + */ +const GEMINI_REASONING_METADATA = { + google: { + groundingMetadata: null, + safetyRatings: null, + usageMetadata: { + promptTokenCount: 14, + candidatesTokenCount: 1, + thoughtsTokenCount: 100, + totalTokenCount: 115, + }, + }, +}; + +function processSpans(spans: SpanJSON[]): SpanJSON[] { + const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 }); + const client = new TestClient(options); + client.init(); + addVercelAiProcessors(client); + + const eventProcessor = client['_eventProcessors'].find(processor => processor.id === 'VercelAiEventProcessor'); + expect(eventProcessor).toBeDefined(); + + return eventProcessor!({ type: 'transaction' as const, spans }, {})!.spans!; +} + +function span(description: string, data: SpanJSON['data'], spanId = 'span-1', parentSpanId?: string): SpanJSON { + return { + description, + span_id: spanId, + parent_span_id: parentSpanId, + trace_id: 'test-trace-id', + start_timestamp: 1000, + timestamp: 2000, + origin: 'auto.vercelai.otel', + data, + }; +} + +describe('vercel-ai Gemini reasoning tokens', () => { + it('adds reasoning tokens into output and total on a model span', () => { + const [processed] = processSpans([ + span('ai.generateText.doGenerate', { + 'ai.usage.promptTokens': 14, + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }), + ]); + + expect(processed?.data?.['gen_ai.usage.input_tokens']).toBe(14); + // Output covers the 100 reasoning tokens, not just the single candidate token. + expect(processed?.data?.['gen_ai.usage.output_tokens']).toBe(101); + expect(processed?.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(100); + // The real Gemini total, not input plus candidate-only output, which would be 15. + expect(processed?.data?.['gen_ai.usage.total_tokens']).toBe(115); + }); + + it('reads the v6 vertex key the same way', () => { + const [processed] = processSpans([ + span('ai.generateText.doGenerate', { + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify({ vertex: GEMINI_REASONING_METADATA.google }), + }), + ]); + + expect(processed?.data?.['gen_ai.usage.output_tokens']).toBe(101); + expect(processed?.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(100); + }); + + it('counts an absent candidatesTokenCount as zero', () => { + // Gemini omits `candidatesTokenCount` when the response is truncated during thinking, which + // means no candidate tokens were produced. Treating it as zero keeps the reasoning tokens in + // the output rather than dropping the whole recompute and reporting the candidate-only count. + const [processed] = processSpans([ + span('ai.generateText.doGenerate', { + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify({ + google: { usageMetadata: { promptTokenCount: 14, thoughtsTokenCount: 100, totalTokenCount: 115 } }, + }), + }), + ]); + + expect(processed?.data?.['gen_ai.usage.output_tokens']).toBe(100); + expect(processed?.data?.['gen_ai.usage.total_tokens']).toBe(115); + expect(processed?.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(100); + }); + + it('does not overwrite an invoke_agent parent aggregate from the last step', () => { + // The parent carries the summed usage of both steps; providerMetadata describes step two only. + // Writing output or total from it would report step two's figures against the summed input. + const [parent] = processSpans([ + span('ai.generateText', { + 'operation.name': 'ai.generateText', + 'ai.usage.promptTokens': 900, + 'ai.usage.completionTokens': 350, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }), + ]); + + expect(parent?.data?.['gen_ai.operation.name']).toBe('invoke_agent'); + expect(parent?.data?.['gen_ai.usage.input_tokens']).toBe(900); + expect(parent?.data?.['gen_ai.usage.output_tokens']).toBe(350); + expect(parent?.data?.['gen_ai.usage.total_tokens']).toBe(1250); + // Reasoning is a subset of an output this span never recomputes, so it is left off entirely + // rather than reporting the last step's count against the summed output. + expect(parent?.data?.['gen_ai.usage.reasoning.output_tokens']).toBeUndefined(); + }); + + it('keeps a multi-step call consistent: the parent sums, each step reports its own reasoning', () => { + const stepOne = { + google: { + usageMetadata: { + promptTokenCount: 400, + candidatesTokenCount: 20, + thoughtsTokenCount: 80, + totalTokenCount: 500, + }, + }, + }; + const processed = processSpans([ + span( + 'ai.generateText', + { + 'operation.name': 'ai.generateText', + 'ai.usage.promptTokens': 900, + 'ai.usage.completionTokens': 350, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }, + 'parent', + ), + span( + 'ai.generateText.doGenerate', + { + 'ai.usage.promptTokens': 400, + 'ai.usage.completionTokens': 20, + 'ai.response.providerMetadata': JSON.stringify(stepOne), + }, + 'step-1', + 'parent', + ), + span( + 'ai.generateText.doGenerate', + { + 'ai.usage.promptTokens': 500, + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }, + 'step-2', + 'parent', + ), + ]); + + const [parent, first, second] = processed; + + // Each step gets its own reasoning-inclusive figures. + expect(first?.data?.['gen_ai.usage.output_tokens']).toBe(100); + expect(first?.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(80); + expect(second?.data?.['gen_ai.usage.output_tokens']).toBe(101); + expect(second?.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(100); + + // The parent keeps an aggregate whose parts add up, rather than step two's 101 and 115. + expect(parent?.data?.['gen_ai.usage.input_tokens']).toBe(900); + expect(parent?.data?.['gen_ai.usage.output_tokens']).toBe(350); + expect(parent?.data?.['gen_ai.usage.total_tokens']).toBe(1250); + // The last step's 100 would stand in for the call's real 180, and nothing sums it, so it is + // left off rather than reported. + expect(parent?.data?.['gen_ai.usage.reasoning.output_tokens']).toBeUndefined(); + }); + + it('holds on the streamed path, which has no accumulation pass to repair the parent', () => { + // The event processor sees the whole transaction and re-derives an `invoke_agent` parent from + // its children; `processSpan` sees one span at a time and ships whatever it produced. This is + // the path the gate actually protects. + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1.0 })); + client.init(); + addVercelAiProcessors(client); + + const streamed = (attrs: Record): Record => { + const span = { span_id: 's', trace_id: 't', attributes: { 'sentry.origin': 'auto.vercelai.otel', ...attrs } }; + client.emit('processSpan', span as never); + return span.attributes; + }; + + const parent = streamed({ + 'operation.name': 'ai.generateText', + 'ai.usage.inputTokens': 14, + 'ai.usage.outputTokens': 1, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }); + const child = streamed({ + 'operation.name': 'ai.generateText.doGenerate', + 'ai.usage.promptTokens': 14, + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }); + + expect(child['gen_ai.usage.output_tokens']).toBe(101); + expect(child['gen_ai.usage.total_tokens']).toBe(115); + expect(child['gen_ai.usage.reasoning.output_tokens']).toBe(100); + + // No reasoning count larger than the output it is meant to be a subset of. + expect(parent['gen_ai.usage.output_tokens']).toBe(1); + expect(parent['gen_ai.usage.reasoning.output_tokens']).toBeUndefined(); + }); +}); diff --git a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts index fead1c746383..b9c8e0ad1798 100644 --- a/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts +++ b/packages/server-utils/src/vercel-ai/vercel-ai-dc-subscriber.ts @@ -33,6 +33,7 @@ import { GEN_AI_CONVERSATION_ID_ATTRIBUTE, GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, + LAST_STEP_ONLY_USAGE_KEYS, getClient, getProviderMetadataAttributes, getTruncatedJsonString, @@ -130,6 +131,20 @@ export function clearOperationCallId(callId: string): void { invokeAgentSpanByCallId.delete(callId); } +/** + * `providerMetadata` is last-step only; drop derived usage on spans that report an aggregate. + * Matches `addProviderMetadataToAttributes`. + */ +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)) { @@ -576,6 +591,7 @@ export function enrichSpanOnEnd( // oxlint-disable-next-line typescript/no-dynamic-delete delete providerAttributes[GEN_AI_CONVERSATION_ID_ATTRIBUTE]; } + dropLastStepOnlyUsage(providerAttributes, type); span.setAttributes(providerAttributes); if (recordOutputs) {