diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 5e0515b3c3..5b58d8f128 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -3,7 +3,7 @@ import sentry_sdk from ..spans import invoke_agent_span, update_invoke_agent_span -from ..utils import _capture_exception +from ..utils import _capture_exception, pop_agent, push_agent from typing import TYPE_CHECKING from pydantic_ai.agent import Agent # type: ignore @@ -41,17 +41,20 @@ async def __aenter__(self): self._isolation_scope = sentry_sdk.isolation_scope() self._isolation_scope.__enter__() - # Store agent reference and streaming flag - sentry_sdk.get_current_scope().set_context( - "pydantic_ai_agent", {"_agent": self.agent, "_streaming": self.is_streaming} - ) - # Create invoke_agent span (will be closed in __aexit__) self._span = invoke_agent_span( - self.user_prompt, self.agent, self.model, self.model_settings + self.user_prompt, + self.agent, + self.model, + self.model_settings, + self.is_streaming, ) 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, self.is_streaming) + # Enter the original context manager result = await self.original_ctx_manager.__aenter__() self._result = result @@ -71,7 +74,9 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if self._span is not None: update_invoke_agent_span(self._span, output) finally: - sentry_sdk.get_current_scope().remove_context("pydantic_ai_agent") + # Pop agent from contextvar stack + pop_agent() + # Clean up invoke span if self._span: self._span.__exit__(exc_type, exc_val, exc_tb) @@ -97,19 +102,19 @@ async def wrapper(self, *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(): - # Store agent reference and streaming flag in Sentry scope for access in nested spans - # We store the full agent to allow access to tools and system prompts - sentry_sdk.get_current_scope().set_context( - "pydantic_ai_agent", {"_agent": self, "_streaming": is_streaming} - ) - # Extract parameters for the span user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) model = kwargs.get("model") model_settings = kwargs.get("model_settings") # Create invoke_agent span - with invoke_agent_span(user_prompt, self, model, model_settings) as span: + 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, is_streaming) + try: result = await original_func(self, *args, **kwargs) @@ -122,7 +127,8 @@ async def wrapper(self, *args, **kwargs): _capture_exception(exc) raise exc from None finally: - sentry_sdk.get_current_scope().remove_context("pydantic_ai_agent") + # 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 671b00ec95..1940be811f 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -5,7 +5,10 @@ import sentry_sdk from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import _capture_exception +from ..utils import ( + _capture_exception, + get_current_agent, +) from typing import TYPE_CHECKING @@ -49,29 +52,43 @@ async def wrapped_call_tool(self, call, *args, **kwargs): if tool and HAS_MCP and isinstance(tool.toolset, MCPServer): tool_type = "mcp" - # Get agent from Sentry scope - current_span = sentry_sdk.get_current_span() - if current_span and tool: - agent_data = ( - sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} - ) - agent = agent_data.get("_agent") + # Get agent from contextvar + agent = get_current_agent() + if agent and tool: try: args_dict = call.args_as_dict() except Exception: args_dict = call.args if isinstance(call.args, dict) else {} - with execute_tool_span(name, args_dict, agent, tool_type=tool_type) as span: - try: - result = await original_call_tool(self, call, *args, **kwargs) - update_execute_tool_span(span, result) - return result - except Exception as exc: - _capture_exception(exc) - raise exc from None + # Create execute_tool span + # Nesting is handled by isolation_scope() to ensure proper parent-child relationships + with sentry_sdk.isolation_scope(): + with execute_tool_span( + name, + args_dict, + agent, + tool_type=tool_type, + ) as span: + try: + result = await original_call_tool( + self, + call, + *args, + **kwargs, + ) + update_execute_tool_span(span, result) + return result + except Exception as exc: + _capture_exception(exc) + raise exc from None # No span context - just call original - return await original_call_tool(self, call, *args, **kwargs) + return await original_call_tool( + self, + call, + *args, + **kwargs, + ) ToolManager._call_tool = wrapped_call_tool diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 735e814acd..a2bd0272d4 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -10,6 +10,8 @@ _set_model_data, _should_send_prompts, _get_model_name, + get_current_agent, + get_is_streaming, ) from typing import TYPE_CHECKING @@ -216,20 +218,11 @@ def ai_client_span(messages, agent, model, model_settings): _set_agent_data(span, agent) _set_model_data(span, model, model_settings) - # Set streaming flag - agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} - is_streaming = agent_data.get("_streaming", False) - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming) + # Set streaming flag from contextvar + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) # Add available tools if agent is available - agent_obj = agent - if not agent_obj: - # Try to get from Sentry scope - agent_data = ( - sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} - ) - agent_obj = agent_data.get("_agent") - + agent_obj = agent or get_current_agent() _set_available_tools(span, agent_obj) # Set input messages (full conversation history) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 2094c53a40..329895778d 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any + from typing import Any, Optional def execute_tool_span(tool_name, tool_args, agent, tool_type="function"): diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index d6fb86918c..f5e22fb346 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -16,8 +16,8 @@ from typing import Any -def invoke_agent_span(user_prompt, agent, model, model_settings): - # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span +def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=False): + # type: (Any, Any, Any, Any, bool) -> sentry_sdk.tracing.Span """Create a span for invoking the agent.""" # Determine agent name for span name = "agent" diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index a7f5290d50..532fb7ddb6 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -1,4 +1,5 @@ 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 @@ -7,7 +8,47 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any + 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("pydantic_ai_agent_context_stack", default=[]) # type: ContextVar[list[dict[str, Any]]] + + +def push_agent(agent, is_streaming=False): + # type: (Any, bool) -> None + """Push an agent context onto the stack along with its streaming flag.""" + stack = _agent_context_stack.get().copy() + stack.append({"agent": agent, "is_streaming": is_streaming}) + _agent_context_stack.set(stack) + + +def pop_agent(): + # type: () -> 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(): + # type: () -> Any + """Get the current agent from the contextvar stack.""" + stack = _agent_context_stack.get() + if stack: + return stack[-1]["agent"] + return None + + +def get_is_streaming(): + # type: () -> bool + """Get the streaming flag from the contextvar stack.""" + stack = _agent_context_stack.get() + if stack: + return stack[-1].get("is_streaming", False) + return False def _should_send_prompts(): @@ -37,23 +78,20 @@ def _set_agent_data(span, agent): Args: span: The span to set data on - agent: Agent object (can be None, will try to get from Sentry scope if not provided) + agent: Agent object (can be None, will try to get from contextvar if not provided) """ - # Extract agent name from agent object or Sentry scope + # Extract agent name from agent object or contextvar agent_obj = agent if not agent_obj: - # Try to get from Sentry scope - agent_data = ( - sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} - ) - agent_obj = agent_data.get("_agent") + # 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) def _get_model_name(model_obj): - # type: (Any) -> str | None + # type: (Any) -> Optional[str] """Extract model name from a model object. Args: @@ -87,9 +125,8 @@ def _set_model_data(span, model, model_settings): 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 Sentry scope if we need it - agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {} - agent_obj = agent_data.get("_agent") + # Try to get agent from contextvar if we need it + agent_obj = get_current_agent() # Extract model information model_obj = model