fix: use isolation_scopes to keep execute_tool spans separate

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, push_invoke_agent_span, pop_invoke_agent_span
+from ..utils import _capture_exception, pop_agent
 
 from typing import TYPE_CHECKING
 from pydantic_ai.agent import Agent  # type: ignore
@@ -70,8 +70,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
                 if self._span is not None:
                     update_invoke_agent_span(self._span, output)
         finally:
-            # Pop span from contextvar stack
-            pop_invoke_agent_span()
+            # Pop agent from contextvar stack
+            pop_agent()
 
             # Clean up invoke span
             if self._span:
@@ -119,8 +119,8 @@ async def wrapper(self, *args, **kwargs):
                     _capture_exception(exc)
                     raise exc from None
                 finally:
-                    # Pop span from contextvar stack
-                    pop_invoke_agent_span()
+                    # Pop agent from contextvar stack
+                    pop_agent()
 
     return wrapper
 

sentry_sdk/integrations/pydantic_ai/patches/tools.py

@@ -8,7 +8,6 @@
 from ..utils import (
     _capture_exception,
     get_current_agent,
-    get_current_invoke_agent_span,
 )
 
 from typing import TYPE_CHECKING
@@ -53,30 +52,33 @@ 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 and invoke_agent span from contextvar
+        # Get agent from contextvar
         agent = get_current_agent()
-        invoke_span = get_current_invoke_agent_span()
 
-        if invoke_span and tool:
+        if agent and tool:
             try:
                 args_dict = call.args_as_dict()
             except Exception:
                 args_dict = call.args if isinstance(call.args, dict) else {}
 
-            # 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
-                    )
-                    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, allow_partial, wrap_validation_errors
+                        )
+                        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(

sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py

@@ -11,35 +11,21 @@
     from typing import Any, Optional
 
 
-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
+def execute_tool_span(tool_name, tool_args, agent, tool_type="function"):
+    # type: (str, Any, Any, str) -> 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.
     """
-    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 = 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/spans/invoke_agent.py

@@ -8,7 +8,7 @@
     _set_available_tools,
     _set_model_data,
     _should_send_prompts,
-    push_invoke_agent_span,
+    push_agent,
 )
 
 from typing import TYPE_CHECKING
@@ -33,9 +33,9 @@ def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=Fa
 
     span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
 
-    # Push span and agent to contextvar stack immediately after span creation
+    # Push agent to contextvar stack immediately after span creation
     # This ensures the agent is available in get_current_agent() before _set_model_data is called
-    push_invoke_agent_span(span, agent, is_streaming)
+    push_agent(agent, is_streaming)
 
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -11,41 +11,32 @@
     from typing import Any, Optional
 
 
-# Store the current invoke_agent span in a contextvar for re-entrant safety
+# Store the current agent context 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]]]
+_agent_context_stack = ContextVar("pydantic_ai_agent_context_stack", default=[])  # type: ContextVar[list[dict[str, Any]]]
 
 
-def push_invoke_agent_span(span, agent, is_streaming=False):
-    # type: (sentry_sdk.tracing.Span, Any, bool) -> None
-    """Push an invoke_agent span onto the stack along with its agent and streaming flag."""
-    stack = _invoke_agent_span_stack.get().copy()
-    stack.append({"span": span, "agent": agent, "is_streaming": is_streaming})
-    _invoke_agent_span_stack.set(stack)
+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_invoke_agent_span():
+def pop_agent():
     # type: () -> None
-    """Pop an invoke_agent span from the stack."""
-    stack = _invoke_agent_span_stack.get().copy()
+    """Pop an agent context from the stack."""
+    stack = _agent_context_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
+    _agent_context_stack.set(stack)
 
 
 def get_current_agent():
     # type: () -> Any
     """Get the current agent from the contextvar stack."""
-    stack = _invoke_agent_span_stack.get()
+    stack = _agent_context_stack.get()
     if stack:
         return stack[-1]["agent"]
     return None
@@ -54,7 +45,7 @@ def get_current_agent():
 def get_is_streaming():
     # type: () -> bool
     """Get the streaming flag from the contextvar stack."""
-    stack = _invoke_agent_span_stack.get()
+    stack = _agent_context_stack.get()
     if stack:
         return stack[-1].get("is_streaming", False)
     return False