ref(pydantic-ai): Resolve version divergence in _compat and restructure modules

Introduce _compat.py, which resolves the installed pydantic-ai version and every version-dependent decision once at import time: the hooks-vs-graph-nodes model backend, the ToolManager tool-call method name, and the message part classes (now resolved individually so one upstream rename degrades only the paths needing that class). The using_request_hooks class flag and the circular-import workarounds are gone; setup_once is three composed calls with deferred imports, and importing the package loads nothing but __init__.

Restructure to the target layout: _spans.py absorbs utils.py and the spans/ package, _wrap_agent.py and _wrap_model.py (both model backends behind one install_model_backend()) and _wrap_tools.py (the two duplicated tool wrappers unified into one) replace patches/.

Review-driven fixes folded in: unknown ToolManager method names and a missing _agent_graph module now degrade gracefully instead of crashing sentry_sdk.init(), agent_run_scope removes exactly its own run from the stack so non-LIFO streaming exits cannot corrupt it, the streaming wrapper propagates exception suppression from the wrapped context manager, the after_model_request span close is guarded, and chat spans extract model info once instead of twice.

Author: ericapisani

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -1,148 +1,30 @@
-import functools
-
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.utils import capture_internal_exceptions, parse_version
 
 try:
     import pydantic_ai  # noqa: F401
-    from pydantic_ai import Agent
 except ImportError:
     raise DidNotEnable("pydantic-ai not installed")
 
 
-from importlib.metadata import PackageNotFoundError, version
-from typing import TYPE_CHECKING
-
-from .patches import (
-    _patch_agent_run,
-    _patch_graph_nodes,
-    _patch_tool_execution,
-)
-from .spans.ai_client import ai_client_span, update_ai_client_span
-
-if TYPE_CHECKING:
-    from typing import Any
-
-    from pydantic_ai import ModelRequestContext, RunContext
-    from pydantic_ai.capabilities import Hooks
-    from pydantic_ai.messages import ModelResponse
-
-
-def register_hooks(hooks: "Hooks") -> None:
-    """
-    Creates hooks for chat model calls and register the hooks by adding the hooks to the `capabilities` argument passed to `Agent.__init__()`.
-
-    The chat span opened in on_request is stored in the run's `RunContext.metadata`
-    dict, which pydantic-ai shares by reference between the hooks of one run. This
-    keeps span pairing correct per run (even for overlapping runs in one task) and
-    covers every entry point that fires request hooks (including `Agent.iter()`,
-    which the Agent.run/run_stream wrappers never see). It requires seeding a
-    metadata dict in `patched_init` below when the user did not provide one.
-    """
-
-    @hooks.on.before_model_request
-    async def on_request(
-        ctx: "RunContext[None]", request_context: "ModelRequestContext"
-    ) -> "ModelRequestContext":
-        run_context_metadata = ctx.metadata
-        if not isinstance(run_context_metadata, dict):
-            return request_context
-
-        span = None
-        with capture_internal_exceptions():
-            span = ai_client_span(
-                messages=request_context.messages,
-                agent=None,
-                model=request_context.model,
-                model_settings=request_context.model_settings,
-            )
-
-        if span is None:
-            return request_context
-
-        run_context_metadata["_sentry_span"] = span
-        span.__enter__()
-
-        return request_context
-
-    @hooks.on.after_model_request
-    async def on_response(
-        ctx: "RunContext[None]",
-        *,
-        request_context: "ModelRequestContext",
-        response: "ModelResponse",
-    ) -> "ModelResponse":
-        run_context_metadata = ctx.metadata
-        if not isinstance(run_context_metadata, dict):
-            return response
-
-        span = run_context_metadata.pop("_sentry_span", None)
-        if span is None:
-            return response
-
-        with capture_internal_exceptions():
-            update_ai_client_span(span, response)
-        span.__exit__(None, None, None)
-
-        return response
-
-    @hooks.on.model_request_error
-    async def on_error(
-        ctx: "RunContext[None]",
-        *,
-        request_context: "ModelRequestContext",
-        error: "Exception",
-    ) -> "ModelResponse":
-        run_context_metadata = ctx.metadata
-
-        if not isinstance(run_context_metadata, dict):
-            raise error
-
-        span = run_context_metadata.pop("_sentry_span", None)
-        if span is None:
-            raise error
-
-        with capture_internal_exceptions():
-            span.__exit__(type(error), error, error.__traceback__)
-
-        raise error
-
-    original_init = Agent.__init__
-
-    @functools.wraps(original_init)
-    def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None:
-        caps = list(kwargs.get("capabilities") or [])
-        caps.append(hooks)
-        kwargs["capabilities"] = caps
-
-        metadata = kwargs.get("metadata")
-        if metadata is None:
-            kwargs["metadata"] = {}  # Used as shared reference between hooks
-
-        return original_init(self, *args, **kwargs)
-
-    Agent.__init__ = patched_init  # type: ignore[method-assign]
-
-
 class PydanticAIIntegration(Integration):
     """
     Typical interaction with the library:
     1. The user creates an Agent instance with configuration, including system instructions sent to every model call.
     2. The user calls `Agent.run()` or `Agent.run_stream()` to start an agent run. The latter can be used to incrementally receive progress.
     3. In a loop, the agent repeatedly calls the model, maintaining a conversation history that includes previous messages and tool results, which is passed to each call.
 
-    Internally, Pydantic AI maintains an execution graph in which ModelRequestNode are responsible for model calls, including retries.
-    Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions);
-    older versions are instrumented by patching the graph nodes directly (see patches/graph_nodes.py).
-
-    The wrappers around `Agent.run()` and `Agent.run_stream()` track each in-flight run on a contextvar stack (see _run_context.py); the tool patches and span
-    helpers read the current agent from there. The request hooks pair each chat span with its model request through the run's `RunContext.metadata` dict
-    (see register_hooks), which stays correct per run and also covers entry points the wrappers don't instrument, such as `Agent.iter()`.
+    How the integration is put together:
+    - _compat.py resolves the installed pydantic-ai version and every version-dependent decision, once, at import time.
+    - _extract.py is the only module that reads pydantic-ai object internals; it returns plain data structures.
+    - _spans.py creates spans and writes extracted data onto them.
+    - _run_context.py tracks each in-flight run on a contextvar stack; the tool wrapper and span helpers read the current agent from there.
+    - _wrap_agent.py instruments Agent.run / Agent.run_stream (invoke_agent spans, isolation scopes, run tracking).
+    - _wrap_model.py emits chat spans for model requests via one of two backends chosen in _compat: request hooks (>= 1.73), paired per run through RunContext.metadata, or graph-node patching (older versions).
+    - _wrap_tools.py instruments the single ToolManager method all tool calls flow through (execute_tool spans).
     """
 
     identifier = "pydantic_ai"
     origin = f"auto.ai.{identifier}"
-    using_request_hooks = False
 
     def __init__(
         self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True
@@ -169,32 +51,14 @@ def setup_once() -> None:
         - Model requests (AI client calls)
         - Tool executions
         """
+        # Deferred imports keep `import sentry_sdk.integrations.pydantic_ai`
+        # cheap when the integration is never enabled; they are the only
+        # intra-package imports in this module, keeping the import graph
+        # acyclic.
+        from ._wrap_agent import _patch_agent_run
+        from ._wrap_model import install_model_backend
+        from ._wrap_tools import _patch_tool_execution
+
         _patch_agent_run()
         _patch_tool_execution()
-
-        PydanticAIIntegration.using_request_hooks = False
-        try:
-            PYDANTIC_AI_VERSION = version("pydantic-ai-slim")
-        except PackageNotFoundError:
-            return
-
-        PYDANTIC_AI_VERSION = parse_version(PYDANTIC_AI_VERSION)
-        if PYDANTIC_AI_VERSION is None:
-            return
-
-        # ModelRequestContext.model added in https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c
-        if PYDANTIC_AI_VERSION < (
-            1,
-            73,
-        ):
-            _patch_graph_nodes()
-            return
-
-        try:
-            from pydantic_ai.capabilities import Hooks
-        except ImportError:
-            return
-
-        PydanticAIIntegration.using_request_hooks = True
-        hooks = Hooks()
-        register_hooks(hooks)
+        install_model_backend()

sentry_sdk/integrations/pydantic_ai/_compat.py

@@ -0,0 +1,70 @@
+"""Version detection and version-dependent imports for pydantic-ai.
+
+Everything here is resolved once at import time. The rest of the integration
+consumes the resulting constants instead of probing versions or attributes at
+call time, so "what runs on version X" is answered entirely by this module.
+"""
+
+from typing import TYPE_CHECKING
+
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.utils import package_version
+
+try:
+    from pydantic_ai import messages as _messages
+    from pydantic_ai.agent import Agent  # noqa: F401
+    from pydantic_ai.exceptions import ToolRetryError  # noqa: F401
+
+    try:
+        from pydantic_ai.tool_manager import ToolManager
+    except ImportError:
+        # older versions
+        from pydantic_ai._tool_manager import ToolManager  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+if TYPE_CHECKING:
+    from typing import Optional
+
+# Message part classes are resolved individually so that a single upstream
+# rename degrades only the extraction paths that need that class, instead of
+# silently disabling all of them at once.
+BaseToolCallPart = getattr(_messages, "BaseToolCallPart", None)
+BaseToolReturnPart = getattr(_messages, "BaseToolReturnPart", None)
+BinaryContent = getattr(_messages, "BinaryContent", None)
+ImageUrl = getattr(_messages, "ImageUrl", None)
+SystemPromptPart = getattr(_messages, "SystemPromptPart", None)
+TextPart = getattr(_messages, "TextPart", None)
+ThinkingPart = getattr(_messages, "ThinkingPart", None)
+
+PYDANTIC_AI_VERSION = package_version("pydantic-ai-slim")
+
+# The ToolManager method through which all tool calls flow; renamed from
+# _call_tool to execute_tool_call in newer versions. None means the method
+# could not be found and tool instrumentation is skipped.
+TOOL_CALL_METHOD: "Optional[str]" = None
+if hasattr(ToolManager, "execute_tool_call"):
+    TOOL_CALL_METHOD = "execute_tool_call"
+elif hasattr(ToolManager, "_call_tool"):
+    TOOL_CALL_METHOD = "_call_tool"
+
+# Request hooks (pydantic_ai.capabilities) are usable from 1.73 on, when
+# ModelRequestContext.model was added:
+# https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c
+USES_REQUEST_HOOKS = False
+if PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION >= (1, 73):
+    try:
+        from pydantic_ai.capabilities import Hooks  # noqa: F401
+
+        USES_REQUEST_HOOKS = True
+    except ImportError:
+        USES_REQUEST_HOOKS = False
+
+# Which mechanism emits chat spans for model requests: request hooks on new
+# versions, graph-node patching on old ones. None (unknown version, or hooks
+# unavailable on a new version) means only agent and tool spans are emitted.
+MODEL_BACKEND: "Optional[str]" = None
+if USES_REQUEST_HOOKS:
+    MODEL_BACKEND = "hooks"
+elif PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION < (1, 73):
+    MODEL_BACKEND = "graph_nodes"

sentry_sdk/integrations/pydantic_ai/_extract.py

@@ -4,8 +4,8 @@
 private attributes and version-dependent shapes) so that upstream library
 changes are absorbed here rather than throughout the integration. The one
 exception is control-flow state read at the patch points themselves (e.g.
-ModelRequestNode._did_stream in patches/graph_nodes.py and Tool.tool_def in
-patches/tools.py); everything else consumes the plain data structures
+ModelRequestNode._did_stream in _wrap_model.py and Tool.tool_def in
+_wrap_tools.py); everything else consumes the plain data structures
 returned here.
 """
 
@@ -18,25 +18,15 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.utils import safe_serialize
 
-try:
-    from pydantic_ai.messages import (
-        BaseToolCallPart,
-        BaseToolReturnPart,
-        BinaryContent,
-        ImageUrl,
-        SystemPromptPart,
-        TextPart,
-        ThinkingPart,
-    )
-except ImportError:
-    # Fallback if these classes are not available
-    BaseToolCallPart = None  # type: ignore[misc,assignment]
-    BaseToolReturnPart = None  # type: ignore[misc,assignment]
-    BinaryContent = None  # type: ignore[misc,assignment]
-    ImageUrl = None  # type: ignore[misc,assignment]
-    SystemPromptPart = None  # type: ignore[misc,assignment]
-    TextPart = None  # type: ignore[misc,assignment]
-    ThinkingPart = None  # type: ignore[misc,assignment]
+from ._compat import (
+    BaseToolCallPart,
+    BaseToolReturnPart,
+    BinaryContent,
+    ImageUrl,
+    SystemPromptPart,
+    TextPart,
+    ThinkingPart,
+)
 
 if TYPE_CHECKING:
     from typing import Any, Dict, List, Optional

sentry_sdk/integrations/pydantic_ai/_run_context.py

@@ -46,18 +46,18 @@ def get_is_streaming() -> bool:
 @contextmanager
 def agent_run_scope(agent: "Any", is_streaming: bool = False) -> "Iterator[AgentRun]":
     """Track an agent run on the contextvar stack for the duration of the
-    with block."""
+    with block.
+
+    On exit, exactly this run is removed from the stack (by identity, not a
+    token reset), so streaming runs that exit out of LIFO order or in a
+    different asyncio task never erase other still-active runs.
+    """
     run = AgentRun(agent=agent, is_streaming=is_streaming)
-    token = _agent_run_stack.set(_agent_run_stack.get() + (run,))
+    _agent_run_stack.set(_agent_run_stack.get() + (run,))
     try:
         yield run
     finally:
-        try:
-            _agent_run_stack.reset(token)
-        except (LookupError, ValueError):
-            # A streaming run's context manager can be exited in a different
-            # asyncio task (and therefore a different Context) than it was
-            # entered in, in which case the token cannot be reset. The stack
-            # entry only lives in the entering task's context copy, so there
-            # is nothing to clean up.
-            pass
+        stack = _agent_run_stack.get()
+        new_stack = tuple(r for r in stack if r is not run)
+        if len(new_stack) != len(stack):
+            _agent_run_stack.set(new_stack)