From df06590cd73938e93551ca69886234d487c1ec92 Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Mon, 10 Nov 2025 17:12:29 +0100 Subject: [PATCH 1/8] 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. --- .../pydantic_ai/patches/agent_run.py | 25 ++++---- .../integrations/pydantic_ai/patches/tools.py | 23 ++++--- .../pydantic_ai/spans/execute_tool.py | 30 +++++++--- sentry_sdk/integrations/pydantic_ai/utils.py | 60 +++++++++++++++---- 4 files changed, 96 insertions(+), 42 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 5e0515b3c3..bd856f9e8d 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, 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,12 +97,6 @@ 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") @@ -110,6 +104,8 @@ async def wrapper(self, *args, **kwargs): # 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 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 25c2cd6afd..a571e337ab 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/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 diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 2094c53a40..b4d7ae324b 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -8,11 +8,13 @@ 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: @@ -20,12 +22,24 @@ def execute_tool_span(tool_name, tool_args, agent, tool_type="function"): 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) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index a7f5290d50..3a38d9dc62 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,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 From 45df059b09369fdc91f448524f1dde653bfb0d94 Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Tue, 11 Nov 2025 08:46:56 +0100 Subject: [PATCH 2/8] fix: correct the propagation of the "is_streaming" flag --- .../pydantic_ai/patches/agent_run.py | 4 ++-- .../integrations/pydantic_ai/spans/ai_client.py | 17 +++++------------ sentry_sdk/integrations/pydantic_ai/utils.py | 17 +++++++++++++---- 3 files changed, 20 insertions(+), 18 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index bd856f9e8d..c37b4b791e 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -48,7 +48,7 @@ async def __aenter__(self): self._span.__enter__() # Push span and agent to contextvar stack - push_invoke_agent_span(self._span, self.agent) + push_invoke_agent_span(self._span, self.agent, self.is_streaming) # Enter the original context manager result = await self.original_ctx_manager.__aenter__() @@ -105,7 +105,7 @@ async def wrapper(self, *args, **kwargs): # 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) + push_invoke_agent_span(span, self, is_streaming) try: result = await original_func(self, *args, **kwargs) 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/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 3a38d9dc62..3ab0c2aa3a 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -17,11 +17,11 @@ # 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.""" +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}) + stack.append({"span": span, "agent": agent, "is_streaming": is_streaming}) _invoke_agent_span_stack.set(stack) @@ -52,6 +52,15 @@ def get_current_agent(): return None +def get_is_streaming(): + # type: () -> bool + """Get the streaming flag from the contextvar stack.""" + stack = _invoke_agent_span_stack.get() + if stack: + return stack[-1].get("is_streaming", False) + return False + + def _should_send_prompts(): # type: () -> bool """ From 0224373caa93153432619a0a43efc5b662d5f3ce Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Tue, 11 Nov 2025 10:20:02 +0100 Subject: [PATCH 3/8] fix: type checking issue --- sentry_sdk/integrations/pydantic_ai/utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 3ab0c2aa3a..88a7362bf4 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -13,8 +13,7 @@ # 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]]] +_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, is_streaming=False): @@ -101,7 +100,7 @@ def _set_agent_data(span, agent): def _get_model_name(model_obj): - # type: (Any) -> str | None + # type: (Any) -> Optional[str] """Extract model name from a model object. Args: From c7b7ae73f7f99bba1bed98d9b434e786b3a758c8 Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Tue, 11 Nov 2025 10:49:08 +0100 Subject: [PATCH 4/8] fix: reshuffling where the _invoke_agent_span_stack context variable is set --- .../integrations/pydantic_ai/patches/agent_run.py | 15 ++++++++------- .../pydantic_ai/spans/invoke_agent.py | 9 +++++++-- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index c37b4b791e..3c2c54d946 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -43,13 +43,14 @@ async def __aenter__(self): # 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 span and agent to contextvar stack - push_invoke_agent_span(self._span, self.agent, self.is_streaming) - # Enter the original context manager result = await self.original_ctx_manager.__aenter__() self._result = result @@ -103,9 +104,9 @@ async def wrapper(self, *args, **kwargs): 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, is_streaming) + with invoke_agent_span( + user_prompt, self, model, model_settings, is_streaming + ) as span: try: result = await original_func(self, *args, **kwargs) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index d6fb86918c..c8c8ba57b5 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -8,6 +8,7 @@ _set_available_tools, _set_model_data, _should_send_prompts, + push_invoke_agent_span, ) from typing import TYPE_CHECKING @@ -16,8 +17,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" @@ -32,6 +33,10 @@ def invoke_agent_span(user_prompt, agent, model, model_settings): span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + # Push span and 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) + _set_agent_data(span, agent) _set_model_data(span, model, model_settings) _set_available_tools(span, agent) From 741e16f38b7a0b010e0a6639c7a8d52fea879316 Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Wed, 12 Nov 2025 14:30:40 +0100 Subject: [PATCH 5/8] fix: use isolation_scopes to keep execute_tool spans separate --- .../pydantic_ai/patches/agent_run.py | 10 ++--- .../integrations/pydantic_ai/patches/tools.py | 38 ++++++++++--------- .../pydantic_ai/spans/execute_tool.py | 28 ++++---------- .../pydantic_ai/spans/invoke_agent.py | 6 +-- sentry_sdk/integrations/pydantic_ai/utils.py | 37 +++++++----------- 5 files changed, 49 insertions(+), 70 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 3c2c54d946..5d61b842b9 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, 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 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index a571e337ab..aa017fab17 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/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( diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index b4d7ae324b..329895778d 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -11,10 +11,8 @@ 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: @@ -22,24 +20,12 @@ def execute_tool_span( 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) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index c8c8ba57b5..a09ad97198 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/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) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 88a7362bf4..532fb7ddb6 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/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 From 5408b8645a0c4807ef507d946218b770ba96a47d Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Thu, 13 Nov 2025 09:37:39 +0100 Subject: [PATCH 6/8] fix: update agent context management in invoke_agent spans --- .../integrations/pydantic_ai/patches/agent_run.py | 10 +++++++++- .../integrations/pydantic_ai/spans/invoke_agent.py | 5 ----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 5d61b842b9..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, pop_agent +from ..utils import _capture_exception, pop_agent, push_agent from typing import TYPE_CHECKING from pydantic_ai.agent import Agent # type: ignore @@ -51,6 +51,10 @@ async def __aenter__(self): ) 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 @@ -107,6 +111,10 @@ async def wrapper(self, *args, **kwargs): 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) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index a09ad97198..f5e22fb346 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -8,7 +8,6 @@ _set_available_tools, _set_model_data, _should_send_prompts, - push_agent, ) from typing import TYPE_CHECKING @@ -33,10 +32,6 @@ def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=Fa span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - # 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_agent(agent, is_streaming) - _set_agent_data(span, agent) _set_model_data(span, model, model_settings) _set_available_tools(span, agent) From 262c6c1a1f80ad1e8a7051e1144d3cc531874a57 Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Thu, 13 Nov 2025 11:12:51 +0100 Subject: [PATCH 7/8] fix: enhance wrapped_call_tool to include approved parameter --- .../integrations/pydantic_ai/patches/tools.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 03f3b2210f..30d01e032c 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -40,8 +40,10 @@ def _patch_tool_execution(): original_call_tool = ToolManager._call_tool @wraps(original_call_tool) - async def wrapped_call_tool(self, call, *args, **kwargs): - # type: (Any, Any, *Any, **Any) -> Any + async def wrapped_call_tool( + self, call, *, allow_partial, wrap_validation_errors, approved + ): + # type: (Any, Any, bool, bool, bool) -> Any # Extract tool info before calling original name = call.tool_name @@ -72,7 +74,11 @@ async def wrapped_call_tool(self, call, *args, **kwargs): ) as span: try: result = await original_call_tool( - self, call, wrap_validation_errors + self, + call, + allow_partial=allow_partial, + wrap_validation_errors=wrap_validation_errors, + approved=approved, ) update_execute_tool_span(span, result) return result @@ -81,6 +87,12 @@ async def wrapped_call_tool(self, call, *args, **kwargs): 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, + allow_partial=allow_partial, + wrap_validation_errors=wrap_validation_errors, + approved=approved, + ) ToolManager._call_tool = wrapped_call_tool From b3f7d05cc076468ce339d3a45f2e42144ddb8e0a Mon Sep 17 00:00:00 2001 From: Fabian Schindler Date: Thu, 13 Nov 2025 13:11:06 +0100 Subject: [PATCH 8/8] fix: errors for earlier versions of pydantic-ai --- .../integrations/pydantic_ai/patches/tools.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 30d01e032c..1940be811f 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -40,10 +40,8 @@ def _patch_tool_execution(): original_call_tool = ToolManager._call_tool @wraps(original_call_tool) - async def wrapped_call_tool( - self, call, *, allow_partial, wrap_validation_errors, approved - ): - # type: (Any, Any, bool, bool, bool) -> Any + async def wrapped_call_tool(self, call, *args, **kwargs): + # type: (Any, Any, *Any, **Any) -> Any # Extract tool info before calling original name = call.tool_name @@ -76,9 +74,8 @@ async def wrapped_call_tool( result = await original_call_tool( self, call, - allow_partial=allow_partial, - wrap_validation_errors=wrap_validation_errors, - approved=approved, + *args, + **kwargs, ) update_execute_tool_span(span, result) return result @@ -90,9 +87,8 @@ async def wrapped_call_tool( return await original_call_tool( self, call, - allow_partial=allow_partial, - wrap_validation_errors=wrap_validation_errors, - approved=approved, + *args, + **kwargs, ) ToolManager._call_tool = wrapped_call_tool