feat(openai-agents): Support span streaming

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 (
@@ -12,7 +13,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Awaitable, Callable, Optional
+    from typing import Any, Awaitable, Callable, Optional, Union
 
     from agents.run_internal.run_steps import SingleStepResult
 
@@ -50,7 +51,7 @@ def _maybe_start_agent_span(
     should_run_agent_start_hooks: bool,
     span_kwargs: "dict[str, Any]",
     is_streaming: bool = False,
-) -> "Optional[Span]":
+) -> "Optional[Union[Span, StreamedSpan]]":
     """
     Start an agent invocation span if conditions are met.
     Handles ending any existing span for a different agent.
@@ -78,7 +79,12 @@ def _maybe_start_agent_span(
     context_wrapper._sentry_agent_span = span
     agent._sentry_agent_span = span
 
-    if is_streaming:
+    if not is_streaming:
+        return span
+
+    if isinstance(span, StreamedSpan):
+        span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+    else:
         span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
 
     return span
@@ -108,7 +114,11 @@ 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 +198,11 @@ 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:

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,
@@ -16,7 +17,7 @@
 from ..spans import ai_client_span, update_ai_client_span
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional
+    from typing import Any, Callable, Optional, Union
 
     from sentry_sdk.tracing import Span
 
@@ -34,11 +35,14 @@ 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(
-    hosted_tool: "HostedMCPTool", span: "Span"
+    hosted_tool: "HostedMCPTool", span: "Union[Span, StreamedSpan]"
 ) -> None:
     headers = hosted_tool.tool_config.get("headers")
     if headers is None:
@@ -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

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:
@@ -61,8 +68,10 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                                 context_wrapper, "_sentry_agent_span", None
                             )
 
-                            if (
-                                invoke_agent_span is not None
+                            if invoke_agent_span is not None and (
+                                isinstance(invoke_agent_span, StreamedSpan)
+                                and invoke_agent_span.end_timestamp is None
+                                or not isinstance(invoke_agent_span, StreamedSpan)
                                 and invoke_agent_span.timestamp is None
                             ):
                                 update_invoke_agent_span(
@@ -135,7 +144,12 @@ def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
 
         # Set conversation ID on workflow span early so it's captured even on errors
         if 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)
 
         # Store span on agent for cleanup
         agent._sentry_workflow_span = workflow_span

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

@@ -2,15 +2,28 @@
 
 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
 
 if TYPE_CHECKING:
+    from typing import Union
+
     import agents
 
 
-def agent_workflow_span(agent: "agents.Agent") -> "sentry_sdk.tracing.Span":
+def agent_workflow_span(
+    agent: "agents.Agent",
+) -> "Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]":
     # 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,

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,14 +14,14 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Optional, Union
 
     from agents import Agent
 
 
 def ai_client_span(
     agent: "Agent", get_response_kwargs: "dict[str, Any]"
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
     # Get model name from agent.model or fall back to request model (for when agent.model is None/default)
     model_name = None
@@ -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)

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,
-    )
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
+    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)