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

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,11 +31,14 @@ 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
-        invoke_agent_span(context_wrapper, agent, kwargs)
+        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
@@ -64,6 +71,7 @@ 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_agent_span", 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
@@ -72,10 +80,19 @@ async def patched_run_single_turn(cls, *args, **kwargs):
                 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)
+                delattr(context_wrapper, "_sentry_agent_span")
+
+            _capture_exception(exc)
+            raise exc from None
 
         return result
 

sentry_sdk/integrations/openai_agents/patches/error_tracing.py

@@ -3,6 +3,7 @@
 import sentry_sdk
 from sentry_sdk.consts import SPANSTATUS
 from sentry_sdk.tracing_utils import set_span_errored
+from ..utils import _record_exception_on_span
 
 from typing import TYPE_CHECKING
 
@@ -58,16 +59,7 @@ def sentry_attach_error_to_current_span(error, *args, **kwargs):
         # Set the current Sentry span to errored
         current_span = sentry_sdk.get_current_span()
         if current_span is not None:
-            set_span_errored(current_span)
-            current_span.set_data("span.status", "error")
-
-            # Optionally capture the error details if we have them
-            if hasattr(error, "__class__"):
-                current_span.set_data("error.type", error.__class__.__name__)
-            if hasattr(error, "__str__"):
-                error_message = str(error)
-                if error_message:
-                    current_span.set_data("error.message", error_message)
+            _record_exception_on_span(current_span, error)
 
         # Call the original function
         return original_attach_error(error, *args, **kwargs)

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/spans/invoke_agent.py

@@ -75,7 +75,7 @@ def invoke_agent_span(context, agent, kwargs):
 
 def update_invoke_agent_span(context, agent, output):
     # type: (agents.RunContextWrapper, agents.Agent, Any) -> None
-    span = sentry_sdk.get_current_span()
+    span = getattr(context, "_sentry_agent_span", None)
 
     if span:
         if should_send_default_pii():
@@ -84,3 +84,4 @@ def update_invoke_agent_span(context, agent, output):
             )
 
         span.__exit__(None, None, None)
+        delattr(context, "_sentry_agent_span")

sentry_sdk/integrations/openai_agents/utils.py

@@ -17,6 +17,8 @@
     from typing import Any
     from agents import Usage
 
+    from sentry_sdk.tracing import Span
+
 try:
     import agents
 
@@ -36,6 +38,20 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
+def _record_exception_on_span(span, error):
+    # type: (Span, Exception) -> Any
+    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(