diff --git a/sentry_sdk/integrations/openai_agents/_context_vars.py b/sentry_sdk/integrations/openai_agents/_context_vars.py deleted file mode 100644 index 83746ecee6..0000000000 --- a/sentry_sdk/integrations/openai_agents/_context_vars.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -Context variables for passing data between nested calls in the OpenAI Agents integration. -""" - -from contextvars import ContextVar - -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - pass - -# Context variable to pass response model between nested calls (for gen_ai.chat spans) -_response_model_context = ContextVar("openai_agents_response_model", default=None) # type: ContextVar[str | None] - -# Context variable to store the last response model for invoke_agent spans -_invoke_agent_response_model_context = ContextVar( - "openai_agents_invoke_agent_response_model", default=None -) # type: ContextVar[str | None] diff --git a/sentry_sdk/integrations/openai_agents/patches/agent_run.py b/sentry_sdk/integrations/openai_agents/patches/agent_run.py index b25bf82ad5..57a68f2f5d 100644 --- a/sentry_sdk/integrations/openai_agents/patches/agent_run.py +++ b/sentry_sdk/integrations/openai_agents/patches/agent_run.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) diff --git a/sentry_sdk/integrations/openai_agents/patches/models.py b/sentry_sdk/integrations/openai_agents/patches/models.py index aa6371302e..c7b8ed1ca5 100644 --- a/sentry_sdk/integrations/openai_agents/patches/models.py +++ b/sentry_sdk/integrations/openai_agents/patches/models.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) diff --git a/sentry_sdk/integrations/openai_agents/patches/runner.py b/sentry_sdk/integrations/openai_agents/patches/runner.py index 745f30a38e..05c15da4d1 100644 --- a/sentry_sdk/integrations/openai_agents/patches/runner.py +++ b/sentry_sdk/integrations/openai_agents/patches/runner.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 diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py index 9ae18c0451..5d1731f247 100644 --- a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py +++ b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.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