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):