From ad4ccab85835b88314f42af2de9267899a58116e Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 14 Aug 2026 10:02:32 -0400 Subject: [PATCH 1/2] fix(claude_agent_sdk): prevent orphan spans for racing MCP tools Local MCP handlers could run before the assistant tool_use reached the application-consumed stream, creating duplicate root traces and attaching nested work or errors to the wrong span. This cluttered projects and made tool failures hard to diagnose. Observe raw SDK messages before control dispatch so handlers re-enter the canonical nested tool span without an arbitrary wait. Add deterministic cassette-backed race coverage. --- .../claude_agent_sdk/_test_transport.py | 15 ++ .../claude_agent_sdk/integration.py | 14 +- .../integrations/claude_agent_sdk/patchers.py | 21 ++- .../claude_agent_sdk/test_claude_agent_sdk.py | 24 ++- .../integrations/claude_agent_sdk/tracing.py | 160 +++++++++++++----- 5 files changed, 190 insertions(+), 44 deletions(-) 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 2ed74a976..9397b860c 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.""" @@ -227,6 +233,7 @@ def __init__( self._ready = False self._cursor_lock = anyio.Lock() self._cursor_changed = anyio.Event() + self._mcp_tool_call_read = anyio.Event() self._control_request_ids: dict[str, str] = {} async def connect(self) -> None: @@ -273,6 +280,8 @@ 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() yield message return @@ -300,6 +309,10 @@ 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() + async def end_input(self) -> None: if self._recording: assert self._delegate is not None @@ -326,6 +339,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 diff --git a/py/src/braintrust/integrations/claude_agent_sdk/integration.py b/py/src/braintrust/integrations/claude_agent_sdk/integration.py index f6a994462..168274406 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 f5659584b..4e2b8007f 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 b7fbae46e..2cc7d87b0 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 @@ -80,12 +80,22 @@ def _patched_claude_sdk(*, wrap_client: bool = False, wrap_tool_class: bool = Fa @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"] @@ -152,6 +162,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 +224,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"] diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index c852d2a33..304c78e77 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -60,10 +60,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 +75,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 +255,8 @@ 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 = getattr(_thread_local, "tool_span_tracker", None) + if tool_span_tracker is None: with start_span( name=str(tool_name), span_attributes={"type": SpanTypeAttribute.TOOL}, @@ -285,6 +266,18 @@ async def wrapped_handler(args: Any) -> Any: span.log(output=result) return result + active_tool_span = tool_span_tracker.acquire_span_for_handler(tool_name, args) + if active_tool_span is None: + # A request tracker exists, so the raw SDK message observer normally + # creates the canonical span before MCP dispatch. If instrumentation + # cannot match it, avoid creating a duplicate fallback span and park + # any error for the eventual canonical span. + try: + return await handler(args) + except Exception as exc: + tool_span_tracker.record_handler_error(tool_name, args, exc) + raise + try: return await handler(args) except Exception as exc: @@ -310,10 +303,12 @@ 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]] = {} + self._pending_handler_errors: dict[tuple[str, str], collections.deque[str]] = {} + self._closed = False def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: if llm_span_export is None or not hasattr(message, "content"): @@ -362,9 +357,15 @@ 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) + pending_errors = self._pending_handler_errors.get(dispatch_key) + if pending_errors: + tool_span.log(error=pending_errors.popleft()) + if not pending_errors: + del self._pending_handler_errors[dispatch_key] + def finish_tool_spans(self, message: Any) -> None: if not hasattr(message, "content"): return @@ -402,6 +403,8 @@ def cleanup_all(self, end_time: float | None = None) -> None: """Close all remaining active spans. Called at end-of-stream.""" for tool_use_id in list(self._active_spans): self._end_tool_span(tool_use_id, end_time=end_time) + self._closed = True + self._pending_handler_errors.clear() @property def has_active_spans(self) -> bool: @@ -420,7 +423,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: @@ -429,12 +432,29 @@ def acquire_span_for_handler(self, tool_name: Any, args: Any) -> _ActiveToolSpan matched_span.activate() return matched_span + def record_handler_error(self, tool_name: Any, args: Any, exc: Exception) -> None: + """Log a late handler error now, or park it until its canonical span arrives.""" + active_tool_span = self.acquire_span_for_handler(tool_name, args) + if active_tool_span is not None: + try: + active_tool_span.log_error(exc) + finally: + active_tool_span.release() + return + + if self._closed: + return + + parsed_tool_name = _parse_tool_name(tool_name) + dispatch_key = _make_dispatch_key(parsed_tool_name.display_name, args) + self._pending_handler_errors.setdefault(dispatch_key, collections.deque()).append(str(exc)) + 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 +481,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 +531,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 +681,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 +962,36 @@ def __init__( query_start_time=query_start_time, captured_messages=captured_messages, ) + self._pretraced_message_types: collections.deque[str] = collections.deque() self._finished = False + _thread_local.request_tracker = self 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 + + 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 +1030,9 @@ def finish(self, *, log_output: bool = False) -> None: self._context_tracker.cleanup() self._root_span.end() self._finished = True + self._pretraced_message_types.clear() + if getattr(_thread_local, "request_tracker", None) is self: + delattr(_thread_local, "request_tracker") 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 +1050,41 @@ 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 None: + request_tracker = getattr(_thread_local, "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 raw message observation before the SDK reader starts.""" + _install_query_message_tracing(instance) + return wrapped(*args, **kwargs) + + def _prepare_prompt_for_tracing(prompt: Any) -> tuple[Any, str | None, list[dict[str, Any]] | None]: if prompt is None: return None, None, None @@ -1139,6 +1214,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 +1225,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 +1276,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 From 7af12dae5fd7762daffae68cc0111051cc38bdf2 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 14 Aug 2026 10:40:16 -0400 Subject: [PATCH 2/2] fix(claude_agent_sdk): isolate concurrent tool tracing Bind raw SDK readers to request-owned trackers without retaining ContextVar tokens across async generator yields. Restore fallback tracing for unmatched local handlers and add deterministic cassette coverage for concurrent clients, query helpers, cancellation, and cross-context cleanup. --- .../claude_agent_sdk/_test_transport.py | 41 ++- .../claude_agent_sdk/test_claude_agent_sdk.py | 271 ++++++++++++++++-- .../integrations/claude_agent_sdk/tracing.py | 90 +++--- 3 files changed, 324 insertions(+), 78 deletions(-) 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 9397b860c..a8bb4d42c 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/_test_transport.py @@ -219,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 @@ -234,6 +236,9 @@ def __init__( 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: @@ -282,6 +287,10 @@ async def _read_messages_impl(self): 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 @@ -289,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 @@ -313,11 +327,32 @@ 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 @@ -396,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/test_claude_agent_sdk.py b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 2cc7d87b0..982497374 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,6 +78,38 @@ 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): @@ -123,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="", @@ -248,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}} @@ -2259,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 304c78e77..059c8e540 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) @@ -255,8 +264,14 @@ def _wrap_tool_handler(handler: Any, tool_name: Any) -> Any: return handler async def wrapped_handler(args: Any) -> Any: - tool_span_tracker = getattr(_thread_local, "tool_span_tracker", None) + 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}, @@ -266,18 +281,6 @@ async def wrapped_handler(args: Any) -> Any: span.log(output=result) return result - active_tool_span = tool_span_tracker.acquire_span_for_handler(tool_name, args) - if active_tool_span is None: - # A request tracker exists, so the raw SDK message observer normally - # creates the canonical span before MCP dispatch. If instrumentation - # cannot match it, avoid creating a duplicate fallback span and park - # any error for the eventual canonical span. - try: - return await handler(args) - except Exception as exc: - tool_span_tracker.record_handler_error(tool_name, args, exc) - raise - try: return await handler(args) except Exception as exc: @@ -307,8 +310,6 @@ def __init__(self): # 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]] = {} - self._pending_handler_errors: dict[tuple[str, str], collections.deque[str]] = {} - self._closed = False def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: if llm_span_export is None or not hasattr(message, "content"): @@ -360,12 +361,6 @@ def start_tool_spans(self, message: Any, llm_span_export: str | None) -> None: dispatch_key = _make_dispatch_key(parsed_tool_name.display_name, tool_input) self._dispatch_queues.setdefault(dispatch_key, collections.deque()).append(tool_use_id) - pending_errors = self._pending_handler_errors.get(dispatch_key) - if pending_errors: - tool_span.log(error=pending_errors.popleft()) - if not pending_errors: - del self._pending_handler_errors[dispatch_key] - def finish_tool_spans(self, message: Any) -> None: if not hasattr(message, "content"): return @@ -403,8 +398,6 @@ def cleanup_all(self, end_time: float | None = None) -> None: """Close all remaining active spans. Called at end-of-stream.""" for tool_use_id in list(self._active_spans): self._end_tool_span(tool_use_id, end_time=end_time) - self._closed = True - self._pending_handler_errors.clear() @property def has_active_spans(self) -> bool: @@ -432,23 +425,6 @@ def acquire_span_for_handler(self, tool_name: Any, args: Any) -> _ActiveToolSpan matched_span.activate() return matched_span - def record_handler_error(self, tool_name: Any, args: Any, exc: Exception) -> None: - """Log a late handler error now, or park it until its canonical span arrives.""" - active_tool_span = self.acquire_span_for_handler(tool_name, args) - if active_tool_span is not None: - try: - active_tool_span.log_error(exc) - finally: - active_tool_span.release() - return - - if self._closed: - return - - parsed_tool_name = _parse_tool_name(tool_name) - dispatch_key = _make_dispatch_key(parsed_tool_name.display_name, args) - self._pending_handler_errors.setdefault(dispatch_key, collections.deque()).append(str(exc)) - def _match_via_dispatch_queue( self, display_name: str, args: Any, candidates: list[_ActiveToolSpan] ) -> _ActiveToolSpan | None: @@ -964,7 +940,6 @@ def __init__( ) self._pretraced_message_types: collections.deque[str] = collections.deque() self._finished = False - _thread_local.request_tracker = self def add_message(self, message: Any) -> None: message_type = type(message).__name__ @@ -978,6 +953,8 @@ def add_raw_message(self, data: Any) -> None: 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 @@ -1031,8 +1008,6 @@ def finish(self, *, log_output: bool = False) -> None: self._root_span.end() self._finished = True self._pretraced_message_types.clear() - if getattr(_thread_local, "request_tracker", None) is self: - delattr(_thread_local, "request_tracker") def _hook_parent_export(self, tool_use_id: str | None) -> str: tool_export = self._context_tracker.get_tool_span_export(tool_use_id) @@ -1060,8 +1035,6 @@ def __init__(self, send_stream: Any, query: Any) -> None: async def send(self, value: Any) -> None: request_tracker = getattr(self._query, "_braintrust_request_tracker", None) - if request_tracker is None: - request_tracker = getattr(_thread_local, "request_tracker", None) if request_tracker is not None: request_tracker.add_raw_message(value) await self._send_stream.send(value) @@ -1080,11 +1053,33 @@ def _install_query_message_tracing(query: Any) -> None: def _wrap_query_read_messages(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: - """Install raw message observation before the SDK reader starts.""" + """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 @@ -1148,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, ):