feat(integrations): implement context management for invoke_agent spans

- 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.

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, push_invoke_agent_span, pop_invoke_agent_span
 
 from typing import TYPE_CHECKING
 from pydantic_ai.agent import Agent  # type: ignore
@@ -41,17 +41,15 @@ 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._span.__enter__()
 
+        # Push span and agent to contextvar stack
+        push_invoke_agent_span(self._span, self.agent)
+
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
         self._result = result
@@ -71,7 +69,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 span from contextvar stack
+            pop_invoke_agent_span()
+
             # Clean up invoke span
             if self._span:
                 self._span.__exit__(exc_type, exc_val, exc_tb)
@@ -97,19 +97,15 @@ 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:
+                # Push span and agent to contextvar stack
+                push_invoke_agent_span(span, self)
                 try:
                     result = await original_func(self, *args, **kwargs)
 
@@ -122,7 +118,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 span from contextvar stack
+                    pop_invoke_agent_span()
 
     return wrapper
 

sentry_sdk/integrations/pydantic_ai/patches/tools.py

@@ -5,7 +5,11 @@
 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,
+    get_current_invoke_agent_span,
+)
 
 from typing import TYPE_CHECKING
 
@@ -49,20 +53,21 @@ async def wrapped_call_tool(self, call, allow_partial, wrap_validation_errors):
         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 and invoke_agent span from contextvar
+        agent = get_current_agent()
+        invoke_span = get_current_invoke_agent_span()
 
+        if invoke_span 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:
+            # Create execute_tool span as a child of invoke_agent span
+            # Passing parent_span ensures parallel tools are siblings under the same parent
+            with execute_tool_span(
+                name, args_dict, agent, tool_type=tool_type, parent_span=invoke_span
+            ) as span:
                 try:
                     result = await original_call_tool(
                         self, call, allow_partial, wrap_validation_errors

sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py

@@ -8,24 +8,38 @@
 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"):
-    # type: (str, Any, Any, str) -> sentry_sdk.tracing.Span
+def execute_tool_span(
+    tool_name, tool_args, agent, tool_type="function", parent_span=None
+):
+    # type: (str, Any, Any, str, Optional[sentry_sdk.tracing.Span]) -> sentry_sdk.tracing.Span
     """Create a span for tool execution.
 
     Args:
         tool_name: The name of the tool being executed
         tool_args: The arguments passed to the tool
         agent: The agent executing the tool
         tool_type: The type of tool ("function" for regular tools, "mcp" for MCP services)
+        parent_span: Optional parent span to create this as a child of. If provided,
+            uses parent_span.start_child() to ensure parallel tools are siblings.
     """
-    span = sentry_sdk.start_span(
-        op=OP.GEN_AI_EXECUTE_TOOL,
-        name=f"execute_tool {tool_name}",
-        origin=SPAN_ORIGIN,
-    )
+    if parent_span:
+        # Create as child of the specified parent span
+        # This ensures parallel tool calls are siblings under the same parent
+        span = parent_span.start_child(
+            op=OP.GEN_AI_EXECUTE_TOOL,
+            name=f"execute_tool {tool_name}",
+            origin=SPAN_ORIGIN,
+        )
+    else:
+        # Create as child of current span
+        span = sentry_sdk.start_span(
+            op=OP.GEN_AI_EXECUTE_TOOL,
+            name=f"execute_tool {tool_name}",
+            origin=SPAN_ORIGIN,
+        )
 
     span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
     span.set_data(SPANDATA.GEN_AI_TOOL_TYPE, tool_type)

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,48 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional
+
+
+# Store the current invoke_agent span in a contextvar for re-entrant safety
+# Using a list as a stack to support nested agent calls
+_invoke_agent_span_stack = ContextVar("pydantic_ai_invoke_agent_span_stack", default=[])
+# type: ContextVar[list[dict[str, Any]]]
+
+
+def push_invoke_agent_span(span, agent):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Push an invoke_agent span onto the stack along with its agent."""
+    stack = _invoke_agent_span_stack.get().copy()
+    stack.append({"span": span, "agent": agent})
+    _invoke_agent_span_stack.set(stack)
+
+
+def pop_invoke_agent_span():
+    # type: () -> None
+    """Pop an invoke_agent span from the stack."""
+    stack = _invoke_agent_span_stack.get().copy()
+    if stack:
+        stack.pop()
+    _invoke_agent_span_stack.set(stack)
+
+
+def get_current_invoke_agent_span():
+    # type: () -> Optional[sentry_sdk.tracing.Span]
+    """Get the current invoke_agent span (top of stack)."""
+    stack = _invoke_agent_span_stack.get()
+    if stack:
+        return stack[-1]["span"]
+    return None
+
+
+def get_current_agent():
+    # type: () -> Any
+    """Get the current agent from the contextvar stack."""
+    stack = _invoke_agent_span_stack.get()
+    if stack:
+        return stack[-1]["agent"]
+    return None
 
 
 def _should_send_prompts():
@@ -37,16 +79,13 @@ 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)
@@ -87,9 +126,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