feat(openai-agents): Support span streaming

alexander-alderman-webb · alexander-alderman-webb · commit 0f53ac0f13b9 · 2026-05-26T08:09:28.000+02:00
diff --git a/sentry_sdk/integrations/openai_agents/patches/agent_run.py b/sentry_sdk/integrations/openai_agents/patches/agent_run.py
@@ -3,6 +3,7 @@
 
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions, reraise
 
 from ..spans import (
@@ -108,7 +109,13 @@ async def _run_single_turn(
         context_wrapper, agent, should_run_agent_start_hooks, kwargs
     )
 
-    if span is None or span.timestamp is not None:
+    if (
+        span is None
+        or isinstance(span, StreamedSpan)
+        and span.end_timestamp is not None
+        or not isinstance(span, StreamedSpan)
+        and span.timestamp is not None
+    ):
         return await original_run_single_turn(*args, **kwargs)
 
     try:
@@ -188,7 +195,13 @@ async def _run_single_turn_streamed(
         is_streaming=True,
     )
 
-    if span is None or span.timestamp is not None:
+    if (
+        span is None
+        or isinstance(span, StreamedSpan)
+        and span.end_timestamp is not None
+        or not isinstance(span, StreamedSpan)
+        and span.timestamp is not None
+    ):
         return await original_run_single_turn_streamed(*args, **kwargs)
 
     try:
diff --git a/sentry_sdk/integrations/openai_agents/patches/models.py b/sentry_sdk/integrations/openai_agents/patches/models.py
@@ -6,6 +6,7 @@
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
 from sentry_sdk.tracing_utils import (
     add_sentry_baggage_to_headers,
@@ -34,7 +35,10 @@ def _set_response_model_on_agent_span(
     if response_model:
         agent_span = getattr(agent, "_sentry_agent_span", None)
         if agent_span:
-            agent_span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
+            if isinstance(agent_span, StreamedSpan):
+                agent_span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
+            else:
+                agent_span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
 
 
 def _inject_trace_propagation_headers(
@@ -151,7 +155,12 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any":
                 for hosted_tool in hosted_tools:
                     _inject_trace_propagation_headers(hosted_tool, span=span)
 
-                span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+                set_on_span = (
+                    span.set_attribute
+                    if isinstance(span, StreamedSpan)
+                    else span.set_data
+                )
+                set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
 
                 streaming_response = None
                 ttft_recorded = False
@@ -162,9 +171,7 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any":
                     # Detect first content token (text delta event)
                     if not ttft_recorded and hasattr(event, "delta"):
                         ttft = time.perf_counter() - start_time
-                        span.set_data(
-                            SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft
-                        )
+                        set_on_span(SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft)
                         ttft_recorded = True
 
                     # Capture the full response from ResponseCompletedEvent
diff --git a/sentry_sdk/integrations/openai_agents/patches/runner.py b/sentry_sdk/integrations/openai_agents/patches/runner.py
@@ -4,6 +4,7 @@
 import sentry_sdk
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions, reraise
 
 from ..spans import agent_workflow_span, update_invoke_agent_span
@@ -43,9 +44,15 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                 conversation_id = kwargs.get("conversation_id")
                 if conversation_id:
                     agent._sentry_conversation_id = conversation_id
-                    workflow_span.set_data(
-                        SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
-                    )
+
+                    if isinstance(workflow_span, StreamedSpan):
+                        workflow_span.set_attribute(
+                            SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
+                        )
+                    else:
+                        workflow_span.set_data(
+                            SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
+                        )
 
                 args = (agent, *args[1:])
                 try:
diff --git a/sentry_sdk/integrations/openai_agents/spans/agent_workflow.py b/sentry_sdk/integrations/openai_agents/spans/agent_workflow.py
@@ -2,6 +2,7 @@
 
 import sentry_sdk
 from sentry_sdk.ai.utils import get_start_span_function
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 from ..consts import SPAN_ORIGIN
 
@@ -11,6 +12,14 @@
 
 def agent_workflow_span(agent: "agents.Agent") -> "sentry_sdk.tracing.Span":
     # Create a transaction or a span if an transaction is already active
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"{agent.name} workflow", attributes={"sentry.origin": SPAN_ORIGIN}
+        )
+
+        return span
+
     span = get_start_span_function()(
         name=f"{agent.name} workflow",
         origin=SPAN_ORIGIN,
diff --git a/sentry_sdk/integrations/openai_agents/spans/ai_client.py b/sentry_sdk/integrations/openai_agents/spans/ai_client.py
@@ -2,6 +2,8 @@
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 from ..consts import SPAN_ORIGIN
 from ..utils import (
@@ -12,7 +14,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Optional, Union
 
     from agents import Agent
 
@@ -28,13 +30,24 @@ def ai_client_span(
     elif hasattr(agent, "_sentry_request_model"):
         model_name = agent._sentry_request_model
 
-    span = sentry_sdk.start_span(
-        op=OP.GEN_AI_CHAT,
-        name=f"chat {model_name}",
-        origin=SPAN_ORIGIN,
-    )
-    # TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on
-    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"chat {model_name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_CHAT,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "chat",
+            },
+        )
+    else:
+        span = sentry_sdk.start_span(
+            op=OP.GEN_AI_CHAT,
+            name=f"chat {model_name}",
+            origin=SPAN_ORIGIN,
+        )
+        # TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
 
     _set_agent_data(span, agent)
     _set_input_data(span, get_response_kwargs)
@@ -43,7 +56,7 @@ def ai_client_span(
 
 
 def update_ai_client_span(
-    span: "sentry_sdk.tracing.Span",
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
     response: "Any",
     response_model: "Optional[str]" = None,
     agent: "Optional[Agent]" = None,
@@ -55,13 +68,17 @@ def update_ai_client_span(
     if hasattr(response, "output") and response.output:
         _set_output_data(span, response)
 
+    set_on_span = (
+        span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
+    )
+
     if response_model is not None:
-        span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
+        set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
     elif hasattr(response, "model") and response.model:
-        span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, str(response.model))
+        set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, str(response.model))
 
     # Set conversation ID from agent if available
     if agent:
         conv_id = getattr(agent, "_sentry_conversation_id", None)
         if conv_id:
-            span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
+            set_on_span(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
diff --git a/sentry_sdk/integrations/openai_agents/spans/execute_tool.py b/sentry_sdk/integrations/openai_agents/spans/execute_tool.py
@@ -3,39 +3,58 @@
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import SpanStatus, StreamedSpan
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 from ..consts import SPAN_ORIGIN
 from ..utils import _set_agent_data
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Union
 
     import agents
 
 
 def execute_tool_span(
     tool: "agents.Tool", *args: "Any", **kwargs: "Any"
 ) -> "sentry_sdk.tracing.Span":
-    span = sentry_sdk.start_span(
-        op=OP.GEN_AI_EXECUTE_TOOL,
-        name=f"execute_tool {tool.name}",
-        origin=SPAN_ORIGIN,
-    )
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"execute_tool {tool.name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_EXECUTE_TOOL,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool",
+                SPANDATA.GEN_AI_TOOL_NAME: tool.name,
+                SPANDATA.GEN_AI_TOOL_DESCRIPTION: tool.description,
+            },
+        )
+
+        set_on_span = span.set_attribute
+    else:
+        span = sentry_sdk.start_span(
+            op=OP.GEN_AI_EXECUTE_TOOL,
+            name=f"execute_tool {tool.name}",
+            origin=SPAN_ORIGIN,
+        )
+
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
 
-    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
+        span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
+        span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
 
-    span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
-    span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
+        set_on_span = span.set_data
 
     if should_send_default_pii():
         input = args[1]
-        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, input)
+        set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, input)
 
     return span
 
 
 def update_execute_tool_span(
-    span: "sentry_sdk.tracing.Span",
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
     agent: "agents.Agent",
     tool: "agents.Tool",
     result: "Any",
@@ -45,12 +64,19 @@ def update_execute_tool_span(
     if isinstance(result, str) and result.startswith(
         "An error occurred while running the tool"
     ):
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        if isinstance(span, StreamedSpan):
+            span.status = SpanStatus.ERROR
+        else:
+            span.set_status(SPANSTATUS.INTERNAL_ERROR)
+
+    set_on_span = (
+        span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
+    )
 
     if should_send_default_pii():
-        span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)
+        set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, result)
 
     # 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)
+        set_on_span(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
diff --git a/sentry_sdk/integrations/openai_agents/spans/handoff.py b/sentry_sdk/integrations/openai_agents/spans/handoff.py
@@ -2,6 +2,7 @@
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 from ..consts import SPAN_ORIGIN
 
@@ -12,14 +13,29 @@
 def handoff_span(
     context: "agents.RunContextWrapper", from_agent: "agents.Agent", to_agent_name: str
 ) -> None:
-    with sentry_sdk.start_span(
-        op=OP.GEN_AI_HANDOFF,
-        name=f"handoff from {from_agent.name} to {to_agent_name}",
-        origin=SPAN_ORIGIN,
-    ) as span:
-        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        with sentry_sdk.traces.start_span(
+            name=f"handoff from {from_agent.name} to {to_agent_name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_HANDOFF,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "handoff",
+            },
+        ) as span:
+            # Add conversation ID from agent
+            conv_id = getattr(from_agent, "_sentry_conversation_id", None)
+            if conv_id:
+                span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
+    else:
+        with sentry_sdk.start_span(
+            op=OP.GEN_AI_HANDOFF,
+            name=f"handoff from {from_agent.name} to {to_agent_name}",
+            origin=SPAN_ORIGIN,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")
 
-        # Add conversation ID from agent
-        conv_id = getattr(from_agent, "_sentry_conversation_id", None)
-        if conv_id:
-            span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
+            # Add conversation ID from agent
+            conv_id = getattr(from_agent, "_sentry_conversation_id", None)
+            if conv_id:
+                span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py
diff --git a/sentry_sdk/integrations/openai_agents/utils.py b/sentry_sdk/integrations/openai_agents/utils.py
diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py
