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

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -1,15 +1,25 @@
 from functools import wraps
 
 from sentry_sdk.integrations import DidNotEnable
-from ..spans import invoke_agent_span, update_invoke_agent_span, handoff_span
+from sentry_sdk.tracing_utils import set_span_errored
+from ..spans import (
+    invoke_agent_span,
+    update_invoke_agent_span,
+    end_invoke_agent_span,
+    handoff_span,
+)
+from ..utils import _capture_exception, _record_exception_on_span, _SingleTurnException
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Optional
 
+    from sentry_sdk.tracing import Span
+
 try:
     import agents
+    from agents.exceptions import AgentsException
 except ImportError:
     raise DidNotEnable("OpenAI Agents not installed")
 
@@ -27,13 +37,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 +77,33 @@ 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
             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)
+                    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 AgentsException:
+            # AgentsException is caught on AgentRunner.run().
+            # Exceptions are captured and agent invocation spans are explicitly finished
+            # as long as only AgentRunner.run() invokes AgentRunner._run_single_turn().
+            raise
+        except Exception as exc:
+            _capture_exception(exc)
+
+            if span is not None and span.timestamp is None:
+                _record_exception_on_span(span, exc)
+                end_invoke_agent_span(context_wrapper, agent)
+
+            raise _SingleTurnException(exc)
 
         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

@@ -1,9 +1,15 @@
 from functools import wraps
 
 import sentry_sdk
+from sentry_sdk.integrations import DidNotEnable
 
-from ..spans import agent_workflow_span
-from ..utils import _capture_exception
+from ..spans import agent_workflow_span, end_invoke_agent_span
+from ..utils import _capture_exception, _record_exception_on_span, _SingleTurnException
+
+try:
+    from agents.exceptions import AgentsException
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
 
 from typing import TYPE_CHECKING
 
@@ -28,18 +34,36 @@ 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:
+                    run_result = await original_func(*args, **kwargs)
+                except AgentsException 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)
+                    context_wrapper = getattr(exc.run_data, "context_wrapper", None)
+                    if context_wrapper is not None:
+                        invoke_agent_span = getattr(
+                            context_wrapper, "_sentry_agent_span", None
+                        )
+
+                        if (
+                            invoke_agent_span is not None
+                            and invoke_agent_span.timestamp is None
+                        ):
+                            _record_exception_on_span(invoke_agent_span, exc)
+                            end_invoke_agent_span(context_wrapper, agent)
 
                     raise exc from None
+                except _SingleTurnException as exc:
+                    # Handled in _run_single_turn() patch.
+                    raise exc.original from None
+                except Exception as exc:
+                    # Invoke agent span is not finished in this case.
+                    # This is much less likely to occur than other cases because
+                    # AgentRunner.run() is "just" a while loop around _run_single_turn.
+                    _capture_exception(exc)
+                    raise exc from None
+
+                end_invoke_agent_span(run_result.context_wrapper, agent)
+                return run_result
 
     return wrapper

sentry_sdk/integrations/openai_agents/spans/__init__.py

@@ -2,4 +2,8 @@
 from .ai_client import ai_client_span, update_ai_client_span  # noqa: F401
 from .execute_tool import execute_tool_span, update_execute_tool_span  # noqa: F401
 from .handoff import handoff_span  # noqa: F401
-from .invoke_agent import invoke_agent_span, update_invoke_agent_span  # noqa: F401
+from .invoke_agent import (
+    invoke_agent_span,
+    update_invoke_agent_span,
+    end_invoke_agent_span,
+)  # noqa: F401

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -16,7 +16,7 @@
 
 if TYPE_CHECKING:
     import agents
-    from typing import Any
+    from typing import Any, Optional
 
 
 def invoke_agent_span(context, agent, kwargs):
@@ -91,3 +91,13 @@ def update_invoke_agent_span(context, agent, output):
 
         span.__exit__(None, None, None)
         delattr(context, "_sentry_agent_span")
+
+
+def end_invoke_agent_span(context_wrapper, agent, output=None):
+    # type: (agents.RunContextWrapper, agents.Agent, Optional[Any]) -> None
+    """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)

sentry_sdk/integrations/openai_agents/utils.py

@@ -18,13 +18,21 @@
     from typing import Any
     from agents import Usage
 
+    from sentry_sdk.tracing import Span
+
 try:
     import agents
 
 except ImportError:
     raise DidNotEnable("OpenAI Agents not installed")
 
 
+class _SingleTurnException(Exception):
+    def __init__(self, original):
+        # type: (Exception) -> None
+        self.original = original
+
+
 def _capture_exception(exc):
     # type: (Any) -> None
     set_span_errored()
@@ -37,6 +45,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(