feat: Add support for claude agent sdk hooks

Author: AbhiPrasad

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

@@ -19,7 +19,7 @@
 
 from braintrust.logger import NOOP_SPAN, current_span, init_logger
 
-from ._wrapper import _create_client_wrapper_class, _create_tool_wrapper_class, _wrap_tool_factory
+from ._wrapper import _create_client_wrapper_class, _create_tool_wrapper_class, _wrap_module_query, _wrap_tool_factory
 
 logger = logging.getLogger(__name__)
 
@@ -69,6 +69,7 @@ def setup_claude_agent_sdk(
         original_client = claude_agent_sdk.ClaudeSDKClient if hasattr(claude_agent_sdk, "ClaudeSDKClient") 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_query_fn = claude_agent_sdk.query if hasattr(claude_agent_sdk, "query") else None
 
         if original_client:
             wrapped_client = _create_client_wrapper_class(original_client)
@@ -97,6 +98,15 @@ def setup_claude_agent_sdk(
                     if getattr(module, "tool", None) is original_tool_fn:
                         setattr(module, "tool", wrapped_tool_fn)
 
+        if original_query_fn:
+            wrapped_query_fn = _wrap_module_query(original_query_fn)
+            claude_agent_sdk.query = wrapped_query_fn
+
+            for module in list(sys.modules.values()):
+                if module and hasattr(module, "query"):
+                    if getattr(module, "query", None) is original_query_fn:
+                        setattr(module, "query", wrapped_query_fn)
+
         return True
     except ImportError:
         # Not installed - this is expected when using auto_instrument()