fix(server-utils): Include Gemini reasoning tokens in Vercel AI token usage

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,

packages/core/src/tracing/vercel-ai/index.ts

@@ -601,21 +601,60 @@ 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<string>([
+  GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
+  'gen_ai.usage.reasoning.output_tokens',
+]);
+
 function addProviderMetadataToAttributes(attributes: Record<string, unknown>): 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)) {
       // Preserve the original behaviour of not overwriting an already-set conversation id.
       if (key === GEN_AI_CONVERSATION_ID_ATTRIBUTE && attributes[key]) {
         continue;
       }
+      if (lastStepOnly && LAST_STEP_ONLY_USAGE_KEYS.has(key)) {
+        continue;
+      }
       attributes[key] = value;
     }
   } catch {

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;
+  };
 }
 
 /**

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<string, unknown>): Record<string, unknown> => {
+      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();
+  });
+});

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<string, number | string>, 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) {