Finish invoke span in case of errors

Author: sentrivana

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -68,7 +68,6 @@ async def patched_run_single_turn(
         cls: agents.Runner, *args: Any, **kwargs: Any
     ) -> Any:
         """Patched _run_single_turn that creates agent invocation spans"""
-
         agent = kwargs.get("agent")
         context_wrapper = kwargs.get("context_wrapper")
         should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks")
@@ -84,7 +83,11 @@ async def patched_run_single_turn(
             _start_invoke_agent_span(context_wrapper, agent)
 
         # 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)
+        finally:
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent)
 
         return result
 
@@ -97,7 +100,6 @@ async def patched_execute_handoffs(
         cls: agents.Runner, *args: Any, **kwargs: Any
     ) -> Any:
         """Patched execute_handoffs that creates handoff spans and ends agent span for handoffs"""
-
         context_wrapper = kwargs.get("context_wrapper")
         run_handoffs = kwargs.get("run_handoffs")
         agent = kwargs.get("agent")
@@ -128,7 +130,6 @@ async def patched_execute_final_output(
         cls: agents.Runner, *args: Any, **kwargs: Any
     ) -> Any:
         """Patched execute_final_output that ends agent span for final outputs"""
-
         agent = kwargs.get("agent")
         context_wrapper = kwargs.get("context_wrapper")
         final_output = kwargs.get("final_output")

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -33,9 +33,9 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
                 _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)
+                span = sentry_sdk.get_current_span()
+                if span is not None and span.timestamp is None:
+                    span.__exit__(None, None, None)
 
                 raise exc from None
 

sentry_sdk/integrations/openai_agents/spans/handoff.py

@@ -12,7 +12,9 @@
 
 
 def handoff_span(
-    context: agents.RunContextWrapper, from_agent: agents.Agent, to_agent_name: str
+    context_wrapper: agents.RunContextWrapper,
+    from_agent: agents.Agent,
+    to_agent_name: str,
 ) -> None:
     with sentry_sdk.start_span(
         op=OP.GEN_AI_HANDOFF,

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -14,7 +14,7 @@
 
 
 def invoke_agent_span(
-    context: agents.RunContextWrapper, agent: agents.Agent
+    context_wrapper: agents.RunContextWrapper, agent: agents.Agent
 ) -> sentry_sdk.tracing.Span:
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_INVOKE_AGENT,
@@ -27,15 +27,15 @@ def invoke_agent_span(
 
     _set_agent_data(span, agent)
 
-    context._sentry_invoke_agent_span = span
+    context_wrapper._sentry_invoke_agent_span = span
 
     return span
 
 
 def update_invoke_agent_span(
-    context: agents.RunContextWrapper, agent: agents.Agent, output: Any
+    context_wrapper: agents.RunContextWrapper, agent: agents.Agent, output: Any
 ) -> None:
-    span = getattr(context, "_sentry_invoke_agent_span", None)
+    span = getattr(context_wrapper, "_sentry_invoke_agent_span", None)
     if span is not None:
         span.__exit__(None, None, None)
-        del context._sentry_invoke_agent_span
+        del context_wrapper._sentry_invoke_agent_span

tests/integrations/openai_agents/test_openai_agents.py

@@ -586,4 +586,4 @@ async def test_error_handling(sentry_init, capture_events, test_agent):
 
     assert ai_client_span["description"] == "chat gpt-4"
     assert ai_client_span["origin"] == "auto.ai.openai_agents"
-    assert ai_client_span["tags"]["status"] == "internal_error"
+    assert ai_client_span["status"] == "internal_error"