Bring back handoff span

Author: antonpirker

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -10,6 +10,7 @@
     invoke_agent_span,
     update_invoke_agent_span,
 )
+from ..utils import _wrap_hooks
 
 from typing import TYPE_CHECKING
 
@@ -36,6 +37,7 @@ async def async_wrapper(cls, *args, **kwargs):
         # type: (agents.Runner, *Any, **Any) -> Any
         agent = args[0]
         with agent_workflow_span(agent) as span:
+            kwargs["hooks"] = _wrap_hooks(kwargs.get("hooks"))
             result = await original_func(*args, **kwargs)
             update_agent_workflow_span(span, agent, result)
 
@@ -46,6 +48,7 @@ async def async_wrapper(cls, *args, **kwargs):
     )
     def sync_wrapper(cls, *args, **kwargs):
         # type: (agents.Runner, *Any, **Any) -> Any
+        kwargs["hooks"] = _wrap_hooks(kwargs.get("hooks"))
         # The sync version (.run_sync()) is just a wrapper around the async version (.run())
         return original_func(*args, **kwargs)
 

sentry_sdk/integrations/openai_agents/spans/handoff.py

@@ -1,6 +1,8 @@
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 
+from ..consts import SPAN_ORIGIN
+
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -9,12 +11,9 @@
 
 def handoff_span(context, from_agent, to_agent):
     # type: (agents.RunContextWrapper, agents.Agent, agents.Agent) -> None
-    current_span = sentry_sdk.get_current_span()
-    if current_span:
-        with current_span.start_child(
-            op=OP.GEN_AI_HANDOFF,
-            name=f"handoff from {from_agent.name} to {to_agent.name}",
-        ) as span:
-            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")
-
-        current_span.__exit__(None, None, None)
+    with sentry_sdk.start_span(
+        op=OP.GEN_AI_HANDOFF,
+        name=f"handoff from {from_agent.name} to {to_agent.name}",
+        origin=SPAN_ORIGIN,
+    ) as span:
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")

sentry_sdk/integrations/openai_agents/utils.py

@@ -1,3 +1,4 @@
+from functools import wraps
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
@@ -155,3 +156,52 @@ 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
+    """
+    import ipdb
+
+    ipdb.set_trace()
+    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()