fix(server-utils): Count Google GenAI tool-use prompt tokens as input tokens

packages/server-utils/src/ai/google-genai/index.ts

@@ -39,6 +39,7 @@ import type { ContentListUnion, Message, PartListUnion } from './utils';
 import {
   candidatesToMessageParts,
   contentUnionToMessages,
+  getInputTokens,
   setOutputMessagesAttribute,
   systemInstructionToText,
 } from './utils';
@@ -206,7 +207,7 @@ export function addResponseAttributes(span: Span, response: GoogleGenAIResponse,
     const hasOutput = typeof usage.candidatesTokenCount === 'number' || typeof usage.thoughtsTokenCount === 'number';
     span.setAttributes(
       getTokenUsageAttributes(
-        usage.promptTokenCount,
+        getInputTokens(usage.promptTokenCount, usage.toolUsePromptTokenCount),
         hasOutput ? (usage.candidatesTokenCount ?? 0) + (usage.thoughtsTokenCount ?? 0) : undefined,
         undefined,
         usage.cachedContentTokenCount,

packages/server-utils/src/ai/google-genai/streaming.ts

@@ -3,7 +3,7 @@ import type { Span } from '@sentry/core';
 import { endStreamSpan } from '../core/utils';
 import type { GoogleGenAIResponse } from './types';
 import type { MessagePart } from './utils';
-import { candidatesToMessageParts, setOutputMessagesAttribute } from './utils';
+import { candidatesToMessageParts, getInputTokens, setOutputMessagesAttribute } from './utils';
 
 /**
  * State object used to accumulate information from a stream of Google GenAI events.
@@ -17,8 +17,10 @@ interface StreamingState {
   responseId?: string;
   /** The model name. */
   responseModel?: string;
-  /** Number of prompt/input tokens used. */
+  /** Number of prompt/input tokens used, including the tool-use prompt tokens. */
   promptTokens?: number;
+  promptTokenCount?: number;
+  toolUsePromptTokenCount?: number;
   /** Number of completion/output tokens used. */
   completionTokens?: number;
   candidateTokens?: number;
@@ -62,7 +64,11 @@ function handleResponseMetadata(chunk: GoogleGenAIResponse, state: StreamingStat
 
   const usage = chunk.usageMetadata;
   if (usage) {
-    if (typeof usage.promptTokenCount === 'number') state.promptTokens = usage.promptTokenCount;
+    if (typeof usage.promptTokenCount === 'number') state.promptTokenCount = usage.promptTokenCount;
+    if (typeof usage.toolUsePromptTokenCount === 'number') {
+      state.toolUsePromptTokenCount = usage.toolUsePromptTokenCount;
+    }
+    state.promptTokens = getInputTokens(state.promptTokenCount, state.toolUsePromptTokenCount);
     if (typeof usage.candidatesTokenCount === 'number') state.candidateTokens = usage.candidatesTokenCount;
     if (typeof usage.thoughtsTokenCount === 'number') state.reasoningOutputTokens = usage.thoughtsTokenCount;
     if (state.candidateTokens !== undefined || state.reasoningOutputTokens !== undefined) {

packages/server-utils/src/ai/google-genai/utils.ts

@@ -205,3 +205,16 @@ export function systemInstructionToText(systemInstruction: unknown): string | un
 
   return texts.length ? texts.join('\n') : undefined;
 }
+
+/**
+ * Input tokens of a Gemini response. `toolUsePromptTokenCount` (results of built-in tools such as Google Search,
+ * URL context or code execution, fed back to the model) is reported apart from `promptTokenCount`, but
+ * `totalTokenCount` includes it, so it is part of the input.
+ */
+export function getInputTokens(promptTokenCount?: number, toolUsePromptTokenCount?: number): number | undefined {
+  if (typeof promptTokenCount !== 'number' && typeof toolUsePromptTokenCount !== 'number') return undefined;
+  return (
+    (typeof promptTokenCount === 'number' ? promptTokenCount : 0) +
+    (typeof toolUsePromptTokenCount === 'number' ? toolUsePromptTokenCount : 0)
+  );
+}

packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts

@@ -1,6 +1,11 @@
 import { describe, expect, it } from 'vitest';
 import type { Span } from '@sentry/core';
-import { GEN_AI_RESPONSE_TOOL_CALLS } from '@sentry/conventions/attributes';
+import {
+  GEN_AI_RESPONSE_TOOL_CALLS,
+  GEN_AI_USAGE_INPUT_TOKENS,
+  GEN_AI_USAGE_OUTPUT_TOKENS,
+  GEN_AI_USAGE_TOTAL_TOKENS,
+} from '@sentry/conventions/attributes';
 import { addResponseAttributes } from '../../../../src/ai/google-genai/index';
 import { instrumentStream } from '../../../../src/ai/google-genai/streaming';
 import type { ContentPart, GoogleGenAIResponse } from '../../../../src/ai/google-genai/types';
@@ -126,3 +131,56 @@ describe('instrumentStream (google-genai tool calls)', () => {
     expect(toolCalls).toEqual([functionCall]);
   });
 });
+
+describe('google-genai tool-use prompt tokens', () => {
+  // Google Search, URL context and code execution results are fed back to the model and counted in
+  // `toolUsePromptTokenCount`, apart from `promptTokenCount`; `totalTokenCount` includes both.
+  const usageMetadata = {
+    promptTokenCount: 100,
+    toolUsePromptTokenCount: 40,
+    candidatesTokenCount: 20,
+    totalTokenCount: 160,
+  };
+
+  it('counts them as input tokens on a response', () => {
+    const { span, attributes } = createMockSpan();
+
+    addResponseAttributes(span, { usageMetadata } as GoogleGenAIResponse, false);
+
+    expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(140);
+    expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(20);
+    expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(160);
+  });
+
+  it('leaves the input tokens alone when there are no tool-use prompt tokens', () => {
+    const { span, attributes } = createMockSpan();
+
+    addResponseAttributes(
+      span,
+      {
+        usageMetadata: { promptTokenCount: 100, candidatesTokenCount: 20, totalTokenCount: 120 },
+      } as GoogleGenAIResponse,
+      false,
+    );
+
+    expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(100);
+    expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(120);
+  });
+
+  it('counts them as input tokens on a stream, also when a later chunk only repeats one of the counts', async () => {
+    const { span, attributes } = createMockSpan();
+
+    const chunks = [
+      chunkWithParts([{ text: 'Hi' }], { usageMetadata: { promptTokenCount: 100, toolUsePromptTokenCount: 40 } }),
+      chunkWithParts([{ text: ' there' }], {
+        usageMetadata: { candidatesTokenCount: 20, totalTokenCount: 160 },
+      }),
+    ];
+
+    await drain(instrumentStream(streamOf(chunks), span, false));
+
+    expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(140);
+    expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(20);
+    expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(160);
+  });
+});