REDROOM

Edit File Size: 6.49 KB

/opt/hc_python/lib/python3.12/site-packages/sentry_sdk/integrations/pydantic_ai/patches/tools.py

Text Base64

import sys
from functools import wraps
from typing import TYPE_CHECKING

import sentry_sdk
from sentry_sdk.integrations import DidNotEnable
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

if TYPE_CHECKING:
    from typing import Any

try:
    try:
        from pydantic_ai.tool_manager import ToolManager  # type: ignore
    except ImportError:
        from pydantic_ai._tool_manager import ToolManager  # type: ignore

from pydantic_ai.exceptions import ToolRetryError  # type: ignore
except ImportError:
    raise DidNotEnable("pydantic-ai not installed")

def _patch_tool_execution() -> None:
    if hasattr(ToolManager, "execute_tool_call"):
        _patch_execute_tool_call()

elif hasattr(ToolManager, "_call_tool"):
        # older versions
        _patch_call_tool()

def _patch_execute_tool_call() -> None:
    original_execute_tool_call = ToolManager.execute_tool_call

@wraps(original_execute_tool_call)
    async def wrapped_execute_tool_call(
        self: "Any", validated: "Any", *args: "Any", **kwargs: "Any"
    ) -> "Any":
        if not validated or not hasattr(validated, "call"):
            return await original_execute_tool_call(self, validated, *args, **kwargs)

# Extract tool info before calling original
        call = validated.call
        name = call.tool_name
        tool = self.tools.get(name) if self.tools else None
        selected_tool_definition = getattr(tool, "tool_def", None)

# 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 {}

# 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_definition=selected_tool_definition,
                ) as span:
                    try:
                        result = await original_execute_tool_call(
                            self,
                            validated,
                            *args,
                            **kwargs,
                        )
                        update_execute_tool_span(span, result)
                        return result
                    except ToolRetryError as exc:
                        exc_info = sys.exc_info()
                        with capture_internal_exceptions():
                            # Avoid circular import due to multi-file integration structure
                            from sentry_sdk.integrations.pydantic_ai import (
                                PydanticAIIntegration,
                            )

integration = sentry_sdk.get_client().get_integration(
                                PydanticAIIntegration
                            )
                            if (
                                integration is not None
                                and integration.handled_tool_call_exceptions
                            ):
                                _capture_exception(exc, handled=True)
                        reraise(*exc_info)

return await original_execute_tool_call(self, validated, *args, **kwargs)

ToolManager.execute_tool_call = wrapped_execute_tool_call

def _patch_call_tool() -> None:
    """
    Patch ToolManager._call_tool to create execute_tool spans.

This is the single point where ALL tool calls flow through in pydantic_ai,
    regardless of toolset type (function, MCP, combined, wrapper, etc.).

By patching here, we avoid:
    - Patching multiple toolset classes
    - Dealing with signature mismatches from instrumented MCP servers
    - Complex nested toolset handling
    """
    original_call_tool = ToolManager._call_tool

@wraps(original_call_tool)
    async def wrapped_call_tool(
        self: "Any", call: "Any", *args: "Any", **kwargs: "Any"
    ) -> "Any":
        # Extract tool info before calling original
        name = call.tool_name
        tool = self.tools.get(name) if self.tools else None
        selected_tool_definition = getattr(tool, "tool_def", None)

# 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 {}

# 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_definition=selected_tool_definition,
                ) as span:
                    try:
                        result = await original_call_tool(
                            self,
                            call,
                            *args,
                            **kwargs,
                        )
                        update_execute_tool_span(span, result)
                        return result
                    except ToolRetryError as exc:
                        exc_info = sys.exc_info()
                        with capture_internal_exceptions():
                            # Avoid circular import due to multi-file integration structure
                            from sentry_sdk.integrations.pydantic_ai import (
                                PydanticAIIntegration,
                            )

# No span context - just call original
        return await original_call_tool(
            self,
            call,
            *args,
            **kwargs,
        )

ToolManager._call_tool = wrapped_call_tool

Exit & Reset Text mode: syntax highlighting auto-detects file type.

Name	Size	Perms	Modified	Actions
__pycache__ DIR	-	drwxr-xr-x	2026-06-11 06:30:31	Edit Download recursive recursive
agent_run.py	6.85 KB	lrw-r--r--	2026-06-11 06:30:31	Edit Download recursive recursive
graph_nodes.py	3.69 KB	lrw-r--r--	2026-06-11 06:30:31	Edit Download recursive recursive
tools.py	6.49 KB	lrw-r--r--	2026-06-11 06:30:31	Edit Download recursive recursive
__init__.py	167 B	lrw-r--r--	2026-06-11 06:30:31	Edit Download recursive recursive

Directory Contents