From bd9f3dbe531da08fe67be1a531f147540bb88fab Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 13 Jul 2026 16:31:29 -0400 Subject: [PATCH] fix(claude_agent_sdk): attribute usage to model call spans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Claude Agent runs may make several orchestrator and delegated model calls, but cumulative result usage was being assigned to one final span. This hid delegated usage and could price tokens using the wrong model. Before: Claude Agent [TASK] ├── root LLM [opus] metrics: none ├── Agent [TOOL] │ └── subagent [TASK] │ └── LLM [haiku] metrics: none └── final LLM [opus] metrics: cumulative usage for every call After: Claude Agent [TASK] ├── root LLM [opus] metrics: root call usage ├── Agent [TOOL] │ └── subagent [TASK] │ └── LLM [haiku] metrics: delegated call usage └── final LLM [opus] metrics: reconciled final-call remainder Each LLM span now keeps its own metadata.model and token/cache metrics, including cache-read, cache-creation, 5m/1h cache TTL, and server-tool usage. Repeated cumulative snapshots are converted to deltas so trace totals do not double-count them. Result usage also provides a safe fallback when assistant usage is unavailable. Result total_cost_usd remains request metadata on the root task rather than a parent-owned estimated cost. This keeps trace totals and per-model cost estimates accurate across root, subagent, concurrent-agent, and older SDK flows. Extend cassette-backed coverage for multi-turn and concurrent-agent requests on both supported Claude Agent SDK versions. --- .../claude_agent_sdk/test_claude_agent_sdk.py | 151 ++++++++++++++- .../integrations/claude_agent_sdk/tracing.py | 183 +++++++++++++++++- 2 files changed, 318 insertions(+), 16 deletions(-) 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..3c62b370 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 @@ -149,11 +149,14 @@ async def calculator_handler(args): options=options, ) + assistant_messages = [] result_message = None async with claude_agent_sdk.ClaudeSDKClient(options=options, transport=transport) as client: await client.query("What is 15 multiplied by 7? Then subtract 5 from the result.") async for message in client.receive_response(): - if type(message).__name__ == "ResultMessage": + if type(message).__name__ == "AssistantMessage": + assistant_messages.append(message) + elif type(message).__name__ == "ResultMessage": result_message = message spans = memory_logger.pop() @@ -186,8 +189,23 @@ async def calculator_handler(args): llm_span_ids = {span["span_id"] for span in llm_spans} _assert_llm_spans_have_time_to_first_token(llm_spans) - llm_spans_with_metrics = [s for s in llm_spans if "prompt_tokens" in s.get("metrics", {})] - assert len(llm_spans_with_metrics) >= 1, "At least one LLM span should have token metrics" + _assert_per_span_usage_and_result_reconciliation( + llm_spans, + assistant_messages, + result_message, + task_span, + ) + + first_llm_span = min(llm_spans, key=lambda span: span["metrics"]["start"]) + first_assistant_completion = _assistant_usage_metrics(assistant_messages[0])["completion_tokens"] + assert first_llm_span["metrics"]["completion_tokens"] == first_assistant_completion + + result_usage_metrics, _ = extract_anthropic_usage(result_message.usage) + for metric_name, expected_total in result_usage_metrics.items(): + observed_total = sum(llm_span.get("metrics", {}).get(metric_name, 0) for llm_span in llm_spans) + assert observed_total == expected_total, [ + (metric_name, llm_span.get("metrics", {}).get(metric_name)) for llm_span in llm_spans + ] for llm_span in llm_spans: assert llm_span["span_attributes"]["name"] == "anthropic.messages.create" @@ -195,7 +213,7 @@ async def calculator_handler(args): assert len(llm_span["output"]) > 0 for metric_name in ("prompt_tokens", "completion_tokens", "tokens"): if metric_name in llm_span.get("metrics", {}): - assert llm_span["metrics"][metric_name] > 0 + assert llm_span["metrics"][metric_name] >= 0 assert any(llm_span.get("metadata", {}).get("usage_service_tier") == "standard" for llm_span in llm_spans) if any("usage_inference_geo" in llm_span.get("metadata", {}) for llm_span in llm_spans): assert all( @@ -242,6 +260,86 @@ def _assert_llm_spans_have_time_to_first_token(llm_spans: list[dict[str, Any]]) assert llm_span["metrics"]["time_to_first_token"] >= 0 +def _assistant_usage_metrics(message: Any) -> dict[str, float]: + metrics, _ = extract_anthropic_usage(getattr(message, "usage", None)) + return metrics + + +def _assert_per_span_usage_and_result_reconciliation( + llm_spans: list[dict[str, Any]], + assistant_messages: list[Any], + result_message: Any, + root_task_span: dict[str, Any], +) -> None: + assert llm_spans + assert assistant_messages + assert result_message is not None + + assistant_models = {getattr(message, "model", None) for message in assistant_messages} + assistant_models.discard(None) + assistant_completion_by_model: dict[str, set[float]] = {} + for message in assistant_messages: + model = getattr(message, "model", None) + completion_tokens = _assistant_usage_metrics(message).get("completion_tokens") + if model is not None and completion_tokens is not None: + assistant_completion_by_model.setdefault(model, set()).add(completion_tokens) + + for llm_span in llm_spans: + assert llm_span["span_attributes"]["name"] == "anthropic.messages.create" + metrics = llm_span.get("metrics", {}) + for metric_name in ("prompt_tokens", "completion_tokens", "tokens"): + assert metric_name in metrics, f"Missing {metric_name} on LLM span {llm_span}" + assert metrics[metric_name] >= 0 + model = llm_span.get("metadata", {}).get("model") + assert model in assistant_models + + assert any(root_task_span["span_id"] in span.get("span_parents", []) for span in llm_spans) + final_assistant = next( + message for message in reversed(assistant_messages) if getattr(message, "parent_tool_use_id", None) is None + ) + final_assistant_content = _serialize_content_blocks(final_assistant.content) + final_root_span = next( + span + for span in llm_spans + if any( + isinstance(output.get("content"), list) + and output["content"][-len(final_assistant_content) :] == final_assistant_content + for output in span.get("output", []) + ) + ) + + for llm_span in llm_spans: + if llm_span is final_root_span: + continue + model = llm_span["metadata"]["model"] + assert ( + llm_span["metrics"]["completion_tokens"] == 0 + or llm_span["metrics"]["completion_tokens"] in assistant_completion_by_model[model] + ) + + result_metrics, result_metadata = extract_anthropic_usage(result_message.usage) + result_completion = result_metrics["completion_tokens"] + prior_completion = sum(span["metrics"]["completion_tokens"] for span in llm_spans if span is not final_root_span) + final_assistant_completion = _assistant_usage_metrics(final_assistant)["completion_tokens"] + expected_final_completion = ( + result_completion - prior_completion if result_completion >= prior_completion else final_assistant_completion + ) + assert final_root_span["metrics"]["completion_tokens"] == expected_final_completion + + if result_completion >= prior_completion: + assert sum(span["metrics"]["completion_tokens"] for span in llm_spans) == result_completion + + for metric_name, value in result_metrics.items(): + if metric_name.startswith("server_tool_use_"): + assert final_root_span["metrics"][metric_name] == value + for metadata_name, value in result_metadata.items(): + assert final_root_span["metadata"][metadata_name] == value + + assert final_root_span["metrics"]["tokens"] == ( + final_root_span["metrics"]["prompt_tokens"] + expected_final_completion + ) + + def _sdk_cassette_name(base: str, *, min_version: str) -> str: """Return base cassette name for SDK >= min_version, else a version-specific variant.""" if _sdk_version_at_least(min_version): @@ -665,6 +763,8 @@ async def test_bundled_subagent_creates_task_span(memory_logger): options=options, ) + assistant_messages = [] + result_message = None async with claude_agent_sdk.ClaudeSDKClient(options=options, transport=transport) as client: await client.query( "You must delegate this task to the bundled general-purpose agent. " @@ -672,7 +772,10 @@ async def test_bundled_subagent_creates_task_span(memory_logger): "Do not answer directly without using the subagent." ) async for message in client.receive_response(): - if type(message).__name__ == "ResultMessage": + if type(message).__name__ == "AssistantMessage": + assistant_messages.append(message) + elif type(message).__name__ == "ResultMessage": + result_message = message break spans = memory_logger.pop() @@ -704,6 +807,12 @@ async def test_bundled_subagent_creates_task_span(memory_logger): llm_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.LLM] _assert_llm_spans_have_time_to_first_token(llm_spans) + _assert_per_span_usage_and_result_reconciliation( + llm_spans, + assistant_messages, + result_message, + root_task_span, + ) assert any( subagent_span["span_id"] in llm_span["span_parents"] for subagent_span in subagent_spans @@ -717,6 +826,22 @@ async def test_bundled_subagent_creates_task_span(memory_logger): ] assert delegated_llm_spans, "Expected at least one delegated LLM span nested under a subagent task span" + delegated_assistant_messages = [ + message for message in assistant_messages if getattr(message, "parent_tool_use_id", None) is not None + ] + delegated_usage_by_model = { + ( + getattr(message, "model", None), + _assistant_usage_metrics(message).get("completion_tokens"), + ) + for message in delegated_assistant_messages + } + for delegated_llm_span in delegated_llm_spans: + assert ( + delegated_llm_span["metadata"]["model"], + delegated_llm_span["metrics"]["completion_tokens"], + ) in delegated_usage_by_model + assert any( any(llm_span["span_id"] in tool_span["span_parents"] for llm_span in delegated_llm_spans) for tool_span in tool_spans @@ -2271,6 +2396,8 @@ async def test_concurrent_subagents_produce_parallel_llm_spans_with_correct_pare options=options, ) + assistant_messages = [] + result_message = None async with claude_agent_sdk.ClaudeSDKClient(options=options, transport=transport) as client: await client.query( "Use exactly three bundled general-purpose subagents and start all three Agent tool calls " @@ -2283,7 +2410,10 @@ async def test_concurrent_subagents_produce_parallel_llm_spans_with_correct_pare "Do not ask clarifying questions. Do not answer directly without using all three subagents." ) async for message in client.receive_response(): - if type(message).__name__ == "ResultMessage": + if type(message).__name__ == "AssistantMessage": + assistant_messages.append(message) + elif type(message).__name__ == "ResultMessage": + result_message = message break spans = memory_logger.pop() @@ -2292,8 +2422,17 @@ async def test_concurrent_subagents_produce_parallel_llm_spans_with_correct_pare tool_spans = find_spans_by_type(spans, SpanTypeAttribute.TOOL) root_task_span = find_span_by_name(task_spans, "Claude Agent") + _assert_per_span_usage_and_result_reconciliation( + llm_spans, + assistant_messages, + result_message, + root_task_span, + ) if not _sdk_version_at_least("0.1.11"): + assert {span["metadata"]["model"] for span in llm_spans} == { + getattr(message, "model", None) for message in assistant_messages + } return subagent_spans = [span for span in task_spans if span["span_id"] != root_task_span["span_id"]] diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index 631b6f99..31ff1e46 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -1,13 +1,15 @@ import asyncio import collections +import copy import dataclasses +import importlib import json import threading import time 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, @@ -540,6 +542,128 @@ def _task_output(message: Any) -> dict[str, Any] | None: } +def _ensure_assistant_usage_parsing() -> None: + """Retain assistant usage that claude-agent-sdk 0.1.10 drops while parsing.""" + try: + types_module = importlib.import_module("claude_agent_sdk.types") + assistant_message_class = getattr(types_module, MessageClassName.ASSISTANT) + if "usage" in getattr(assistant_message_class, "__dataclass_fields__", {}): + return + parser_module = importlib.import_module("claude_agent_sdk._internal.message_parser") + except ImportError: + return + + original_parse_message = getattr(parser_module, "parse_message", None) + if original_parse_message is None or getattr(original_parse_message, "_braintrust_retains_usage", False): + return + + def parse_message_with_usage(data: Any) -> Any: + message = original_parse_message(data) + if type(message).__name__ == MessageClassName.ASSISTANT and getattr(message, "usage", None) is None: + raw_message = data.get("message") if isinstance(data, dict) else None + usage = raw_message.get("usage") if isinstance(raw_message, dict) else None + if usage is not None: + message.usage = copy.deepcopy(usage) + return message + + parse_message_with_usage._braintrust_retains_usage = True # type: ignore[attr-defined] + parser_module.parse_message = parse_message_with_usage + + # The one-shot query path imports parse_message at module import time in + # claude-agent-sdk 0.1.10, while ClaudeSDKClient imports it lazily. + try: + client_module = importlib.import_module("claude_agent_sdk._internal.client") + except ImportError: + return + if getattr(client_module, "parse_message", None) is original_parse_message: + client_module.parse_message = parse_message_with_usage + + +def _assistant_usage(message: Any) -> dict[str, Any] | None: + usage = _try_to_dict(getattr(message, "usage", None)) + if not usage: + return None + + metrics, metadata = extract_anthropic_usage(usage) + if not metrics and not metadata: + return None + return copy.deepcopy(usage) + + +def _numeric_usage_value(usage: dict[str, Any] | None, *path: str) -> float | None: + value: Any = usage + for name in path: + value = value.get(name) if isinstance(value, dict) else None + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return None + + +def _set_reconciled_usage_value( + merged_usage: dict[str, Any], + result_usage: dict[str, Any], + prior_usages: list[dict[str, Any]], + *path: str, +) -> None: + result_value = _numeric_usage_value(result_usage, *path) + if result_value is None: + return + + prior_value = sum(_numeric_usage_value(usage, *path) or 0.0 for usage in prior_usages) + adjusted_value = result_value - prior_value + if adjusted_value < 0: + return + + target = merged_usage + for name in path[:-1]: + nested = target.get(name) + if not isinstance(nested, dict): + nested = {} + target[name] = nested + target = nested + target[path[-1]] = adjusted_value + + +def _reconciled_usage( + base_usage: dict[str, Any] | None, + result_usage: dict[str, Any], + prior_usages: list[dict[str, Any]], +) -> dict[str, Any]: + merged_usage = copy.deepcopy(base_usage or {}) + for field_name in ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + ): + _set_reconciled_usage_value(merged_usage, result_usage, prior_usages, field_name) + + for container_name in ("cache_creation", "server_tool_use"): + container = _try_to_dict(result_usage.get(container_name)) or {} + for field_name in container: + _set_reconciled_usage_value( + merged_usage, + result_usage, + prior_usages, + container_name, + field_name, + ) + + for field_name in ("service_tier", "inference_geo"): + if field_name in result_usage: + merged_usage[field_name] = copy.deepcopy(result_usage[field_name]) + return merged_usage + + +def _log_anthropic_usage(span: Any, usage: dict[str, Any] | None) -> None: + if span is None or usage is None: + return + + metrics, metadata = extract_anthropic_usage(usage) + if metrics or metadata: + span.log(metrics=metrics or None, metadata=metadata or None) + + def _message_starts_subagent_tool(message: Any) -> bool: if not hasattr(message, "content"): return False @@ -560,6 +684,8 @@ class _AgentContext: llm_span: Any | None = None llm_parent_export: str | None = None llm_output: list[dict[str, Any]] | None = None + llm_usage: dict[str, Any] | None = None + last_usage_snapshot: dict[str, Any] | None = None next_llm_start: float | None = None task_span: Any | None = None task_confirmed: bool = False @@ -584,6 +710,7 @@ def __init__( self._contexts: dict[str | None, _AgentContext] = {None: _AgentContext(next_llm_start=query_start_time)} self._active_key: str | None = None self._task_order: list[str | None] = [] + self._finalized_llm_usages: list[dict[str, Any]] = [] self._final_results: list[dict[str, Any]] = [] self._result_output: Any | None = None @@ -627,9 +754,7 @@ def log_tasks(self) -> None: def cleanup(self) -> None: """End all open LLM spans, TASK spans, and TOOL spans; clear thread-local.""" for ctx in self._contexts.values(): - if ctx.llm_span: - ctx.llm_span.end() - ctx.llm_span = None + self._finish_llm_span(ctx) if ctx.task_span: ctx.task_span.end() ctx.task_span = None @@ -698,6 +823,21 @@ def _handle_assistant(self, message: Any) -> None: parent_export = self._llm_parent_for_message(message) final_content, extended = self._start_or_merge_llm_span(message, parent_export, ctx) + usage = _assistant_usage(message) + if usage is not None and ctx.llm_span is not None: + span_usage = usage + # The CLI may replay the same cumulative usage snapshot after a + # tool result while dispatching another block from one model call. + # Log the delta so a new span does not duplicate that model usage. + if ctx.llm_usage is None and ctx.last_usage_snapshot is not None and usage == ctx.last_usage_snapshot: + span_usage = _reconciled_usage(None, usage, [ctx.last_usage_snapshot]) + ctx.llm_usage = span_usage + ctx.last_usage_snapshot = usage + model = getattr(message, "model", None) + if model is not None: + ctx.llm_span.log(metadata={"model": model}) + _log_anthropic_usage(ctx.llm_span, span_usage) + message_error = getattr(message, "error", None) if message_error and ctx.llm_span is not None: ctx.llm_span.log(error=str(message_error)) @@ -727,11 +867,20 @@ def _handle_user(self, message: Any) -> None: def _handle_result(self, message: Any) -> None: self._active_key = None - if hasattr(message, "usage"): - usage_metrics, usage_metadata = extract_anthropic_usage(message.usage) - 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) + result_usage = _assistant_usage(message) + root_ctx = self._get_context(None) + if result_usage is not None and root_ctx.llm_span is not None: + prior_usages = [ + *self._finalized_llm_usages, + *( + ctx.llm_usage + for key, ctx in self._contexts.items() + if key is not None and ctx.llm_span is not None and ctx.llm_usage is not None + ), + ] + merged_usage = _reconciled_usage(root_ctx.llm_usage, result_usage, prior_usages) + root_ctx.llm_usage = merged_usage + _log_anthropic_usage(root_ctx.llm_span, merged_usage) result_value = getattr(message, "result", None) if result_value is not None: @@ -840,7 +989,7 @@ def _start_or_merge_llm_span( first_token_time = time.time() if ctx.llm_span: - ctx.llm_span.end(end_time=resolved_start) + self._finish_llm_span(ctx, end_time=resolved_start) final_content, span = _create_llm_span_for_messages( [message], @@ -857,6 +1006,18 @@ def _start_or_merge_llm_span( ctx.next_llm_start = None return final_content, False + def _finish_llm_span(self, ctx: _AgentContext, *, end_time: float | None = None) -> None: + if ctx.llm_span is None: + return + + if ctx.llm_usage is not None: + self._finalized_llm_usages.append(copy.deepcopy(ctx.llm_usage)) + ctx.llm_span.end(end_time=end_time) + ctx.llm_span = None + ctx.llm_parent_export = None + ctx.llm_output = None + ctx.llm_usage = None + def _process_task_event(self, message: Any, agent_span_export: str | None) -> None: """Handle TaskStarted / TaskProgress / TaskNotification system messages.""" task_id = _msg_field(message, "task_id") @@ -1022,6 +1183,7 @@ async def _stream_messages_with_tracing( def _create_query_wrapper_function(original_query: Any) -> Any: """Create a tracing wrapper for the exported one-shot ``query()`` helper.""" + _ensure_assistant_usage_parsing() async def wrapped_query(*args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]: query_start_time = time.time() @@ -1052,6 +1214,7 @@ async def wrapped_query(*args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]: def _create_client_wrapper_class(original_client_class: Any) -> Any: """Creates a wrapper class for ClaudeSDKClient that wraps query and receive_response.""" + _ensure_assistant_usage_parsing() class WrappedClaudeSDKClient(Wrapper): def __init__(self, *args: Any, **kwargs: Any):