fix(pydantic-ai): Remove agent stack context variable

Author: alexander-alderman-webb

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, self.is_streaming)
-
         # 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, is_streaming)
-
                 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
 

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:

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:

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

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", is_streaming: bool = False) -> 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() -> 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]["agent"]
-    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 = {