fix: significantly simplifying instrumentation

Removing the workflow spans, using better suited hook functions to patch into

Author: constantinius

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -35,8 +35,7 @@ def setup_once():
         Set up the pydantic-ai integration.
 
         This patches the key methods in pydantic-ai to create Sentry spans for:
-        - Agent workflow execution (root span)
-        - Individual agent invocations
+        - Agent invocations (Agent.run methods)
         - Model requests (AI client calls)
         - Tool executions
         """

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -4,7 +4,7 @@
 from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception
 
-from ..spans import agent_workflow_span
+from ..spans import invoke_agent_span, update_invoke_agent_span
 
 from typing import TYPE_CHECKING
 from pydantic_ai.agent import Agent
@@ -28,17 +28,29 @@ def _capture_exception(exc):
 class _StreamingContextManagerWrapper:
     """Wrapper for streaming methods that return async context managers."""
 
-    def __init__(self, agent, original_ctx_manager, is_streaming=True):
-        # type: (Any, Any, bool) -> None
+    def __init__(
+        self,
+        agent,
+        original_ctx_manager,
+        user_prompt,
+        model,
+        model_settings,
+        is_streaming=True,
+    ):
+        # type: (Any, Any, Any, Any, Any, bool) -> None
         self.agent = agent
         self.original_ctx_manager = original_ctx_manager
+        self.user_prompt = user_prompt
+        self.model = model
+        self.model_settings = model_settings
         self.is_streaming = is_streaming
         self._isolation_scope = None  # type: Any
-        self._workflow_span = None  # type: Optional[sentry_sdk.tracing.Span]
+        self._span = None  # type: Optional[sentry_sdk.tracing.Span]
+        self._result = None  # type: Any
 
     async def __aenter__(self):
         # type: () -> Any
-        # Set up isolation scope and workflow span
+        # Set up isolation scope and invoke_agent span
         self._isolation_scope = sentry_sdk.isolation_scope()
         self._isolation_scope.__enter__()
 
@@ -47,23 +59,33 @@ async def __aenter__(self):
             "pydantic_ai_agent", {"_agent": self.agent, "_streaming": self.is_streaming}
         )
 
-        # Create workflow span
-        self._workflow_span = agent_workflow_span(self.agent)
-        self._workflow_span.__enter__()
+        # 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__()
 
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
+        self._result = result
         return result
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         # type: (Any, Any, Any) -> None
         try:
             # Exit the original context manager first
             await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb)
+
+            # Update span with output if successful
+            if exc_type is None and self._result and hasattr(self._result, "output"):
+                output = (
+                    self._result.output if hasattr(self._result, "output") else None
+                )
+                update_invoke_agent_span(self._span, output)
         finally:
-            # Clean up workflow span
-            if self._workflow_span:
-                self._workflow_span.__exit__(exc_type, exc_val, exc_tb)
+            # Clean up invoke span
+            if self._span:
+                self._span.__exit__(exc_type, exc_val, exc_tb)
 
             # Clean up isolation scope
             if self._isolation_scope:
@@ -73,7 +95,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 def _create_run_wrapper(original_func, is_streaming=False):
     # type: (Callable[..., Any], bool) -> Callable[..., Any]
     """
-    Wraps the Agent.run method to create a root span for the agent workflow.
+    Wraps the Agent.run method to create an invoke_agent span.
 
     Args:
         original_func: The original run method
@@ -92,19 +114,23 @@ async def wrapper(self, *args, **kwargs):
                 "pydantic_ai_agent", {"_agent": self, "_streaming": is_streaming}
             )
 
-            with agent_workflow_span(self):
-                result = None
+            # Extract parameters for the span
+            user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+            model = kwargs.get("model")
+            model_settings = kwargs.get("model_settings")
+
+            # Create invoke_agent span
+            with invoke_agent_span(user_prompt, self, model, model_settings) as span:
                 try:
                     result = await original_func(self, *args, **kwargs)
+
+                    # Update span with output
+                    output = result.output if hasattr(result, "output") else None
+                    update_invoke_agent_span(span, output)
+
                     return result
                 except Exception as exc:
                     _capture_exception(exc)
-
-                    # It could be that there is an "invoke agent" span still open
-                    current_span = sentry_sdk.get_current_span()
-                    if current_span is not None and current_span.timestamp is None:
-                        current_span.__exit__(None, None, None)
-
                     raise exc from None
 
     return wrapper
@@ -113,36 +139,22 @@ async def wrapper(self, *args, **kwargs):
 def _create_run_sync_wrapper(original_func):
     # type: (Callable[..., Any]) -> Callable[..., Any]
     """
-    Wraps the Agent.run_sync method to create a root span for the agent workflow.
-    Note: run_sync is always non-streaming.
+    Wraps the Agent.run_sync method - no span needed as it delegates to run().
+
+    Note: run_sync just calls self.run() via run_until_complete, so the
+    invoke_agent span will be created by the run() wrapper.
     """
 
     @wraps(original_func)
     def wrapper(self, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
-        # Isolate each workflow so that when agents are run 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": False}
-            )
-
-            with agent_workflow_span(self):
-                result = None
-                try:
-                    result = original_func(self, *args, **kwargs)
-                    return result
-                except Exception as exc:
-                    _capture_exception(exc)
-
-                    # It could be that there is an "invoke agent" span still open
-                    current_span = sentry_sdk.get_current_span()
-                    if current_span is not None and current_span.timestamp is None:
-                        current_span.__exit__(None, None, None)
-
-                    raise exc from None
+        # Just call the original function - it will call run() which has the instrumentation
+        try:
+            result = original_func(self, *args, **kwargs)
+            return result
+        except Exception as exc:
+            _capture_exception(exc)
+            raise exc from None
 
     return wrapper
 
@@ -156,12 +168,22 @@ def _create_streaming_wrapper(original_func):
     @wraps(original_func)
     def wrapper(self, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
+        # Extract parameters for the span
+        user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+        model = kwargs.get("model")
+        model_settings = kwargs.get("model_settings")
+
         # Call original function to get the context manager
         original_ctx_manager = original_func(self, *args, **kwargs)
 
         # Wrap it with our instrumentation
         return _StreamingContextManagerWrapper(
-            agent=self, original_ctx_manager=original_ctx_manager, is_streaming=True
+            agent=self,
+            original_ctx_manager=original_ctx_manager,
+            user_prompt=user_prompt,
+            model=model,
+            model_settings=model_settings,
+            is_streaming=True,
         )
 
     return wrapper
@@ -170,34 +192,22 @@ def wrapper(self, *args, **kwargs):
 def _create_streaming_events_wrapper(original_func):
     # type: (Callable[..., Any]) -> Callable[..., Any]
     """
-    Wraps run_stream_events method that returns an async generator/iterator.
+    Wraps run_stream_events method - no span needed as it delegates to run().
+
+    Note: run_stream_events internally calls self.run() with an event_stream_handler,
+    so the invoke_agent span will be created by the run() wrapper.
     """
 
     @wraps(original_func)
     async def wrapper(self, *args, **kwargs):
         # type: (Any, *Any, **Any) -> Any
-        # 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
-            sentry_sdk.get_current_scope().set_context(
-                "pydantic_ai_agent", {"_agent": self, "_streaming": True}
-            )
-
-            with agent_workflow_span(self):
-                try:
-                    # Call the original generator and yield all events
-                    async for event in original_func(self, *args, **kwargs):
-                        yield event
-                except Exception as exc:
-                    _capture_exception(exc)
-
-                    # It could be that there is an "invoke agent" span still open
-                    current_span = sentry_sdk.get_current_span()
-                    if current_span is not None and current_span.timestamp is None:
-                        current_span.__exit__(None, None, None)
-
-                    raise exc from None
+        # Just call the original generator - it will call run() which has the instrumentation
+        try:
+            async for event in original_func(self, *args, **kwargs):
+                yield event
+        except Exception as exc:
+            _capture_exception(exc)
+            raise exc from None
 
     return wrapper