Skip to content

Google ADK incompletely normalizes Gemini token usage #2305

Description

Summary

Google ADK usage metadata follows Google GenAI's token accounting, but the plugin currently performs only a partial normalization. It records visible candidate tokens as completion usage, excludes internal thinking tokens from completion, excludes tool-use prompt tokens from prompt usage, and ignores supported modality breakdowns.

Google's current usage contract is:

totalTokenCount =
  promptTokenCount
  + toolUsePromptTokenCount
  + candidatesTokenCount
  + thoughtsTokenCount

This underreports normalized prompt/completion usage and estimated cost for thinking-capable Gemini responses and provider-managed tool calls.

Reference: https://ai.google.dev/api/generate-content#UsageMetadata

Current Behavior

populateUsageMetrics() in js/src/instrumentation/plugins/google-adk-plugin.ts maps:

metrics.prompt_tokens = usage.promptTokenCount;
metrics.completion_tokens = usage.candidatesTokenCount;
metrics.completion_reasoning_tokens = usage.thoughtsTokenCount;
metrics.tokens = usage.totalTokenCount;

The ADK event aggregator already sums candidate and thought counts independently before calling this helper, but the final completion metric does not combine them. The vendored ADK usage type and aggregator also omit toolUsePromptTokenCount and modality detail arrays.

Expected Behavior

Normalize aggregate usage as:

prompt_tokens = promptTokenCount + toolUsePromptTokenCount
completion_tokens = candidatesTokenCount + thoughtsTokenCount
completion_reasoning_tokens = thoughtsTokenCount
prompt_cached_tokens = cachedContentTokenCount
tokens = totalTokenCount

Reasoning remains a detail metric and is also included in completion usage. Tool-use prompt tokens belong to normalized prompt usage.

Where Google reports modality details, populate existing standard Braintrust metrics:

promptTokensDetails[AUDIO]     -> prompt_audio_tokens
candidatesTokensDetails[AUDIO] -> completion_audio_tokens
candidatesTokensDetails[IMAGE] -> completion_image_tokens

Other modality details should remain available in provider-native output. Do not invent new metric names without updating the instrumentation spec.

Acceptance Criteria

  • Fold aggregated thoughtsTokenCount into completion_tokens without removing completion_reasoning_tokens.
  • Aggregate and fold toolUsePromptTokenCount into prompt_tokens.
  • Preserve totalTokenCount as tokens and cachedContentTokenCount as prompt_cached_tokens.
  • Extend vendored ADK usage types and aggregation for tool-use prompt counts and modality detail arrays.
  • Populate supported audio/image detail metrics when reported.
  • Preserve missing and zero values correctly without fabricating absent usage.
  • Add focused unit coverage for candidate-only, thought-only, tool-use-only, combined, and modality-detail usage.
  • Add cassette-backed ADK coverage using thinking-capable and provider-managed-tool Gemini responses, including streaming where supported.
  • Assert tokens = prompt_tokens + completion_tokens when all aggregate values are reported.

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions