diff --git a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 1ff40395..7ea396b3 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -30,6 +30,7 @@ from braintrust.integrations.claude_agent_sdk.tracing import ( ContextTracker, ToolSpanTracker, + _aggregate_model_usage, _build_llm_input, _create_client_wrapper_class, _create_tool_wrapper_class, @@ -2899,3 +2900,33 @@ def test_context_tracker_preserves_bash_output_when_next_tool_use_arrives_before read_output = read_spans[0].get("output") assert read_output is not None assert read_output["content"] == "alpha_file_contents" + + +def test_aggregate_model_usage_sums_across_models(): + model_usage = { + "claude-opus-4-8": { + "inputTokens": 100, + "outputTokens": 20, + "cacheReadInputTokens": 5000, + "cacheCreationInputTokens": 300, + }, + "claude-haiku-4-5": { + "inputTokens": 10, + "outputTokens": 40, + "cacheReadInputTokens": 2000, + "cacheCreationInputTokens": 100, + }, + } + + assert _aggregate_model_usage(model_usage) == { + "input_tokens": 110, + "output_tokens": 60, + "cache_read_input_tokens": 7000, + "cache_creation_input_tokens": 400, + } + + +def test_aggregate_model_usage_returns_none_when_missing_or_empty(): + assert _aggregate_model_usage(None) is None + assert _aggregate_model_usage({}) is None + assert _aggregate_model_usage({"claude-opus-4-8": {"inputTokens": 0, "outputTokens": 0}}) is None diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index 631b6f99..98a1fdbf 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -7,7 +7,7 @@ from collections.abc import AsyncGenerator, AsyncIterable from typing import Any -from braintrust.integrations.anthropic._utils import Wrapper, extract_anthropic_usage +from braintrust.integrations.anthropic._utils import Wrapper, _try_to_dict, extract_anthropic_usage from braintrust.integrations.claude_agent_sdk._constants import ( ANTHROPIC_MESSAGES_CREATE_SPAN_NAME, CLAUDE_AGENT_RUN_FAILED_ERROR, @@ -25,11 +25,46 @@ ) from braintrust.logger import start_span from braintrust.span_types import SpanTypeAttribute +from braintrust.util import is_numeric _thread_local = threading.local() +_MODEL_USAGE_TOKEN_FIELDS = ( + ("inputTokens", "input_tokens"), + ("outputTokens", "output_tokens"), + ("cacheReadInputTokens", "cache_read_input_tokens"), + ("cacheCreationInputTokens", "cache_creation_input_tokens"), +) + + +def _aggregate_model_usage(model_usage: Any) -> dict[str, int] | None: + """Sum ``ResultMessage.model_usage`` into an Anthropic-usage-shaped dict. + + ``model_usage`` is the SDK's per-model token breakdown. Unlike + ``ResultMessage.usage`` — which covers only the orchestrator agent — it + includes subagent calls, so summing it recovers the whole turn's usage. + Returns ``None`` when the breakdown is missing or unusable so callers can + fall back to ``usage``. + """ + usage_by_model = _try_to_dict(model_usage) + if not usage_by_model: + return None + + totals = {target: 0 for _, target in _MODEL_USAGE_TOKEN_FIELDS} + for model_usage_entry in usage_by_model.values(): + entry = _try_to_dict(model_usage_entry) + if entry is None: + return None + for source, target in _MODEL_USAGE_TOKEN_FIELDS: + value = entry.get(source) + if is_numeric(value): + totals[target] += int(value) + + return totals if any(totals.values()) else None + + @dataclasses.dataclass(frozen=True) class ParsedToolName: raw_name: str @@ -729,6 +764,15 @@ def _handle_result(self, message: Any) -> None: self._active_key = None if hasattr(message, "usage"): usage_metrics, usage_metadata = extract_anthropic_usage(message.usage) + # ``usage`` covers only the orchestrator agent; prefer the token + # counts from ``model_usage`` (the per-model breakdown, which + # includes subagent calls) so the logged metrics reflect the whole + # turn. Metadata (service tier, etc.) still comes from ``usage``. + aggregated_usage = _aggregate_model_usage(getattr(message, "model_usage", None)) + if aggregated_usage is not None: + model_usage_metrics, _ = extract_anthropic_usage(aggregated_usage) + if model_usage_metrics: + usage_metrics = model_usage_metrics ctx = self._get_context(None) if ctx.llm_span and (usage_metrics or usage_metadata): ctx.llm_span.log(metrics=usage_metrics or None, metadata=usage_metadata or None)