clean up

Author: AbhiPrasad

py/src/braintrust/wrappers/claude_agent_sdk/__init__.py

@@ -71,10 +71,17 @@ def setup_claude_agent_sdk(
 
         import claude_agent_sdk
 
+        def _module_references_any(module: object, references: list[tuple[str, object | None]]) -> bool:
+            for attr_name, original_value in references:
+                if original_value is not None and getattr(module, attr_name, None) is original_value:
+                    return True
+            return False
+
         original_client = claude_agent_sdk.ClaudeSDKClient if hasattr(claude_agent_sdk, "ClaudeSDKClient") else None
         original_query_fn = claude_agent_sdk.query if hasattr(claude_agent_sdk, "query") else None
         original_tool_class = claude_agent_sdk.SdkMcpTool if hasattr(claude_agent_sdk, "SdkMcpTool") else None
         original_tool_fn = claude_agent_sdk.tool if hasattr(claude_agent_sdk, "tool") else None
+        original_options_class = claude_agent_sdk.ClaudeAgentOptions if hasattr(claude_agent_sdk, "ClaudeAgentOptions") else None
 
         wrapped_client = None
         if original_client:
@@ -90,9 +97,15 @@ def setup_claude_agent_sdk(
             wrapped_query_fn = _wrap_query_function(original_query_fn, wrapped_client)
             claude_agent_sdk.query = wrapped_query_fn
 
+            query_patch_anchors = [
+                ("ClaudeSDKClient", original_client),
+                ("ClaudeAgentOptions", original_options_class),
+                ("SdkMcpTool", original_tool_class),
+                ("tool", original_tool_fn),
+            ]
             for module in list(sys.modules.values()):
-                if module and hasattr(module, "query"):
-                    if getattr(module, "query", None) is original_query_fn:
+                if module and getattr(module, "query", None) is original_query_fn:
+                    if _module_references_any(module, query_patch_anchors):
                         setattr(module, "query", wrapped_query_fn)
 
         if original_tool_class:

py/src/braintrust/wrappers/claude_agent_sdk/_wrapper.py

@@ -3,9 +3,10 @@
 import logging
 import threading
 import time
-from collections.abc import AsyncGenerator, AsyncIterable, Mapping
+from collections.abc import AsyncGenerator, AsyncIterable
 from typing import Any
 
+from braintrust.bt_json import bt_safe_deep_copy
 from braintrust.logger import start_span
 from braintrust.span_types import SpanTypeAttribute
 from braintrust.wrappers._anthropic_utils import Wrapper, extract_anthropic_usage, finalize_anthropic_tokens
@@ -165,32 +166,8 @@ def _serialize_system_message(message: Any) -> dict[str, Any]:
 
     return serialized
 
-
-def _serialize_hook_value(value: Any) -> Any:
-    if value is None or isinstance(value, (bool, int, float, str)):
-        return value
-
-    if dataclasses.is_dataclass(value):
-        return _serialize_hook_value(dataclasses.asdict(value))
-
-    if isinstance(value, Mapping):
-        return {str(key): _serialize_hook_value(item) for key, item in value.items()}
-
-    if isinstance(value, (list, tuple)):
-        return [_serialize_hook_value(item) for item in value]
-
-    if hasattr(value, "__dict__"):
-        return {
-            key: _serialize_hook_value(item)
-            for key, item in vars(value).items()
-            if not key.startswith("_") and not callable(item)
-        }
-
-    return str(value)
-
-
 def _serialize_hook_context(context: Any) -> dict[str, Any] | None:
-    serialized = _serialize_hook_value(context)
+    serialized = bt_safe_deep_copy(context)
     if not isinstance(serialized, dict):
         return None
 
@@ -224,7 +201,7 @@ async def wrapped_hook(*args: Any, **kwargs: Any) -> Any:
         tool_use_id = kwargs.get("tool_use_id") if "tool_use_id" in kwargs else (args[1] if len(args) > 1 else None)
         context = kwargs.get("context") if "context" in kwargs else (args[2] if len(args) > 2 else None)
 
-        span_input = {"input": _serialize_hook_value(hook_input)}
+        span_input = {"input": bt_safe_deep_copy(hook_input)}
         if tool_use_id is not None:
             span_input["tool_use_id"] = str(tool_use_id)
 
@@ -249,7 +226,7 @@ async def wrapped_hook(*args: Any, **kwargs: Any) -> Any:
             result = callback(*args, **kwargs)
             if inspect.isawaitable(result):
                 result = await result
-            span.log(output=_serialize_hook_value(result))
+            span.log(output=bt_safe_deep_copy(result))
             return result
 
     wrapped_hook._braintrust_wrapped_claude_hook = True  # type: ignore[attr-defined]
@@ -281,6 +258,8 @@ def _wrap_options_hooks(options: Any) -> None:
     hooks_by_event = getattr(options, "hooks", None)
     if not isinstance(hooks_by_event, dict):
         return
+    if getattr(options, "_braintrust_wrapped_claude_hooks_ref", None) is hooks_by_event:
+        return
 
     for event_name, matchers in list(hooks_by_event.items()):
         if not isinstance(matchers, list):
@@ -295,6 +274,11 @@ def _wrap_options_hooks(options: Any) -> None:
         except Exception:
             continue
 
+    try:
+        setattr(options, "_braintrust_wrapped_claude_hooks_ref", hooks_by_event)
+    except Exception:
+        pass
+
 
 def _wrap_client_hooks(client: Any) -> None:
     _wrap_options_hooks(getattr(client, "options", None))
@@ -406,9 +390,7 @@ async def _trace_message_stream(
                     task_events.append(_serialize_system_message(message))
 
                 yield message
-        except Exception:
-            raise
-        else:
+
             if final_results:
                 span.log(output=final_results[-1])
         finally:
@@ -892,6 +874,7 @@ async def query(self, *args: Any, **kwargs: Any) -> Any:
             """Wrap query to capture the prompt and start time for tracing."""
             # Capture the time when query is called (when LLM call starts)
             self.__query_start_time = time.time()
+            # Re-wrap if client.options.hooks was replaced after client construction.
             _wrap_client_hooks(self.__client)
 
             # Capture the prompt for use in receive_response
@@ -936,16 +919,13 @@ def _wrap_query_function(query_fn: Any, client_class: Any) -> Any:
         return query_fn
 
     def wrapped_query(*args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]:
-        prompt = kwargs.get("prompt")
-        if prompt is None and args:
-            prompt = args[0]
-
-        options = kwargs.get("options")
-        transport = kwargs.get("transport")
+        query_kwargs = dict(kwargs)
+        options = query_kwargs.pop("options", None)
+        transport = query_kwargs.pop("transport", None)
 
         async def traced_generator() -> AsyncGenerator[Any, None]:
             async with client_class(options=options, transport=transport) as client:
-                await client.query(prompt)
+                await client.query(*args, **query_kwargs)
                 async for message in client.receive_response():
                     yield message
 

py/src/braintrust/wrappers/claude_agent_sdk/test_wrapper.py

@@ -36,9 +36,11 @@
     _extract_usage_from_result_message,
     _parse_tool_name,
     _serialize_content_blocks,
+    _serialize_hook_context,
     _serialize_system_message,
     _serialize_tool_result_output,
     _thread_local,
+    _wrap_client_hooks,
     _wrap_query_function,
 )
 from braintrust.wrappers.test_utils import verify_autoinstrument_script
@@ -81,6 +83,120 @@ def _patched_claude_sdk(*, wrap_client: bool = False, wrap_tool_class: bool = Fa
         claude_agent_sdk.tool = original_tool_fn
 
 
+def test_wrap_client_hooks_fast_exits_for_same_hooks_object():
+    class FakeOptions:
+        def __init__(self, hooks: dict[str, list[Any]]):
+            self.hooks = hooks
+
+    class FakeClient:
+        def __init__(self, options: FakeOptions):
+            self.options = options
+
+    def callback(value: Any) -> Any:
+        return value
+
+    matcher = types.SimpleNamespace(hooks=[callback], matcher="tool")
+    hooks = {"PreToolUse": [matcher]}
+    client = FakeClient(FakeOptions(hooks))
+
+    _wrap_client_hooks(client)
+    wrapped_callback = matcher.hooks[0]
+    wrapped_ref = client.options._braintrust_wrapped_claude_hooks_ref
+
+    _wrap_client_hooks(client)
+
+    assert matcher.hooks[0] is wrapped_callback
+    assert client.options._braintrust_wrapped_claude_hooks_ref is wrapped_ref
+
+
+def test_wrap_client_hooks_rewraps_when_hooks_container_is_replaced():
+    class FakeOptions:
+        def __init__(self, hooks: dict[str, list[Any]]):
+            self.hooks = hooks
+
+    class FakeClient:
+        def __init__(self, options: FakeOptions):
+            self.options = options
+
+    def first_callback(value: Any) -> Any:
+        return value
+
+    def second_callback(value: Any) -> Any:
+        return value
+
+    first_matcher = types.SimpleNamespace(hooks=[first_callback], matcher="tool")
+    client = FakeClient(FakeOptions({"PreToolUse": [first_matcher]}))
+
+    _wrap_client_hooks(client)
+    assert hasattr(first_matcher.hooks[0], "_braintrust_wrapped_claude_hook")
+
+    second_matcher = types.SimpleNamespace(hooks=[second_callback], matcher="tool")
+    client.options.hooks = {"PreToolUse": [second_matcher]}
+
+    _wrap_client_hooks(client)
+
+    assert hasattr(second_matcher.hooks[0], "_braintrust_wrapped_claude_hook")
+    assert client.options._braintrust_wrapped_claude_hooks_ref is client.options.hooks
+
+
+@pytest.mark.asyncio
+async def test_wrap_query_function_forwards_unknown_kwargs():
+    call_log: dict[str, Any] = {}
+
+    async def original_query(*args: Any, **kwargs: Any) -> AsyncIterable[str]:
+        del args, kwargs
+        if False:
+            yield ""
+
+    class FakeWrappedClient:
+        def __init__(self, *args: Any, **kwargs: Any):
+            call_log["client_init_kwargs"] = kwargs
+
+        async def __aenter__(self) -> "FakeWrappedClient":
+            return self
+
+        async def __aexit__(self, *args: Any) -> None:
+            del args
+            return None
+
+        async def query(self, *args: Any, **kwargs: Any) -> None:
+            call_log["query_args"] = args
+            call_log["query_kwargs"] = kwargs
+
+        async def receive_response(self) -> AsyncIterable[str]:
+            yield "ok"
+
+    wrapped_query = _wrap_query_function(original_query, FakeWrappedClient)
+
+    messages = [
+        message
+        async for message in wrapped_query(
+            "Say hi",
+            options="opts",
+            transport="transport",
+            session_id="session-123",
+            custom_flag=True,
+        )
+    ]
+
+    assert messages == ["ok"]
+    assert call_log["client_init_kwargs"] == {"options": "opts", "transport": "transport"}
+    assert call_log["query_args"] == ("Say hi",)
+    assert call_log["query_kwargs"] == {"session_id": "session-123", "custom_flag": True}
+
+
+def test_serialize_hook_context_handles_cyclic_containers():
+    data: dict[str, Any] = {"name": "root", "signal": "omit-me"}
+    data["child"] = {"parent": data}
+
+    serialized = _serialize_hook_context(data)
+
+    assert serialized is not None
+    assert serialized["name"] == "root"
+    assert serialized["child"]["parent"] == "<circular reference>"
+    assert "signal" not in serialized
+
+
 @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):