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

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

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)