feat(integrations): implement context management for invoke_agent spans (#5089)

- Introduced context variables to manage nested invoke_agent spans
safely.
- Added functions to push and pop spans from the context stack.
- Updated existing code to utilize context variables instead of Sentry
scope for agent management.
- Enhanced execute_tool_span to support parent-child relationships
between spans.

This change improves the handling of agent spans during nested calls,
ensuring better traceability and isolation of spans in asynchronous
contexts.

### Description
<!-- What changed and why? -->

#### Issues
* Closes
https://linear.app/getsentry/issue/TET-1388/pydantic-ai-replace-current-scope-context-with-contextvar
* Closes
https://linear.app/getsentry/issue/TET-1373/pydantic-ai-wrong-hierarchy-of-execute-tool-spans

#### Reminders
- Please add tests to validate your changes, and lint your code using
`tox -e linters`.
- Add GH Issue ID _&_ Linear ID (if applicable)
- PR title should use [conventional
commit](https://develop.sentry.dev/engineering-practices/commit-messages/#type)
style (`feat:`, `fix:`, `ref:`, `meta:`)
- For external contributors:
[CONTRIBUTING.md](https://github.com/getsentry/sentry-python/blob/master/CONTRIBUTING.md),
[Sentry SDK development docs](https://develop.sentry.dev/sdk/), [Discord
community](https://discord.gg/Ww9hbqr)

Author: constantinius

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
 

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

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)

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"):

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"

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