diff --git a/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py b/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py index 2ed74a97..a8bb4d42 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py @@ -203,6 +203,12 @@ def _sanitize_url_string(value: str) -> str: ) +def _is_mcp_tool_call(message: Any) -> bool: + request = message.get("request") if isinstance(message, dict) else None + mcp_message = request.get("message") if isinstance(request, dict) else None + return isinstance(mcp_message, dict) and mcp_message.get("method") == "tools/call" + + class ClaudeAgentSdkCassetteTransport(Transport): """Record or replay the SDK<->CLI JSON protocol at the transport layer.""" @@ -213,6 +219,8 @@ def __init__( prompt: str | Any, options: ClaudeAgentOptions, record_mode: str | None = None, + pause_before_sdk_messages: bool = False, + pause_before_mcp_tool_calls: bool = False, ) -> None: _require_sdk() self._cassette_name = cassette_name @@ -227,6 +235,10 @@ def __init__( self._ready = False self._cursor_lock = anyio.Lock() self._cursor_changed = anyio.Event() + self._mcp_tool_call_read = anyio.Event() + self._mcp_tool_call_gate = anyio.Event() if pause_before_mcp_tool_calls else None + self._sdk_message_blocked = anyio.Event() if pause_before_sdk_messages else None + self._sdk_message_gate = anyio.Event() if pause_before_sdk_messages else None self._control_request_ids: dict[str, str] = {} async def connect(self) -> None: @@ -273,6 +285,12 @@ async def _read_messages_impl(self): assert self._delegate is not None async for message in self._delegate.read_messages(): self._events.append({"op": "read", "payload": message}) + if _is_mcp_tool_call(message): + self._mcp_tool_call_read.set() + if self._mcp_tool_call_gate is not None: + await self._mcp_tool_call_gate.wait() + if self._sdk_message_gate is not None: + await self._wait_before_sdk_message(message) yield message return @@ -280,7 +298,12 @@ async def _read_messages_impl(self): event = await self._wait_for_event("read", allow_eof=True) if event is None: return - yield self._remap_read_message(event["payload"]) + message = self._remap_read_message(event["payload"]) + if _is_mcp_tool_call(message) and self._mcp_tool_call_gate is not None: + await self._mcp_tool_call_gate.wait() + if self._sdk_message_gate is not None: + await self._wait_before_sdk_message(message) + yield message async def close(self) -> None: self._ready = False @@ -300,11 +323,36 @@ async def close(self) -> None: def is_ready(self) -> bool: return self._ready + async def wait_for_mcp_tool_call(self) -> None: + """Wait until replay delivers a local MCP ``tools/call`` request.""" + await self._mcp_tool_call_read.wait() + + def release_mcp_tool_calls(self) -> None: + assert self._mcp_tool_call_gate is not None + self._mcp_tool_call_gate.set() + + async def wait_until_sdk_message_blocked(self) -> None: + assert self._sdk_message_blocked is not None + await self._sdk_message_blocked.wait() + + def release_sdk_messages(self) -> None: + assert self._sdk_message_gate is not None + self._sdk_message_gate.set() + async def end_input(self) -> None: if self._recording: assert self._delegate is not None await self._delegate.end_input() + async def _wait_before_sdk_message(self, message: Any) -> None: + assert self._sdk_message_gate is not None + assert self._sdk_message_blocked is not None + message_type = message.get("type") if isinstance(message, dict) else None + if message_type in {"control_response", "control_request", "control_cancel_request"}: + return + self._sdk_message_blocked.set() + await self._sdk_message_gate.wait() + def _should_replay(self) -> bool: if self._record_mode == "all": return False @@ -326,6 +374,8 @@ async def _wait_for_event(self, op: str, *, allow_eof: bool = False) -> dict[str old_event = self._cursor_changed self._cursor_changed = anyio.Event() old_event.set() + if op == "read" and _is_mcp_tool_call(event["payload"]): + self._mcp_tool_call_read.set() return event waiter = self._cursor_changed @@ -381,10 +431,14 @@ def make_cassette_transport( prompt: str | Any, options: "ClaudeAgentOptions", record_mode: str | None = None, + pause_before_sdk_messages: bool = False, + pause_before_mcp_tool_calls: bool = False, ) -> ClaudeAgentSdkCassetteTransport: return ClaudeAgentSdkCassetteTransport( cassette_name=cassette_name, prompt=prompt, options=options, record_mode=record_mode, + pause_before_sdk_messages=pause_before_sdk_messages, + pause_before_mcp_tool_calls=pause_before_mcp_tool_calls, ) diff --git a/py/src/braintrust/integrations/claude_agent_sdk/integration.py b/py/src/braintrust/integrations/claude_agent_sdk/integration.py index f6a99446..16827440 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/integration.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/integration.py @@ -2,7 +2,12 @@ from braintrust.integrations.base import BaseIntegration -from .patchers import ClaudeSDKClientPatcher, ClaudeSDKQueryPatcher, SdkMcpToolPatcher +from .patchers import ( + ClaudeSDKClientPatcher, + ClaudeSDKMessageReaderPatcher, + ClaudeSDKQueryPatcher, + SdkMcpToolPatcher, +) class ClaudeAgentSDKIntegration(BaseIntegration): @@ -11,4 +16,9 @@ class ClaudeAgentSDKIntegration(BaseIntegration): name = "claude_agent_sdk" import_names = ("claude_agent_sdk",) min_version = "0.1.10" - patchers = (ClaudeSDKClientPatcher, ClaudeSDKQueryPatcher, SdkMcpToolPatcher) + patchers = ( + ClaudeSDKMessageReaderPatcher, + ClaudeSDKClientPatcher, + ClaudeSDKQueryPatcher, + SdkMcpToolPatcher, + ) diff --git a/py/src/braintrust/integrations/claude_agent_sdk/patchers.py b/py/src/braintrust/integrations/claude_agent_sdk/patchers.py index f5659584..4e2b8007 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/patchers.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/patchers.py @@ -1,8 +1,23 @@ -"""Claude Agent SDK patchers — replacement patchers for ClaudeSDKClient, query, and SdkMcpTool.""" +"""Claude Agent SDK patchers for message reading, clients, queries, and SDK MCP tools.""" -from braintrust.integrations.base import ClassReplacementPatcher +from braintrust.integrations.base import ClassReplacementPatcher, FunctionWrapperPatcher -from .tracing import _create_client_wrapper_class, _create_query_wrapper_function, _create_tool_wrapper_class +from .tracing import ( + _create_client_wrapper_class, + _create_query_wrapper_function, + _create_tool_wrapper_class, + _wrap_query_read_messages, +) + + +class ClaudeSDKMessageReaderPatcher(FunctionWrapperPatcher): + """Observe SDK messages before local MCP control requests can dispatch.""" + + name = "claude_agent_sdk.message_reader" + target_module = "claude_agent_sdk._internal.query" + target_path = "Query._read_messages" + wrapper = staticmethod(_wrap_query_read_messages) + priority = 50 class ClaudeSDKClientPatcher(ClassReplacementPatcher): 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 b7fbae46..98249737 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 @@ -1,6 +1,7 @@ """Tests for the Claude Agent SDK wrapper.""" import asyncio +import contextvars import dataclasses import sys import types @@ -77,15 +78,57 @@ def _patched_claude_sdk(*, wrap_client: bool = False, wrap_tool_class: bool = Fa claude_agent_sdk.query = original_query +def _make_calculator_options(handler: Any) -> Any: + calculator_tool = claude_agent_sdk.SdkMcpTool( + name="calculator", + description="Performs basic arithmetic operations", + input_schema={ + "type": "object", + "properties": { + "operation": { + "type": "string", + "enum": ["add", "subtract", "multiply", "divide"], + "description": "The arithmetic operation to perform", + }, + "a": {"type": "number", "description": "First number"}, + "b": {"type": "number", "description": "Second number"}, + }, + "required": ["operation", "a", "b"], + }, + handler=handler, + ) + return claude_agent_sdk.ClaudeAgentOptions( + model=TEST_MODEL, + permission_mode="bypassPermissions", + mcp_servers={ + "calculator": claude_agent_sdk.create_sdk_mcp_server( + name="calculator", + version="1.0.0", + tools=[calculator_tool], + ) + }, + ) + + @pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") @pytest.mark.asyncio async def test_calculator_with_multiple_operations(memory_logger): - """Test claude_agent.py example - calculator with multiple operations.""" + """Local MCP handlers racing stream consumption reuse the canonical tool spans.""" assert not memory_logger.pop() with _patched_claude_sdk(wrap_client=True, wrap_tool_class=True): # Create calculator tool + handler_started = asyncio.Event() + async def calculator_handler(args): + handler_started.set() + nested_span = start_span( + name=f"nested_calculator_{args['operation']}", + type=SpanTypeAttribute.FUNCTION, + ) + nested_span.log(input=args) + nested_span.end() + operation = args["operation"] a = args["a"] b = args["b"] @@ -113,36 +156,7 @@ async def calculator_handler(args): "content": [{"type": "text", "text": f"The result of {operation}({a}, {b}) is {result}"}], } - calculator_tool = claude_agent_sdk.SdkMcpTool( - name="calculator", - description="Performs basic arithmetic operations", - input_schema={ - "type": "object", - "properties": { - "operation": { - "type": "string", - "enum": ["add", "subtract", "multiply", "divide"], - "description": "The arithmetic operation to perform", - }, - "a": {"type": "number", "description": "First number"}, - "b": {"type": "number", "description": "Second number"}, - }, - "required": ["operation", "a", "b"], - }, - handler=calculator_handler, - ) - - options = claude_agent_sdk.ClaudeAgentOptions( - model=TEST_MODEL, - permission_mode="bypassPermissions", - mcp_servers={ - "calculator": claude_agent_sdk.create_sdk_mcp_server( - name="calculator", - version="1.0.0", - tools=[calculator_tool], - ) - }, - ) + options = _make_calculator_options(calculator_handler) transport = make_cassette_transport( cassette_name="test_calculator_with_multiple_operations", prompt="", @@ -152,6 +166,13 @@ async def calculator_handler(args): 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.") + + # Deterministically let local MCP dispatch race ahead of application + # consumption of the AssistantMessage containing its tool_use block. + # Before the regression fix, this creates an orphan fallback span. + await asyncio.wait_for(transport.wait_for_mcp_tool_call(), timeout=1) + await asyncio.wait_for(handler_started.wait(), timeout=1) + async for message in client.receive_response(): if type(message).__name__ == "ResultMessage": result_message = message @@ -207,12 +228,17 @@ async def calculator_handler(args): if "usage_inference_geo" in llm_span.get("metadata", {}) ) tool_spans = [s for s in spans if s["span_attributes"]["type"] == SpanTypeAttribute.TOOL] + assert len(tool_spans) == 2, "Each local MCP call should create exactly one canonical tool span" for tool_span in tool_spans: assert tool_span["span_attributes"]["name"] == "calculator" assert tool_span["input"] is not None assert tool_span["output"] is not None + assert tool_span.get("metadata", {}).get("gen_ai.tool.call.id") assert any(parent_id in llm_span_ids for parent_id in tool_span["span_parents"]) + nested_span = find_span_by_name(spans, f"nested_calculator_{tool_span['input']['operation']}") + assert tool_span["span_id"] in nested_span["span_parents"] + # Descendants share the task's trace (``root_span_id``); direct children # reference the task's ``span_id`` in ``span_parents``. task_root_span_id = task_span["root_span_id"] @@ -226,6 +252,81 @@ async def calculator_handler(args): assert any(parent_id in llm_span_ids for parent_id in tool_span["span_parents"]) +@pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") +@pytest.mark.asyncio +async def test_local_mcp_handler_uses_its_own_tracker_with_concurrent_client(memory_logger): + assert not memory_logger.pop() + + with _patched_claude_sdk(wrap_client=True, wrap_tool_class=True): + handler_started = asyncio.Event() + + async def calculator_handler(args): + handler_started.set() + nested_span = start_span(name="nested_concurrent_calculator", type=SpanTypeAttribute.FUNCTION) + nested_span.log(input=args) + nested_span.end() + + operation = args["operation"] + result = args["a"] * args["b"] if operation == "multiply" else args["a"] - args["b"] + return { + "content": [ + {"type": "text", "text": f"The result of {operation}({args['a']}, {args['b']}) is {result}"} + ] + } + + calculator_options = _make_calculator_options(calculator_handler) + other_options = claude_agent_sdk.ClaudeAgentOptions( + model="claude-3-5-haiku-20241022", + permission_mode="bypassPermissions", + ) + calculator_transport = make_cassette_transport( + cassette_name="test_calculator_with_multiple_operations", + prompt="", + options=calculator_options, + pause_before_mcp_tool_calls=True, + ) + other_transport = make_cassette_transport( + cassette_name="test_auto_claude_agent_sdk", + prompt="", + options=other_options, + pause_before_sdk_messages=True, + ) + + async with ( + claude_agent_sdk.ClaudeSDKClient( + options=calculator_options, + transport=calculator_transport, + ) as calculator_client, + claude_agent_sdk.ClaudeSDKClient(options=other_options, transport=other_transport) as other_client, + ): + await calculator_client.query("What is 15 multiplied by 7? Then subtract 5 from the result.") + await asyncio.wait_for(calculator_transport.wait_for_mcp_tool_call(), timeout=1) + + await other_client.query("Say hi") + await asyncio.wait_for(other_transport.wait_until_sdk_message_blocked(), timeout=1) + + calculator_transport.release_mcp_tool_calls() + await asyncio.wait_for(handler_started.wait(), timeout=1) + async for _ in calculator_client.receive_response(): + pass + + other_transport.release_sdk_messages() + async for _ in other_client.receive_response(): + pass + + spans = memory_logger.pop() + calculator_spans = [ + span + for span in find_spans_by_type(spans, SpanTypeAttribute.TOOL) + if span["span_attributes"]["name"] == "calculator" + ] + nested_span = find_span_by_name(spans, "nested_concurrent_calculator") + + assert len(calculator_spans) == 2 + multiply_span = next(span for span in calculator_spans if span["input"]["operation"] == "multiply") + assert multiply_span["span_id"] in nested_span["span_parents"] + + def _make_message(content: str) -> dict: """Create a streaming format message dict.""" return {"type": "user", "message": {"role": "user", "content": content}} @@ -2237,6 +2338,138 @@ async def fake_query(*, prompt, **kwargs): assert len(llm_spans) == 1 +@pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") +@pytest.mark.asyncio +async def test_concurrent_query_helpers_keep_raw_messages_request_scoped(memory_logger): + assert not memory_logger.pop() + + async def user_prompt_hook(input_data: Any, tool_use_id: str | None, context: Any) -> dict[str, Any]: + del input_data, tool_use_id, context + return { + "hookSpecificOutput": { + "hookEventName": "UserPromptSubmit", + "additionalContext": "Remember the answer should stay concise.", + } + } + + first_prompt = "Say hi" + second_prompt = "Say hello in one short sentence." + first_options = claude_agent_sdk.ClaudeAgentOptions( + model="claude-3-5-haiku-20241022", + permission_mode="bypassPermissions", + ) + second_options = claude_agent_sdk.ClaudeAgentOptions( + model=TEST_MODEL, + permission_mode="bypassPermissions", + hooks={ + "UserPromptSubmit": [ + claude_agent_sdk.HookMatcher(hooks=[user_prompt_hook]), + ], + }, + ) + first_transport = make_cassette_transport( + cassette_name="test_auto_claude_agent_sdk", + prompt="", + options=first_options, + pause_before_sdk_messages=True, + ) + second_transport = make_cassette_transport( + cassette_name="test_user_prompt_submit_hook_creates_function_span", + prompt="", + options=second_options, + pause_before_sdk_messages=True, + ) + + async def consume(prompt: str, options: Any, transport: Any) -> str: + async def prompt_stream(): + yield { + "type": "user", + "session_id": "default", + "message": {"role": "user", "content": prompt}, + "parent_tool_use_id": None, + } + + result = None + async for message in claude_agent_sdk.query(prompt=prompt_stream(), options=options, transport=transport): + if type(message).__name__ == "ResultMessage": + result = getattr(message, "result", None) + assert isinstance(result, str) + return result + + with _patched_claude_sdk(): + assert setup_claude_agent_sdk(project=PROJECT_NAME, api_key=logger.TEST_API_KEY) + + first_task = asyncio.create_task(consume(first_prompt, first_options, first_transport)) + await asyncio.wait_for(first_transport.wait_until_sdk_message_blocked(), timeout=1) + + second_task = asyncio.create_task(consume(second_prompt, second_options, second_transport)) + await asyncio.wait_for(second_transport.wait_until_sdk_message_blocked(), timeout=1) + + first_transport.release_sdk_messages() + first_result = await asyncio.wait_for(first_task, timeout=1) + second_transport.release_sdk_messages() + second_result = await asyncio.wait_for(second_task, timeout=1) + + spans = memory_logger.pop() + task_spans = find_spans_by_type(spans, SpanTypeAttribute.TASK) + llm_spans = find_spans_by_type(spans, SpanTypeAttribute.LLM) + assert len(task_spans) == 2 + + task_spans_by_prompt = {span["input"][0]["message"]["content"]: span for span in task_spans} + first_task_span = task_spans_by_prompt[first_prompt] + second_task_span = task_spans_by_prompt[second_prompt] + + def text_outputs(root_span_id: str) -> list[str]: + texts = [] + for span in llm_spans: + if span["root_span_id"] != root_span_id: + continue + for message in span.get("output") or []: + for block in message.get("content") or []: + text = block.get("text") + if isinstance(text, str): + texts.append(text) + return texts + + assert text_outputs(first_task_span["root_span_id"]) == [first_result] + assert text_outputs(second_task_span["root_span_id"]) == [second_result] + + +@pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") +@pytest.mark.asyncio +async def test_query_helper_can_close_from_different_task(memory_logger): + assert not memory_logger.pop() + + options = claude_agent_sdk.ClaudeAgentOptions( + model="claude-3-5-haiku-20241022", + permission_mode="bypassPermissions", + ) + transport = make_cassette_transport( + cassette_name="test_auto_claude_agent_sdk", + prompt="", + options=options, + ) + + async def prompt_stream(): + yield { + "type": "user", + "session_id": "default", + "message": {"role": "user", "content": "Say hi"}, + "parent_tool_use_id": None, + } + + with _patched_claude_sdk(): + assert setup_claude_agent_sdk(project=PROJECT_NAME, api_key=logger.TEST_API_KEY) + messages = claude_agent_sdk.query(prompt=prompt_stream(), options=options, transport=transport) + await anext(messages) + close_task = contextvars.Context().run(asyncio.create_task, messages.aclose()) + await close_task + + task_spans = find_spans_by_type(memory_logger.pop(), SpanTypeAttribute.TASK) + assert len(task_spans) == 1 + assert task_spans[0]["input"][0]["message"]["content"] == "Say hi" + + @pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") @pytest.mark.asyncio async def test_concurrent_subagents_produce_parallel_llm_spans_with_correct_parenting(memory_logger, tmp_path): diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index c852d2a3..059c8e54 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -1,5 +1,6 @@ import asyncio import collections +import contextvars import dataclasses import json import threading @@ -40,6 +41,14 @@ def start_span(*args, **kwargs): _thread_local = threading.local() +_request_tracker_context: contextvars.ContextVar[Any | None] = contextvars.ContextVar( + "braintrust_claude_agent_sdk_request_tracker", + default=None, +) +_tool_span_tracker_context: contextvars.ContextVar[Any | None] = contextvars.ContextVar( + "braintrust_claude_agent_sdk_tool_span_tracker", + default=None, +) @dataclasses.dataclass(frozen=True) @@ -60,10 +69,6 @@ class _ActiveToolSpan: parent_tool_use_id: str | None = None handler_active: bool = False - @property - def has_span(self) -> bool: - return True - def activate(self) -> None: self.handler_active = True self.span.set_current() @@ -79,21 +84,6 @@ def release(self) -> None: self.span.unset_current() -class _NoopActiveToolSpan: - @property - def has_span(self) -> bool: - return False - - def log_error(self, exc: Exception) -> None: - del exc - - def release(self) -> None: - return - - -_NOOP_ACTIVE_TOOL_SPAN = _NoopActiveToolSpan() - - def _parse_tool_name(tool_name: Any) -> ParsedToolName: raw_name = str(tool_name) if tool_name is not None else DEFAULT_TOOL_NAME @@ -274,8 +264,14 @@ def _wrap_tool_handler(handler: Any, tool_name: Any) -> Any: return handler async def wrapped_handler(args: Any) -> Any: - active_tool_span = _activate_tool_span_for_handler(tool_name, args) - if not active_tool_span.has_span: + tool_span_tracker = _tool_span_tracker_context.get() + if tool_span_tracker is None: + tool_span_tracker = getattr(_thread_local, "tool_span_tracker", None) + + active_tool_span = ( + tool_span_tracker.acquire_span_for_handler(tool_name, args) if tool_span_tracker is not None else None + ) + if active_tool_span is None: with start_span( name=str(tool_name), span_attributes={"type": SpanTypeAttribute.TOOL}, @@ -310,9 +306,9 @@ class ToolSpanTracker: def __init__(self): self._active_spans: dict[str, _ActiveToolSpan] = {} self._completed_span_exports: dict[str, str] = {} - # Per-(tool_name, input_signature) FIFO queue of tool_use_ids. - # Used by acquire_span_for_handler to disambiguate identical concurrent - # tool calls (same name + same input) from sibling subagents. + # Per-(display_name, input_signature) FIFO queue of tool_use_ids. + # SDK MCP handlers receive the bare display name while stream messages + # use names like ``mcp__server__tool``, so both sides key by display name. self._dispatch_queues: dict[tuple[str, str], collections.deque[str]] = {} def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: @@ -362,7 +358,7 @@ def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: tool_use_id=tool_use_id, parent_tool_use_id=message_parent_tool_use_id, ) - dispatch_key = _make_dispatch_key(parsed_tool_name.raw_name, tool_input) + dispatch_key = _make_dispatch_key(parsed_tool_name.display_name, tool_input) self._dispatch_queues.setdefault(dispatch_key, collections.deque()).append(tool_use_id) def finish_tool_spans(self, message: Any) -> None: @@ -420,7 +416,7 @@ def acquire_span_for_handler(self, tool_name: Any, args: Any) -> _ActiveToolSpan and (active_tool_span.raw_name in candidate_names or active_tool_span.display_name in candidate_names) ] - matched_span = self._match_via_dispatch_queue(parsed_tool_name.raw_name, args, candidates) + matched_span = self._match_via_dispatch_queue(parsed_tool_name.display_name, args, candidates) if matched_span is None: matched_span = _match_tool_span_for_handler(candidates, args) if matched_span is None: @@ -430,11 +426,11 @@ def acquire_span_for_handler(self, tool_name: Any, args: Any) -> _ActiveToolSpan return matched_span def _match_via_dispatch_queue( - self, raw_name: str, args: Any, candidates: list[_ActiveToolSpan] + self, display_name: str, args: Any, candidates: list[_ActiveToolSpan] ) -> _ActiveToolSpan | None: """Use the dispatch queue to match by tool_use_id when multiple identical candidates exist (same name + same input from different subagents).""" - dispatch_key = _make_dispatch_key(raw_name, args) + dispatch_key = _make_dispatch_key(display_name, args) queue = self._dispatch_queues.get(dispatch_key) if not queue: return None @@ -461,7 +457,7 @@ def _end_tool_span( self._completed_span_exports[tool_use_id] = active_tool_span.span.export() # Remove from dispatch queue so stale entries don't accumulate. - dispatch_key = _make_dispatch_key(active_tool_span.raw_name, active_tool_span.input) + dispatch_key = _make_dispatch_key(active_tool_span.display_name, active_tool_span.input) queue = self._dispatch_queues.get(dispatch_key) if queue: try: @@ -511,14 +507,6 @@ def _match_tool_span_for_handler(candidates: list[_ActiveToolSpan], args: Any) - return candidates[0] -def _activate_tool_span_for_handler(tool_name: Any, args: Any) -> _ActiveToolSpan | _NoopActiveToolSpan: - tool_span_tracker = getattr(_thread_local, "tool_span_tracker", None) - if tool_span_tracker is None: - return _NOOP_ACTIVE_TOOL_SPAN - - return tool_span_tracker.acquire_span_for_handler(tool_name, args) or _NOOP_ACTIVE_TOOL_SPAN - - def _msg_field(message: Any, field: str) -> Any: """Read a field from a system message, falling back to message.data for older SDK versions. @@ -669,7 +657,7 @@ def cleanup(self) -> None: ctx.task_span = None self._task_order.clear() self._tool_tracker.cleanup_all() - if hasattr(_thread_local, "tool_span_tracker"): + if getattr(_thread_local, "tool_span_tracker", None) is self._tool_tracker: delattr(_thread_local, "tool_span_tracker") def get_tool_span_export(self, tool_use_id: str | None) -> str | None: @@ -950,11 +938,37 @@ def __init__( query_start_time=query_start_time, captured_messages=captured_messages, ) + self._pretraced_message_types: collections.deque[str] = collections.deque() self._finished = False def add_message(self, message: Any) -> None: + message_type = type(message).__name__ + if self._pretraced_message_types and self._pretraced_message_types[0] == message_type: + self._pretraced_message_types.popleft() + return self._context_tracker.add(message) + def add_raw_message(self, data: Any) -> None: + """Trace one raw SDK message before it is queued for application consumption.""" + if self._finished or not isinstance(data, dict): + return + + _tool_span_tracker_context.set(self._context_tracker._tool_tracker) + + try: + from claude_agent_sdk._internal.message_parser import parse_message + + message = parse_message(data) + if message is None: + return + self._context_tracker.add(message) + except Exception: + # Provider message parsing and instrumentation must not interrupt + # the SDK's reader task. + return + + self._pretraced_message_types.append(type(message).__name__) + def log_error(self, exc: Exception) -> None: self._root_span.log(error=str(exc)) @@ -993,6 +1007,7 @@ def finish(self, *, log_output: bool = False) -> None: self._context_tracker.cleanup() self._root_span.end() self._finished = True + self._pretraced_message_types.clear() def _hook_parent_export(self, tool_use_id: str | None) -> str: tool_export = self._context_tracker.get_tool_span_export(tool_use_id) @@ -1010,6 +1025,61 @@ def _hook_parent_export(self, tool_use_id: str | None) -> str: return self._root_span.export() +class _TracingMessageSendStream: + """Observe raw SDK messages before the provider queues them for consumers.""" + + def __init__(self, send_stream: Any, query: Any) -> None: + self._send_stream = send_stream + self._query = query + self._braintrust_wrapped = True + + async def send(self, value: Any) -> None: + request_tracker = getattr(self._query, "_braintrust_request_tracker", None) + if request_tracker is not None: + request_tracker.add_raw_message(value) + await self._send_stream.send(value) + + def __getattr__(self, name: str) -> Any: + return getattr(self._send_stream, name) + + +def _install_query_message_tracing(query: Any) -> None: + if query is None: + return + send_stream = getattr(query, "_message_send", None) + if send_stream is None or getattr(send_stream, "_braintrust_wrapped", False): + return + query._message_send = _TracingMessageSendStream(send_stream, query) + + +def _wrap_query_read_messages(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: + """Install request-context-aware raw message observation on this SDK reader.""" + request_tracker = _request_tracker_context.get() + if request_tracker is not None: + instance._braintrust_request_tracker = request_tracker + _install_query_message_tracing(instance) + return wrapped(*args, **kwargs) + + +async def _bind_request_tracker_to_query( + generator: AsyncIterable[Any], request_tracker: RequestTracker +) -> AsyncGenerator[Any, None]: + """Bind the query reader during startup without retaining a ContextVar token across yields.""" + iterator = generator.__aiter__() + token = _request_tracker_context.set(request_tracker) + try: + try: + first_message = await anext(iterator) + except StopAsyncIteration: + return + finally: + _request_tracker_context.reset(token) + + yield first_message + async for message in iterator: + yield message + + def _prepare_prompt_for_tracing(prompt: Any) -> tuple[Any, str | None, list[dict[str, Any]] | None]: if prompt is None: return None, None, None @@ -1073,9 +1143,10 @@ async def wrapped_query(*args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]: query_start_time=query_start_time, captured_messages=captured_messages, ) + generator = _bind_request_tracker_to_query(original_query(*args, **kwargs), request_tracker) async for message in _stream_messages_with_tracing( - original_query(*args, **kwargs), + generator, request_tracker=request_tracker, finish_request_tracker=request_tracker.finish, ): @@ -1139,6 +1210,10 @@ def __start_request_tracker(self) -> RequestTracker: query_start_time=self.__query_start_time, captured_messages=self.__captured_messages, ) + query = getattr(self.__client, "_query", None) + _install_query_message_tracing(query) + if query is not None: + query._braintrust_request_tracker = self.__request_tracker return self.__request_tracker def __finish_request_tracker(self, *, log_output: bool = False) -> None: @@ -1146,11 +1221,15 @@ def __finish_request_tracker(self, *, log_output: bool = False) -> None: if request_tracker is None: return + query = getattr(self.__client, "_query", None) + if query is not None and getattr(query, "_braintrust_request_tracker", None) is request_tracker: + delattr(query, "_braintrust_request_tracker") request_tracker.finish(log_output=log_output) self.__request_tracker = None async def connect(self, *args: Any, **kwargs: Any) -> Any: result = await self.__client.connect(*args, **kwargs) + _install_query_message_tracing(getattr(self.__client, "_query", None)) self.__instrument_hook_callbacks() return result @@ -1193,6 +1272,7 @@ async def receive_response(self) -> AsyncGenerator[Any, None]: async def __aenter__(self) -> "WrappedClaudeSDKClient": await self.__client.__aenter__() + _install_query_message_tracing(getattr(self.__client, "_query", None)) self.__instrument_hook_callbacks() return self