fix(openai-agents): Stop setting transaction status when child span fails

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -3,13 +3,13 @@
 
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
-from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import capture_internal_exceptions, reraise
 
 from ..spans import (
-    end_invoke_agent_span,
+    _AgentInvocationSpanContext,
     handoff_span,
     invoke_agent_span,
+    update_invoke_agent_span,
 )
 
 if TYPE_CHECKING:
@@ -65,7 +65,9 @@ def _maybe_start_agent_span(
     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)
+            _AgentInvocationSpanContext(
+                context_wrapper=context_wrapper, agent=current_agent
+            ).__exit__(None, None, None)
 
     # Store the agent on the context wrapper so we can access it later
     context_wrapper._sentry_current_agent = agent
@@ -103,15 +105,23 @@ async def _run_single_turn(
         context_wrapper, agent, should_run_agent_start_hooks, kwargs
     )
 
+    if span is None or span.timestamp is not None:
+        return await original_run_single_turn(*args, **kwargs)
+
     try:
-        result = await original_run_single_turn(*args, **kwargs)
+        return await original_run_single_turn(*args, **kwargs)
     except Exception:
-        if span is not None and span.timestamp is None:
-            set_span_errored(span)
-            end_invoke_agent_span(context_wrapper, agent)
-        reraise(*sys.exc_info())
+        exc_info = sys.exc_info()
+        with capture_internal_exceptions():
+            span = getattr(context_wrapper, "_sentry_agent_span", None)
+            if not span:
+                return
 
-    return result
+            update_invoke_agent_span(span=span, context=context_wrapper, agent=agent)
+
+            span.__exit__(*exc_info)
+            delattr(context_wrapper, "_sentry_agent_span")
+        reraise(*exc_info)
 
 
 async def _run_single_turn_streamed(
@@ -174,18 +184,20 @@ async def _run_single_turn_streamed(
         is_streaming=True,
     )
 
-    try:
-        result = await original_run_single_turn_streamed(*args, **kwargs)
-    except Exception:
-        exc_info = sys.exc_info()
-        with capture_internal_exceptions():
-            if span is not None and span.timestamp is None:
-                set_span_errored(span)
-                end_invoke_agent_span(context_wrapper, agent)
-            _close_streaming_workflow_span(agent)
-        reraise(*exc_info)
+    if span is None or span.timestamp is not None:
+        return await original_run_single_turn_streamed(*args, **kwargs)
 
-    return result
+    with _AgentInvocationSpanContext(
+        context_wrapper=context_wrapper,
+        agent=agent,
+    ):
+        try:
+            return await original_run_single_turn_streamed(*args, **kwargs)
+        except Exception:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _close_streaming_workflow_span(agent)
+            reraise(*exc_info)
 
 
 async def _execute_handoffs(
@@ -211,18 +223,30 @@ async def _execute_handoffs(
         handoff_agent_name = first_handoff.handoff.agent_name
         handoff_span(context_wrapper, agent, handoff_agent_name)
 
-    # Call original method with all parameters
-    try:
-        result = await original_execute_handoffs(*args, **kwargs)
-    except Exception:
-        exc_info = sys.exc_info()
-        with capture_internal_exceptions():
-            _close_streaming_workflow_span(agent)
-        reraise(*exc_info)
-    finally:
-        # 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)
+    if not agent or not context_wrapper or not _has_active_agent_span(context_wrapper):
+        # Call original method with all parameters
+        try:
+            result = await original_execute_handoffs(*args, **kwargs)
+        except Exception:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _close_streaming_workflow_span(agent)
+            reraise(*exc_info)
+
+        return result
+
+    with _AgentInvocationSpanContext(
+        context_wrapper=context_wrapper,
+        agent=agent,
+    ):
+        # Call original method with all parameters
+        try:
+            result = await original_execute_handoffs(*args, **kwargs)
+        except Exception:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _close_streaming_workflow_span(agent)
+            reraise(*exc_info)
 
     return result
 
@@ -243,13 +267,22 @@ async def _execute_final_output(
     context_wrapper = kwargs.get("context_wrapper")
     final_output = kwargs.get("final_output")
 
-    try:
-        result = await original_execute_final_output(*args, **kwargs)
-    finally:
-        with capture_internal_exceptions():
-            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
-                end_invoke_agent_span(context_wrapper, agent, final_output)
-            # For streaming, close the workflow span (non-streaming uses context manager in _create_run_wrapper)
-            _close_streaming_workflow_span(agent)
+    if not agent or not context_wrapper or not _has_active_agent_span(context_wrapper):
+        try:
+            result = await original_execute_final_output(*args, **kwargs)
+        finally:
+            with capture_internal_exceptions():
+                # For streaming, close the workflow span (non-streaming uses context manager in _create_run_wrapper)
+                _close_streaming_workflow_span(agent)
+
+    with _AgentInvocationSpanContext(
+        context_wrapper=context_wrapper, agent=agent, output=final_output
+    ):
+        try:
+            result = await original_execute_final_output(*args, **kwargs)
+        finally:
+            with capture_internal_exceptions():
+                # For streaming, close the workflow span (non-streaming uses context manager in _create_run_wrapper)
+                _close_streaming_workflow_span(agent)
 
     return result

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -4,10 +4,9 @@
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
-from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import capture_internal_exceptions, reraise
 
-from ..spans import agent_workflow_span, end_invoke_agent_span
+from ..spans import agent_workflow_span, update_invoke_agent_span
 from ..utils import _capture_exception
 
 try:
@@ -66,8 +65,14 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                                 invoke_agent_span is not None
                                 and invoke_agent_span.timestamp is None
                             ):
-                                set_span_errored(invoke_agent_span)
-                                end_invoke_agent_span(context_wrapper, agent)
+                                update_invoke_agent_span(
+                                    span=invoke_agent_span,
+                                    context=context_wrapper,
+                                    agent=agent,
+                                )
+
+                                invoke_agent_span.__exit__(*exc_info)
+                                delattr(context_wrapper, "_sentry_agent_span")
                     reraise(*exc_info)
                 except Exception as exc:
                     exc_info = sys.exc_info()
@@ -78,7 +83,20 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                         _capture_exception(exc)
                     reraise(*exc_info)
 
-                end_invoke_agent_span(run_result.context_wrapper, agent)
+                invoke_agent_span = getattr(
+                    run_result.context_wrapper, "_sentry_agent_span", None
+                )
+                if not invoke_agent_span:
+                    return run_result
+
+                update_invoke_agent_span(
+                    span=invoke_agent_span,
+                    context=run_result.context_wrapper,
+                    agent=agent,
+                )
+
+                invoke_agent_span.__exit__(None, None, None)
+                delattr(run_result.context_wrapper, "_sentry_agent_span")
                 return run_result
 
     return wrapper

sentry_sdk/integrations/openai_agents/spans/__init__.py

@@ -3,7 +3,7 @@
 from .execute_tool import execute_tool_span, update_execute_tool_span  # noqa: F401
 from .handoff import handoff_span  # noqa: F401
 from .invoke_agent import (
-    end_invoke_agent_span,  # noqa: F401
+    _AgentInvocationSpanContext,  # noqa: F401
     invoke_agent_span,  # noqa: F401
     update_invoke_agent_span,  # noqa: F401
 )

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -9,7 +9,7 @@
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.utils import safe_serialize
+from sentry_sdk.utils import capture_internal_exceptions, safe_serialize
 
 from ..consts import SPAN_ORIGIN
 from ..utils import _set_agent_data, _set_usage_data
@@ -85,37 +85,64 @@ def invoke_agent_span(
 
 
 def update_invoke_agent_span(
-    context: "agents.RunContextWrapper", agent: "agents.Agent", output: "Any"
+    span: "sentry_sdk.tracing.Span",
+    context: "agents.RunContextWrapper",
+    agent: "agents.Agent",
+    output: "Any" = None,
 ) -> None:
-    span = getattr(context, "_sentry_agent_span", None)
-
-    if span:
-        # Add aggregated usage data from context_wrapper
-        if hasattr(context, "usage"):
-            _set_usage_data(span, context.usage)
+    # Add aggregated usage data from context_wrapper
+    if hasattr(context, "usage"):
+        _set_usage_data(span, context.usage)
 
-        if should_send_default_pii():
-            set_data_normalized(
-                span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False
+    if should_send_default_pii():
+        set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False)
+
+    # Add conversation ID from agent
+    conv_id = getattr(agent, "_sentry_conversation_id", None)
+    if conv_id:
+        span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
+
+
+class _AgentInvocationSpanContext:
+    """
+    Sets accumulated data on the Invoke Agent span and finishes the span on exit.
+    Is a no-op if the context wrapper has no span set, i.e., when the span has already been finished.
+    """
+
+    def __init__(
+        self,
+        context_wrapper: "agents.RunContextWrapper",
+        agent: "agents.Agent",
+        output: "Optional[Any]" = None,
+    ) -> None:
+        self._context_wrapper = context_wrapper
+        self._agent = agent
+        self._output = output
+
+    def __enter__(self) -> "_AgentInvocationSpanContext":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: "Optional[type[BaseException]]",
+        exc_val: "Optional[BaseException]",
+        exc_tb: "Optional[Any]",
+    ) -> None:
+        with capture_internal_exceptions():
+            # Clear the stored agent
+            if hasattr(self._context_wrapper, "_sentry_current_agent"):
+                delattr(self._context_wrapper, "_sentry_current_agent")
+
+            span = getattr(self._context_wrapper, "_sentry_agent_span", None)
+            if not span:
+                return
+
+            update_invoke_agent_span(
+                span=span,
+                context=self._context_wrapper,
+                agent=self._agent,
+                output=self._output,
             )
 
-        # Add conversation ID from agent
-        conv_id = getattr(agent, "_sentry_conversation_id", None)
-        if conv_id:
-            span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
-
-        span.__exit__(None, None, None)
-        delattr(context, "_sentry_agent_span")
-
-
-def end_invoke_agent_span(
-    context_wrapper: "agents.RunContextWrapper",
-    agent: "agents.Agent",
-    output: "Optional[Any]" = None,
-) -> 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)
+            span.__exit__(exc_type, exc_val, exc_tb)
+            delattr(self._context_wrapper, "_sentry_agent_span")

sentry_sdk/integrations/openai_agents/utils.py

@@ -17,7 +17,6 @@
 from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception, safe_serialize
 
 if TYPE_CHECKING:
@@ -35,8 +34,6 @@
 
 
 def _capture_exception(exc: "Any") -> None:
-    set_span_errored()
-
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,