Skip to content

fix: log full-turn token usage from model_usage in claude-agent-sdk wrapper#565

Open
lsuarez-ekline wants to merge 1 commit into
braintrustdata:mainfrom
lsuarez-ekline:fix/claude-agent-sdk-model-usage-tokens
Open

fix: log full-turn token usage from model_usage in claude-agent-sdk wrapper#565
lsuarez-ekline wants to merge 1 commit into
braintrustdata:mainfrom
lsuarez-ekline:fix/claude-agent-sdk-model-usage-tokens

Conversation

@lsuarez-ekline

Copy link
Copy Markdown

What

ContextTracker._handle_result logs the turn's token metrics from ResultMessage.usage, which is the cumulative usage for the orchestrator agent only. On a turn that spawns subagents (the Task/Agent tool), the subagent tokens are billed but missing from the trace, so the logged totals under-count — often heavily.

This prefers the token counts from ResultMessage.model_usage (the per-model breakdown, which is per-agent and includes subagent calls), summed across models, for the logged metrics. Metadata (service tier, inference geo) still comes from usage. When model_usage is missing or unusable it falls back to usage, so single-agent turns and older SDKs are unchanged.

Fixes #564.

Details

  • New helper _aggregate_model_usage() sums model_usage into an Anthropic-usage-shaped dict, then reuses extract_anthropic_usage for the metric mapping — so cache-read / cache-creation / completion all flow through the same path as before.
  • Single-agent turns are unaffected (usage == model_usage).
  • A model that only ever runs as a subagent has no dedicated LLM span; its tokens now land on the turn's LLM span via the aggregate, so the trace's token total matches the SDK's model_usage.

Testing

  • nox -s "test_claude_agent_sdk(latest)" — 50 passed (existing suite + 2 new unit tests for _aggregate_model_usage).
  • ruff format / ruff check clean.

Notes

Alternative considered: attribute tokens onto per-subagent spans instead of aggregating onto the turn's LLM span. Happy to take that direction if you'd prefer — this PR takes the minimal approach that makes the trace-level token total correct.

…rapper

ResultMessage.usage covers only the orchestrator agent, so turns that spawn
subagents under-count tokens on the trace. Prefer the per-model model_usage
breakdown (which includes subagent calls) for the logged token metrics;
metadata (service tier, etc.) still comes from usage.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Claude Agent SDK: _handle_result logs main-agent usage, under-counts turns with subagents

1 participant