fix: Make ADK and pydantic integration wrappers thread safe (#23)

ibolmo · AbhiPrasad · web-flow · commit 7da14b1a295c · 2026-03-05T11:57:31.000-05:00
Co-authored-by: Abhijeet Prasad &lt;abhijeet@braintrustdata.com&gt;
diff --git a/py/src/braintrust/wrappers/adk/__init__.py b/py/src/braintrust/wrappers/adk/__init__.py
@@ -1,3 +1,4 @@
+import contextvars
 import inspect
 import logging
 import time
@@ -55,6 +56,15 @@ def setup_adk(
         runners.Runner = wrap_runner(runners.Runner)
         base_llm_flow.BaseLlmFlow = wrap_flow(base_llm_flow.BaseLlmFlow)
 
+        try:
+            from google.adk.platform import thread as adk_thread
+
+            adk_thread.create_thread = _wrap_create_thread(adk_thread.create_thread)
+            runners.create_thread = _wrap_create_thread(runners.create_thread)
+            logger.debug("ADK thread bridge patching successful")
+        except Exception as e:
+            logger.warning(f"Failed to patch ADK thread bridge: {e}")
+
         # Try to patch McpTool if available (MCP is optional)
         try:
             from google.adk.tools.mcp_tool import mcp_tool
@@ -75,6 +85,22 @@ def setup_adk(
         return False
 
 
+def _wrap_create_thread(create_thread):
+    if _is_patched(create_thread):
+        return create_thread
+
+    def _wrapped_create_thread(target: Any, *args: Any, **kwargs: Any):
+        ctx = contextvars.copy_context()
+
+        def _run_in_context(*target_args: Any, **target_kwargs: Any):
+            return ctx.run(target, *target_args, **target_kwargs)
+
+        return create_thread(_run_in_context, *args, **kwargs)
+
+    _wrapped_create_thread._braintrust_patched = True
+    return _wrapped_create_thread
+
+
 def wrap_agent(Agent: Any) -> Any:
     if _is_patched(Agent):
         return Agent
diff --git a/py/src/braintrust/wrappers/adk/test_adk.py b/py/src/braintrust/wrappers/adk/test_adk.py
@@ -7,7 +7,7 @@
 from braintrust.bt_json import bt_safe_deep_copy
 from braintrust.logger import Attachment
 from braintrust.test_helpers import init_test_logger
-from braintrust.wrappers.adk import setup_adk
+from braintrust.wrappers.adk import _wrap_create_thread, setup_adk
 from google.adk import Agent
 
 ADK_VERSION = tuple(int(x) for x in pkg_version("google-adk").split(".")[:3])
@@ -50,6 +50,84 @@ def memory_logger():
         yield bgl
 
 
+def test_adk_thread_context_propagation(memory_logger):
+    """Runner.run should preserve Braintrust context across its thread bridge."""
+    import asyncio
+
+    from braintrust import current_span, start_span
+    from google.adk.agents import LlmAgent
+    from google.adk.models.base_llm import BaseLlm
+    from google.adk.models.llm_request import LlmRequest
+    from google.adk.models.llm_response import LlmResponse
+    from google.adk.models.registry import LLMRegistry
+    from google.adk.runners import Runner
+    from google.adk.sessions import InMemorySessionService
+    from google.genai import types
+
+    assert not memory_logger.pop()
+
+    parent_seen = []
+
+    class TestLlm(BaseLlm):
+        @classmethod
+        def supported_models(cls) -> list[str]:
+            return [r"test-llm-context-prop"]
+
+        async def generate_content_async(self, llm_request: LlmRequest, stream: bool = False):
+            parent_seen.append(current_span())
+            yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="ok")]))
+
+    LLMRegistry.register(TestLlm)
+
+    agent = LlmAgent(
+        name="echo_agent",
+        model="test-llm-context-prop",
+        instruction="Respond with ok.",
+    )
+    session_service = InMemorySessionService()
+    app_name = "thread_bridge_app"
+    user_id = "test-user"
+    session_id = "test-session-thread"
+    asyncio.run(
+        session_service.create_session(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+        )
+    )
+    runner = Runner(agent=agent, app_name=app_name, session_service=session_service)
+    user_msg = types.Content(role="user", parts=[types.Part(text="hello")])
+
+    with start_span(name="adk_thread_parent") as parent_span:
+        events = list(runner.run(user_id=user_id, session_id=session_id, new_message=user_msg))
+
+    assert events
+    assert parent_seen
+    thread_root = getattr(parent_seen[0], "root_span_id", None)
+    assert thread_root is not None
+    assert thread_root == parent_span.root_span_id
+
+
+def test_wrap_create_thread_exception_does_not_double_invoke_target():
+    """Regression test: target exceptions must not cause a second invocation."""
+    call_count = 0
+
+    def create_thread(target, *args, **kwargs):
+        return target(*args, **kwargs)
+
+    wrapped_create_thread = _wrap_create_thread(create_thread)
+
+    def target():
+        nonlocal call_count
+        call_count += 1
+        raise RuntimeError("boom")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        wrapped_create_thread(target)
+
+    assert call_count == 1
+
+
 @pytest.mark.vcr
 @pytest.mark.asyncio
 async def test_adk_braintrust_integration(memory_logger):
diff --git a/py/src/braintrust/wrappers/cassettes/test_model_request_stream_sync_thread_context_propagation.yaml b/py/src/braintrust/wrappers/cassettes/test_model_request_stream_sync_thread_context_propagation.yaml
@@ -0,0 +1,135 @@
+interactions:
+- request:
+    body: '{"messages":[{"role":"user","content":"Hello"}],"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true}}'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '136'
+      Content-Type:
+      - application/json
+      Cookie:
+      - __cf_bm=SnTjRxRlQXUGj3jo4FeKDaKGiTqzcsdhS_LPeF2hwmg-1766265191-1.0.1.1-aN3W.WEa6IjC4znMp0fTYaZOZeW.3xU8ex4fRGFibZ0fuR86Piz9hDDAEvh8Z9ewQ9WHoKwdnsDvPE5KIpoZY4xurXV7VS4TcwfukWZofIk;
+        _cfuvid=tXOZ7vGE2DBF6L6fDg_veKtSaUVC4UPotJDezWYoYXI-1766265191281-0.0.1.1-604800000
+      Host:
+      - api.openai.com
+      User-Agent:
+      - pydantic-ai/1.37.0
+      X-Stainless-Arch:
+      - arm64
+      X-Stainless-Async:
+      - async:asyncio
+      X-Stainless-Lang:
+      - python
+      X-Stainless-OS:
+      - MacOS
+      X-Stainless-Package-Version:
+      - 2.14.0
+      X-Stainless-Runtime:
+      - CPython
+      X-Stainless-Runtime-Version:
+      - 3.13.3
+      x-stainless-read-timeout:
+      - '600'
+      x-stainless-retry-count:
+      - '0'
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: 'data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"n4wIfDywj"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"1"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"M7nBs8tJ4w"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":","},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"3AKnil79q5"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"
+        "},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"tj0YRKsPzd"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"2"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"3G9IJGTD7B"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":","},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"ury0XoZd21"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"
+        "},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"mJX6hpQnaM"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"3"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"qhlqTCCpyt"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"2JL0YHl8QJ"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null,"obfuscation":"DgeyW"}
+
+
+        data: {"id":"chatcmpl-CoyYZrJy9JYFY664IA2WgzVOsQSmj","object":"chat.completion.chunk","created":1766265215,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_ee69c2ef48","choices":[],"usage":{"prompt_tokens":14,"completion_tokens":8,"total_tokens":22,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"obfuscation":"pdj6Hh5uLUK"}
+
+
+        data: [DONE]
+
+
+        '
+    headers:
+      CF-RAY:
+      - 9b1224f9cc5cdbe5-LAX
+      Connection:
+      - keep-alive
+      Content-Type:
+      - text/event-stream; charset=utf-8
+      Date:
+      - Sat, 20 Dec 2025 21:13:35 GMT
+      Server:
+      - cloudflare
+      Strict-Transport-Security:
+      - max-age=31536000; includeSubDomains; preload
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      cf-cache-status:
+      - DYNAMIC
+      openai-organization:
+      - braintrust-data
+      openai-processing-ms:
+      - '159'
+      openai-project:
+      - proj_vsCSXafhhByzWOThMrJcZiw9
+      openai-version:
+      - '2020-10-01'
+      x-envoy-upstream-service-time:
+      - '171'
+      x-openai-proxy-wasm:
+      - v0.1
+      x-ratelimit-limit-requests:
+      - '30000'
+      x-ratelimit-limit-tokens:
+      - '150000000'
+      x-ratelimit-remaining-requests:
+      - '29999'
+      x-ratelimit-remaining-tokens:
+      - '149999992'
+      x-ratelimit-reset-requests:
+      - 2ms
+      x-ratelimit-reset-tokens:
+      - 0s
+      x-request-id:
+      - req_fac92d6c619a4d7e94154752dd20d282
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/py/src/braintrust/wrappers/pydantic_ai.py b/py/src/braintrust/wrappers/pydantic_ai.py
@@ -1,4 +1,5 @@
 import asyncio
+import contextvars
 import logging
 import sys
 import time
@@ -50,6 +51,16 @@ def setup_pydantic_ai(
 
         wrap_model_classes()
 
+        # Patch StreamedResponseSync to propagate context to background threads
+        try:
+            if hasattr(direct_module, "StreamedResponseSync"):
+                wrap_function_wrapper(
+                    direct_module.StreamedResponseSync, "_start_producer", _create_start_producer_wrapper()
+                )
+                logger.debug("Pydantic AI StreamedResponseSync context propagation patching successful")
+        except Exception as e:
+            logger.warning(f"Failed to patch StreamedResponseSync context propagation: {e}")
+
         return True
     except ImportError:
         # Not installed - this is expected when using auto_instrument()
@@ -1230,6 +1241,30 @@ def _extract_response_metrics(
     return metrics if metrics else None
 
 
+def _create_start_producer_wrapper():
+    """Create wrapper for StreamedResponseSync._start_producer to propagate context.
+
+    StreamedResponseSync._start_producer creates a background thread that doesn't
+    inherit contextvars. This wrapper ensures Braintrust context flows to that thread
+    so nested instrumentation (like wrap_openai) creates properly parented spans.
+    """
+
+    def wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> None:
+        ctx = contextvars.copy_context()
+        original_async_producer = instance._async_producer
+
+        def _context_wrapped_async_producer() -> None:
+            ctx.run(original_async_producer)
+
+        instance._async_producer = _context_wrapped_async_producer
+        try:
+            return wrapped(*args, **kwargs)
+        finally:
+            instance._async_producer = original_async_producer
+
+    return wrapper
+
+
 def _is_patched(obj: Any) -> bool:
     """Check if object is already patched.
 
diff --git a/py/src/braintrust/wrappers/test_pydantic_ai_integration.py b/py/src/braintrust/wrappers/test_pydantic_ai_integration.py
