fix(openai-agents): Remove context variables (#5184)

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/_context_vars.py

@@ -8,6 +8,8 @@
 if TYPE_CHECKING:
     from typing import Any, Optional
 
+    from sentry_sdk.tracing import Span
+
 try:
     import agents
 except ImportError:
@@ -27,13 +29,15 @@ def _patch_agent_run():
     original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
 
     def _start_invoke_agent_span(context_wrapper, agent, kwargs):
-        # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> None
+        # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> Span
         """Start an agent invocation span"""
         # Store the agent on the context wrapper so we can access it later
         context_wrapper._sentry_current_agent = agent
         span = invoke_agent_span(context_wrapper, agent, kwargs)
         context_wrapper._sentry_agent_span = span
 
+        return span
+
     def _end_invoke_agent_span(context_wrapper, agent, output=None):
         # type: (agents.RunContextWrapper, agents.Agent, Optional[Any]) -> None
         """End the agent invocation span"""
@@ -73,7 +77,8 @@ async def patched_run_single_turn(cls, *args, **kwargs):
                 if current_agent and current_agent != agent:
                     _end_invoke_agent_span(context_wrapper, current_agent)
 
-            _start_invoke_agent_span(context_wrapper, agent, kwargs)
+            span = _start_invoke_agent_span(context_wrapper, agent, kwargs)
+            agent._sentry_agent_span = span
 
         # Call original method with all the correct parameters
         result = await original_run_single_turn(*args, **kwargs)

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -2,11 +2,8 @@
 
 from sentry_sdk.integrations import DidNotEnable
 
-from .._context_vars import (
-    _invoke_agent_response_model_context,
-    _response_model_context,
-)
 from ..spans import ai_client_span, update_ai_client_span
+from sentry_sdk.consts import SPANDATA
 
 from typing import TYPE_CHECKING
 
@@ -47,7 +44,7 @@ async def wrapped_fetch_response(*args, **kwargs):
                 response = await original_fetch_response(*args, **kwargs)
                 # Store model from raw response in context variable
                 if hasattr(response, "model"):
-                    _response_model_context.set(str(response.model))
+                    agent._sentry_raw_response_model = str(response.model)
                 return response
 
             model._fetch_response = wrapped_fetch_response
@@ -59,13 +56,15 @@ async def wrapped_get_response(*args, **kwargs):
                 result = await original_get_response(*args, **kwargs)
 
                 # Retrieve response model from context and attach to ModelResponse
-                response_model = _response_model_context.get(None)
+                response_model = getattr(agent, "_sentry_raw_response_model", None)
                 if response_model:
                     result._sentry_response_model = response_model
-                    _response_model_context.set(None)  # Clear context
 
-                    # Also store for invoke_agent span (will be the last one used)
-                    _invoke_agent_response_model_context.set(response_model)
+                    agent_span = getattr(agent, "_sentry_agent_span", None)
+                    if agent_span:
+                        agent_span.set_data(
+                            SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
+                        )
 
                 update_ai_client_span(span, agent, kwargs, result)
 

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -26,9 +26,11 @@ async def wrapper(*args, **kwargs):
         # Isolate each workflow so that when agents are run in asyncio tasks they
         # don't touch each other's scopes
         with sentry_sdk.isolation_scope():
-            agent = args[0]
+            # Clone agent because agent invocation spans are attached per run.
+            agent = args[0].clone()
             with agent_workflow_span(agent):
                 result = None
+                args = (agent, *args[1:])
                 try:
                     result = await original_func(*args, **kwargs)
                     return result

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -9,7 +9,6 @@
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.utils import safe_serialize
 
-from .._context_vars import _invoke_agent_response_model_context
 from ..consts import SPAN_ORIGIN
 from ..utils import _set_agent_data, _set_usage_data
 
@@ -89,12 +88,6 @@ def update_invoke_agent_span(context, agent, output):
         if hasattr(context, "usage"):
             _set_usage_data(span, context.usage)
 
-        # Add response model if available (will be the last model used)
-        response_model = _invoke_agent_response_model_context.get(None)
-        if response_model:
-            span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
-            _invoke_agent_response_model_context.set(None)  # Clear after use
-
         if should_send_default_pii():
             set_data_normalized(
                 span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False