fix(openai-agents): Avoid double span exit on exception

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -1,13 +1,17 @@
 from functools import wraps
 
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.tracing_utils import set_span_errored
 from ..spans import invoke_agent_span, update_invoke_agent_span, handoff_span
+from ..utils import _capture_exception, _record_exception_on_span
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Optional
 
+    from sentry_sdk.tracing import Span
+
 try:
     import agents
 except ImportError:
@@ -27,13 +31,15 @@ def _patch_agent_run():
     original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
 
     def _start_invoke_agent_span(context_wrapper, agent, kwargs):
-        # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> None
+        # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> Span
         """Start an agent invocation span"""
         # Store the agent on the context wrapper so we can access it later
         context_wrapper._sentry_current_agent = agent
         span = invoke_agent_span(context_wrapper, agent, kwargs)
         context_wrapper._sentry_agent_span = span
 
+        return span
+
     def _end_invoke_agent_span(context_wrapper, agent, output=None):
         # type: (agents.RunContextWrapper, agents.Agent, Optional[Any]) -> None
         """End the agent invocation span"""
@@ -65,18 +71,21 @@ async def patched_run_single_turn(cls, *args, **kwargs):
         context_wrapper = kwargs.get("context_wrapper")
         should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks")
 
+        span = getattr(context_wrapper, "_sentry_current_agent", None)
         # 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)
-
-            _start_invoke_agent_span(context_wrapper, agent, kwargs)
+            span = _start_invoke_agent_span(context_wrapper, agent, kwargs)
 
         # Call original method with all the correct parameters
-        result = await original_run_single_turn(*args, **kwargs)
+        try:
+            result = await original_run_single_turn(*args, **kwargs)
+        except Exception as exc:
+            if span is not None and span.timestamp is None:
+                _record_exception_on_span(span, exc)
+                span.__exit__(None, None, None)
+
+            _capture_exception(exc)
+            raise exc from None
 
         return result
 

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -28,18 +28,18 @@ async def wrapper(*args, **kwargs):
         with sentry_sdk.isolation_scope():
             agent = args[0]
             with agent_workflow_span(agent):
-                result = None
-                try:
-                    result = await original_func(*args, **kwargs)
-                    return result
-                except Exception as exc:
-                    _capture_exception(exc)
-
-                    # It could be that there is a "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
+                run_result = await original_func(*args, **kwargs)
+
+                invoke_agent_span = getattr(
+                    run_result.context_wrapper, "_sentry_agent_span", None
+                )
+
+                if (
+                    invoke_agent_span is not None
+                    and invoke_agent_span.timestamp is None
+                ):
+                    invoke_agent_span.__exit__(None, None, None)
+
+                return run_result
 
     return wrapper

sentry_sdk/integrations/openai_agents/utils.py

@@ -36,6 +36,21 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
+def _record_exception_on_span(span, error):
+    # type: (Optional[Span], Exception) -> Any
+    if span is not None:
+        set_span_errored(span)
+        span.set_data("span.status", "error")
+
+        # Optionally capture the error details if we have them
+        if hasattr(error, "__class__"):
+            span.set_data("error.type", error.__class__.__name__)
+        if hasattr(error, "__str__"):
+            error_message = str(error)
+            if error_message:
+                span.set_data("error.message", error_message)
+
+
 def _set_agent_data(span, agent):
     # type: (sentry_sdk.tracing.Span, agents.Agent) -> None
     span.set_data(