Better hooking into agent invokation

Author: antonpirker

sentry_sdk/integrations/openai_agents/__init__.py

@@ -4,6 +4,7 @@
     _create_get_model_wrapper,
     _create_get_all_tools_wrapper,
     _create_run_wrapper,
+    _patch_agent_run,
 )
 
 try:
@@ -15,15 +16,16 @@
 
 def _patch_runner():
     # type: () -> None
-
-    # Creating agent workflow spans
+    # Create the root span for one full agent run (including eventual handoffs)
     # Note agents.run.DEFAULT_AGENT_RUNNER.run_sync is a wrapper around
     # agents.run.DEFAULT_AGENT_RUNNER.run. It does not need to be wrapped separately.
+    # TODO-anton: Also patch streaming runner: agents.Runner.run_streamed
     agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
         agents.run.DEFAULT_AGENT_RUNNER.run
     )
 
-    # TODO-anton: Also patch streaming runner: agents.Runner.run_streamed
+    # Creating the actual spans for each agent run.
+    _patch_agent_run()
 
 
 def _patch_model():

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -1,3 +1,4 @@
 from .models import _create_get_model_wrapper  # noqa: F401
 from .tools import _create_get_all_tools_wrapper  # noqa: F401
 from .runner import _create_run_wrapper  # noqa: F401
+from .agent_run import _patch_agent_run  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -0,0 +1,163 @@
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import invoke_agent_span, update_invoke_agent_span
+
+
+try:
+    import agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _patch_agent_run():
+    """
+    Patches AgentRunner methods to create agent invocation spans without using RunHooks.
+    This directly patches the execution flow to track when agents start and stop.
+    """
+
+    # Store original methods
+    original_run_single_turn = agents.run.AgentRunner._run_single_turn
+    original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
+    original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
+
+    def _start_invoke_agent_span(context_wrapper, agent):
+        """Start an agent invocation span"""
+        # Store the agent on the context wrapper so we can access it later
+        context_wrapper._sentry_current_agent = agent
+        invoke_agent_span(context_wrapper, agent)
+
+    def _end_invoke_agent_span(context_wrapper, agent, output):
+        """End the agent invocation span"""
+        # Clear the stored agent
+        if hasattr(context_wrapper, "_sentry_current_agent"):
+            delattr(context_wrapper, "_sentry_current_agent")
+
+        update_invoke_agent_span(context_wrapper, agent, output)
+
+    def _has_active_agent_span(context_wrapper):
+        """Check if there's an active agent span for this context"""
+        return hasattr(context_wrapper, "_sentry_current_agent")
+
+    def _get_current_agent(context_wrapper):
+        """Get the current agent from context wrapper"""
+        return getattr(context_wrapper, "_sentry_current_agent", None)
+
+    @wraps(original_run_single_turn)
+    async def patched_run_single_turn(
+        cls,
+        *,
+        agent,
+        all_tools,
+        original_input,
+        generated_items,
+        hooks,
+        context_wrapper,
+        run_config,
+        should_run_agent_start_hooks,
+        tool_use_tracker,
+        previous_response_id,
+    ):
+        """Patched _run_single_turn that creates agent invocation spans"""
+
+        # Start agent span when agent starts (but only once per agent)
+        if should_run_agent_start_hooks and agent and context_wrapper:
+            # End any existing span for a different agent
+            if _has_active_agent_span(context_wrapper):
+                current_agent = _get_current_agent(context_wrapper)
+                if current_agent and current_agent != agent:
+                    _end_invoke_agent_span(context_wrapper, current_agent, None)
+
+            _start_invoke_agent_span(context_wrapper, agent)
+
+        # Call original method with all the correct parameters
+        result = await original_run_single_turn(
+            agent=agent,
+            all_tools=all_tools,
+            original_input=original_input,
+            generated_items=generated_items,
+            hooks=hooks,
+            context_wrapper=context_wrapper,
+            run_config=run_config,
+            should_run_agent_start_hooks=should_run_agent_start_hooks,
+            tool_use_tracker=tool_use_tracker,
+            previous_response_id=previous_response_id,
+        )
+
+        return result
+
+    @wraps(original_execute_handoffs)
+    async def patched_execute_handoffs(
+        cls,
+        *,
+        agent,
+        original_input,
+        pre_step_items,
+        new_step_items,
+        new_response,
+        run_handoffs,
+        hooks,
+        context_wrapper,
+        run_config,
+    ):
+        """Patched execute_handoffs that ends agent span for handoffs"""
+
+        # Call original method with all parameters
+        result = await original_execute_handoffs(
+            agent=agent,
+            original_input=original_input,
+            pre_step_items=pre_step_items,
+            new_step_items=new_step_items,
+            new_response=new_response,
+            run_handoffs=run_handoffs,
+            hooks=hooks,
+            context_wrapper=context_wrapper,
+            run_config=run_config,
+        )
+
+        # End span for current agent after handoff processing is complete
+        if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+            _end_invoke_agent_span(context_wrapper, agent, None)
+
+        return result
+
+    @wraps(original_execute_final_output)
+    async def patched_execute_final_output(
+        cls,
+        *,
+        agent,
+        original_input,
+        new_response,
+        pre_step_items,
+        new_step_items,
+        final_output,
+        hooks,
+        context_wrapper,
+    ):
+        """Patched execute_final_output that ends agent span for final outputs"""
+
+        # Call original method with all parameters
+        result = await original_execute_final_output(
+            agent=agent,
+            original_input=original_input,
+            new_response=new_response,
+            pre_step_items=pre_step_items,
+            new_step_items=new_step_items,
+            final_output=final_output,
+            hooks=hooks,
+            context_wrapper=context_wrapper,
+        )
+
+        # End span for current agent after final output processing is complete
+        if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+            _end_invoke_agent_span(context_wrapper, agent, final_output)
+
+        return result
+
+    # Apply patches
+    agents.run.AgentRunner._run_single_turn = classmethod(patched_run_single_turn)
+    agents._run_impl.RunImpl.execute_handoffs = classmethod(patched_execute_handoffs)
+    agents._run_impl.RunImpl.execute_final_output = classmethod(
+        patched_execute_final_output
+    )

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -8,7 +8,7 @@
     agent_workflow_span,
     update_agent_workflow_span,
 )
-from ..utils import _capture_exception, _wrap_hooks
+from ..utils import _capture_exception
 
 from typing import TYPE_CHECKING
 
@@ -28,7 +28,6 @@ async def async_wrapper(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         agent = args[0]
         with agent_workflow_span(agent) as span:
-            kwargs["hooks"] = _wrap_hooks(kwargs.get("hooks"))
             result = None
             try:
                 result = await original_func(*args, **kwargs)

sentry_sdk/integrations/openai_agents/run_hooks.py

@@ -1,4 +1,3 @@
-from functools import wraps
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
@@ -156,49 +155,3 @@ def _set_output_data(span, result):
 
     if len(output_messages["response"]) > 0:
         span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"])
-
-
-def _create_hook_wrapper(original_hook, sentry_hook):
-    # type: (Callable[..., Any], Callable[..., Any]) -> Callable[..., Any]
-    @wraps(original_hook)
-    async def async_wrapper(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
-        await sentry_hook(*args, **kwargs)
-        return await original_hook(*args, **kwargs)
-
-    return async_wrapper
-
-
-def _wrap_hooks(hooks):
-    # type: (agents.RunHooks) -> agents.RunHooks
-    """
-    Our integration uses RunHooks to create spans. This function will either
-    enable our SentryRunHooks or if the users has given custom RunHooks wrap
-    them so the Sentry hooks and the users hooks are both called
-    """
-    from .run_hooks import SentryRunHooks
-
-    sentry_hooks = SentryRunHooks()
-
-    if hooks is None:
-        return sentry_hooks
-
-    wrapped_hooks = type("SentryWrappedHooks", (hooks.__class__,), {})
-
-    # Wrap all methods from RunHooks
-    for method_name in dir(agents.RunHooks):
-        if method_name.startswith("on_"):
-            original_method = getattr(hooks, method_name)
-            # Only wrap if the method exists in SentryRunHooks
-            try:
-                sentry_method = getattr(sentry_hooks, method_name)
-                setattr(
-                    wrapped_hooks,
-                    method_name,
-                    _create_hook_wrapper(original_method, sentry_method),
-                )
-            except AttributeError:
-                # If method doesn't exist in SentryRunHooks, just use the original method
-                setattr(wrapped_hooks, method_name, original_method)
-
-    return wrapped_hooks()