diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index cc6bd602b0..928505a0c6 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -6,7 +6,7 @@ from sentry_sdk.utils import capture_internal_exceptions, reraise from ..spans import invoke_agent_span, update_invoke_agent_span -from ..utils import _capture_exception, pop_agent, push_agent +from ..utils import _capture_exception from typing import TYPE_CHECKING @@ -56,10 +56,6 @@ async def __aenter__(self) -> "Any": ) self._span.__enter__() - # Push agent to contextvar stack after span is successfully created and entered - # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur - push_agent(self.agent) - # Enter the original context manager result = await self.original_ctx_manager.__aenter__() self._result = result @@ -74,9 +70,6 @@ async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> Non if exc_type is None and self._result and self._span is not None: update_invoke_agent_span(self._span, self._result) finally: - # Pop agent from contextvar stack - pop_agent() - # Clean up invoke span if self._span: self._span.__exit__(exc_type, exc_val, exc_tb) @@ -111,10 +104,6 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": with invoke_agent_span( user_prompt, self, model, model_settings, is_streaming ) as span: - # Push agent to contextvar stack after span is successfully created and entered - # This ensures proper pairing with pop_agent() in finally even if exceptions occur - push_agent(self) - try: result = await original_func(self, *args, **kwargs) @@ -127,9 +116,6 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": with capture_internal_exceptions(): _capture_exception(exc) reraise(*exc_info) - finally: - # Pop agent from contextvar stack - pop_agent() return wrapper diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 1f5cde8742..67499a0804 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -6,7 +6,7 @@ from sentry_sdk.utils import capture_internal_exceptions, reraise from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import _capture_exception, get_current_agent +from ..utils import _capture_exception from typing import TYPE_CHECKING @@ -57,10 +57,7 @@ async def wrapped_execute_tool_call( if tool and HAS_MCP and isinstance(tool.toolset, MCPServer): tool_type = "mcp" - # Get agent from contextvar - agent = get_current_agent() - - if agent and tool: + if tool: try: args_dict = call.args_as_dict() except Exception: @@ -72,7 +69,7 @@ async def wrapped_execute_tool_call( with execute_tool_span( name, args_dict, - agent, + validated.ctx.agent, tool_type=tool_type, tool_definition=selected_tool_definition, ) as span: @@ -136,10 +133,7 @@ async def wrapped_call_tool( if tool and HAS_MCP and isinstance(tool.toolset, MCPServer): tool_type = "mcp" - # Get agent from contextvar - agent = get_current_agent() - - if agent and tool: + if tool: try: args_dict = call.args_as_dict() except Exception: @@ -151,7 +145,7 @@ async def wrapped_call_tool( with execute_tool_span( name, args_dict, - agent, + call.ctx.agent, tool_type=tool_type, tool_definition=selected_tool_definition, ) as span: diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index e376f8097c..5e35d0b508 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -16,7 +16,6 @@ _set_model_data, _should_send_prompts, _get_model_name, - get_current_agent, ) from .utils import ( _serialize_binary_content_item, @@ -261,11 +260,9 @@ def ai_client_span( span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") _set_agent_data(span, agent) - _set_model_data(span, model, model_settings) + _set_model_data(span, agent, model, model_settings) - # Add available tools if agent is available - agent_obj = agent or get_current_agent() - _set_available_tools(span, agent_obj) + _set_available_tools(span, agent) # Set input messages (full conversation history) if messages: diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index ee08ca7036..248157b5d7 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -54,7 +54,7 @@ def invoke_agent_span( span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") _set_agent_data(span, agent) - _set_model_data(span, model, model_settings) + _set_model_data(span, agent, model, model_settings) _set_available_tools(span, agent) # Add user prompt and system prompts if available and prompts are enabled diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index c6dd658d20..3019249f02 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -1,5 +1,4 @@ import sentry_sdk -from contextvars import ContextVar from sentry_sdk.consts import SPANDATA from sentry_sdk.scope import should_send_default_pii from sentry_sdk.tracing_utils import set_span_errored @@ -11,36 +10,6 @@ from typing import Any, Optional -# Store the current agent context in a contextvar for re-entrant safety -# Using a list as a stack to support nested agent calls -_agent_context_stack: "ContextVar[list[dict[str, Any]]]" = ContextVar( - "pydantic_ai_agent_context_stack", default=[] -) - - -def push_agent(agent: "Any") -> None: - """Push an agent context onto the stack.""" - stack = _agent_context_stack.get().copy() - stack.append(agent) - _agent_context_stack.set(stack) - - -def pop_agent() -> None: - """Pop an agent context from the stack.""" - stack = _agent_context_stack.get().copy() - if stack: - stack.pop() - _agent_context_stack.set(stack) - - -def get_current_agent() -> "Any": - """Get the current agent from the contextvar stack.""" - stack = _agent_context_stack.get() - if stack: - return stack[-1] - return None - - def _should_send_prompts() -> bool: """ Check if prompts should be sent to Sentry. @@ -66,16 +35,10 @@ def _set_agent_data(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: Args: span: The span to set data on - agent: Agent object (can be None, will try to get from contextvar if not provided) + agent: Agent object """ - # Extract agent name from agent object or contextvar - agent_obj = agent - if not agent_obj: - # Try to get from contextvar - agent_obj = get_current_agent() - - if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + if agent and hasattr(agent, "name") and agent.name: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent.name) def _get_model_name(model_obj: "Any") -> "Optional[str]": @@ -104,7 +67,7 @@ def _get_model_name(model_obj: "Any") -> "Optional[str]": def _set_model_data( - span: "sentry_sdk.tracing.Span", model: "Any", model_settings: "Any" + span: "sentry_sdk.tracing.Span", agent: "Any", model: "Any", model_settings: "Any" ) -> None: """Set model-related data on a span. @@ -113,13 +76,10 @@ def _set_model_data( model: Model object (can be None, will try to get from agent if not provided) model_settings: Model settings (can be None, will try to get from agent if not provided) """ - # Try to get agent from contextvar if we need it - agent_obj = get_current_agent() - # Extract model information model_obj = model - if not model_obj and agent_obj and hasattr(agent_obj, "model"): - model_obj = agent_obj.model + if not model_obj and agent and hasattr(agent, "model"): + model_obj = agent.model if model_obj: # Set system from model @@ -133,8 +93,8 @@ def _set_model_data( # Extract model settings settings = model_settings - if not settings and agent_obj and hasattr(agent_obj, "model_settings"): - settings = agent_obj.model_settings + if not settings and agent and hasattr(agent, "model_settings"): + settings = agent.model_settings if settings: settings_map = {