ref(pydantic-ai): Consolidate run-scoped state into one context manager

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -31,6 +31,13 @@
 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
@@ -41,12 +48,17 @@ async def on_request(
         if not isinstance(run_context_metadata, dict):
             return request_context
 
-        span = ai_client_span(
-            messages=request_context.messages,
-            agent=None,
-            model=request_context.model,
-            model_settings=request_context.model_settings,
-        )
+        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__()
@@ -68,7 +80,8 @@ async def on_response(
         if span is None:
             return response
 
-        update_ai_client_span(span, response)
+        with capture_internal_exceptions():
+            update_ai_client_span(span, response)
         span.__exit__(None, None, None)
 
         return response
@@ -116,15 +129,15 @@ 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.
-    - Each run invocation has `RunContext` objects that are passed to the library hooks.
     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).
-    The span is created in `on_request` and stored in the metadata of the `RunContext` object shared with `on_response` and `on_error`.
+    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 metadata dictionary on the RunContext instance is initialized with `{"_sentry_span": None}` in the `_create_run_wrapper()` and `_create_streaming_wrapper()` wrappers that
-    instrument `Agent.run()` and `Agent.run_stream()`, respectively. A non-empty dictionary is required for the metadata object to be a shared reference between hooks.
+    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()`.
     """
 
     identifier = "pydantic_ai"

sentry_sdk/integrations/pydantic_ai/_extract.py

@@ -15,6 +15,7 @@
 from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
 from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX
 from sentry_sdk.ai.utils import get_modality_from_mime_type
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.utils import safe_serialize
 
 try:
@@ -53,24 +54,15 @@ class ModelInfo:
     settings: "Dict[str, Any]" = field(default_factory=dict)
 
 
-@dataclass
-class UsageInfo:
-    input_tokens: "Optional[int]" = None
-    cache_read_tokens: "Optional[int]" = None
-    cache_write_tokens: "Optional[int]" = None
-    output_tokens: "Optional[int]" = None
-    total_tokens: "Optional[int]" = None
-
-
-# Model settings that get mirrored onto spans; values are read with dict
-# access first because ModelSettings is a TypedDict (dict at runtime).
-MODEL_SETTING_NAMES = (
-    "max_tokens",
-    "temperature",
-    "top_p",
-    "frequency_penalty",
-    "presence_penalty",
-)
+# Single source of truth for which model settings get mirrored onto spans
+# and which span attribute each one maps to.
+MODEL_SETTINGS_TO_SPANDATA = {
+    "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
+    "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
+    "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
+    "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
+    "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
+}
 
 
 def get_model_name(model_obj: "Any") -> "Optional[str]":
@@ -97,7 +89,7 @@ def extract_model_settings(settings: "Any") -> "Dict[str, Any]":
     if not settings:
         return extracted
 
-    for setting_name in MODEL_SETTING_NAMES:
+    for setting_name in MODEL_SETTINGS_TO_SPANDATA:
         if isinstance(settings, dict):
             value = settings.get(setting_name)
         else:
@@ -167,8 +159,8 @@ def extract_available_tools(agent: "Any") -> "Optional[List[Dict[str, Any]]]":
         return None
 
 
-def extract_usage(usage: "Any") -> "Optional[UsageInfo]":
-    """Extract token usage counts.
+def extract_usage_kwargs(usage: "Any") -> "Optional[Dict[str, Optional[int]]]":
+    """Extract token usage counts as record_token_usage keyword arguments.
 
     Works with both RequestUsage (single request) and RunUsage (agent run)
     objects from pydantic-ai; note the library uses cache_read_tokens /
@@ -177,13 +169,13 @@ def extract_usage(usage: "Any") -> "Optional[UsageInfo]":
     if usage is None:
         return None
 
-    return UsageInfo(
-        input_tokens=getattr(usage, "input_tokens", None),
-        cache_read_tokens=getattr(usage, "cache_read_tokens", None),
-        cache_write_tokens=getattr(usage, "cache_write_tokens", None),
-        output_tokens=getattr(usage, "output_tokens", None),
-        total_tokens=getattr(usage, "total_tokens", None),
-    )
+    return {
+        "input_tokens": getattr(usage, "input_tokens", None),
+        "input_tokens_cached": getattr(usage, "cache_read_tokens", None),
+        "input_tokens_cache_write": getattr(usage, "cache_write_tokens", None),
+        "output_tokens": getattr(usage, "output_tokens", None),
+        "total_tokens": getattr(usage, "total_tokens", None),
+    }
 
 
 def serialize_image_url_item(item: "Any") -> "Dict[str, Any]":

sentry_sdk/integrations/pydantic_ai/_run_context.py

@@ -0,0 +1,63 @@
+"""Run-scoped state shared between the agent wrappers and the model/tool
+instrumentation.
+
+One agent run corresponds to one AgentRun on the contextvar stack. The stack
+makes nested agent calls re-entrant safe, and the context manager guarantees
+push/pop pairing.
+"""
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Iterator, Optional
+
+
+@dataclass
+class AgentRun:
+    """State for one in-flight agent run."""
+
+    agent: "Any"
+    is_streaming: bool = False
+
+
+_agent_run_stack: "ContextVar[tuple[AgentRun, ...]]" = ContextVar(
+    "pydantic_ai_agent_run_stack", default=()
+)
+
+
+def current_agent_run() -> "Optional[AgentRun]":
+    stack = _agent_run_stack.get()
+    return stack[-1] if stack else None
+
+
+def get_current_agent() -> "Any":
+    run = current_agent_run()
+    return run.agent if run is not None else None
+
+
+def get_is_streaming() -> bool:
+    run = current_agent_run()
+    return run.is_streaming if run is not None else False
+
+
+@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."""
+    run = AgentRun(agent=agent, is_streaming=is_streaming)
+    token = _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