.

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -13,11 +13,10 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Awaitable, Callable, Optional, Union
+    from typing import Any, Awaitable, Callable, Optional
 
     from agents.run_internal.run_steps import SingleStepResult
 
-    from sentry_sdk.tracing import Span
 
 try:
     import agents
@@ -51,7 +50,7 @@ def _maybe_start_agent_span(
     should_run_agent_start_hooks: bool,
     span_kwargs: "dict[str, Any]",
     is_streaming: bool = False,
-) -> "Optional[Union[Span, StreamedSpan]]":
+) -> "Optional[StreamedSpan]":
     """
     Start an agent invocation span if conditions are met.
     Handles ending any existing span for a different agent.
@@ -82,10 +81,7 @@ def _maybe_start_agent_span(
     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)
+    span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
 
     return span
 
@@ -114,11 +110,7 @@ async def _run_single_turn(
         context_wrapper, agent, should_run_agent_start_hooks, kwargs
     )
 
-    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)
-    ):
+    if span is None or span.end_timestamp is not None:
         return await original_run_single_turn(*args, **kwargs)
 
     try:
@@ -198,11 +190,7 @@ async def _run_single_turn_streamed(
         is_streaming=True,
     )
 
-    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)
-    ):
+    if span is None or span.end_timestamp is not None:
         return await original_run_single_turn_streamed(*args, **kwargs)
 
     try:

sentry_sdk/integrations/openai_agents/patches/error_tracing.py

@@ -2,9 +2,7 @@
 from typing import TYPE_CHECKING
 
 import sentry_sdk
-from sentry_sdk.consts import SPANSTATUS
 from sentry_sdk.traces import SpanStatus
-from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 if TYPE_CHECKING:
     from typing import Any
@@ -56,15 +54,9 @@ def sentry_attach_error_to_current_span(
         the agents library swallows exceptions.
         """
         # Set the current Sentry span to errored
-        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
-        if span_streaming:
-            current_span = sentry_sdk.get_current_scope().streamed_span
-            if current_span is not None:
-                current_span.status = SpanStatus.ERROR
-        else:
-            current_span = sentry_sdk.get_current_span()
-            if current_span is not None:
-                current_span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        current_span = sentry_sdk.get_current_scope().streamed_span
+        if current_span is not None:
+            current_span.status = SpanStatus.ERROR
 
         # Call the original function
         return original_attach_error(error, *args, **kwargs)

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -17,9 +17,8 @@
 from ..spans import ai_client_span, update_ai_client_span
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Union
+    from typing import Any, Callable
 
-    from sentry_sdk.tracing import Span
 
 try:
     import agents
@@ -29,7 +28,7 @@
 
 
 def _inject_trace_propagation_headers(
-    hosted_tool: "HostedMCPTool", span: "Union[Span, StreamedSpan]"
+    hosted_tool: "HostedMCPTool", span: "StreamedSpan"
 ) -> None:
     headers = hosted_tool.tool_config.get("headers")
     if headers is None:
@@ -141,12 +140,7 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any":
                 for hosted_tool in hosted_tools:
                     _inject_trace_propagation_headers(hosted_tool, span=span)
 
-                set_on_span = (
-                    span.set_attribute
-                    if isinstance(span, StreamedSpan)
-                    else span.set_data
-                )
-                set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+                span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
 
                 streaming_response = None
                 ttft_recorded = False
@@ -157,7 +151,9 @@ 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
-                        set_on_span(SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft)
+                        span.set_attribute(
+                            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

@@ -64,10 +64,7 @@ async def on_tool_start(
         elif not should_send_default_pii():
             return
 
-        if isinstance(span, StreamedSpan):
-            span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, context.tool_arguments)
-        else:
-            span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, context.tool_arguments)
+        span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, context.tool_arguments)
 
     async def on_tool_end(
         self,
@@ -161,14 +158,9 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                 if conversation_id:
                     agent._sentry_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
-                        )
+                    workflow_span.set_attribute(
+                        SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
+                    )
 
                 if "starting_agent" in kwargs:
                     kwargs["starting_agent"] = agent
@@ -188,15 +180,9 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                                 context_wrapper, "_sentry_agent_span", 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
-                                )
+                            if (
+                                invoke_agent_span is not None
+                                and invoke_agent_span.end_timestamp is None
                             ):
                                 update_invoke_agent_span(
                                     span=invoke_agent_span,

sentry_sdk/integrations/openai_agents/patches/tools.py

@@ -5,7 +5,6 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import has_data_collection_enabled
 
 from ..spans import execute_tool_span, update_execute_tool_span
@@ -65,10 +64,7 @@ async def sentry_wrapped_on_invoke_tool(
                     elif not should_send_default_pii():
                         return result
 
-                    if isinstance(span, StreamedSpan):
-                        span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, args[1])
-                    else:
-                        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, args[1])
+                    span.set_attribute(SPANDATA.GEN_AI_TOOL_INPUT, args[1])
 
                 return result
 

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

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

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -10,14 +10,14 @@
 from ..utils import _set_agent_data
 
 if TYPE_CHECKING:
-    from typing import Any, Union
+    from typing import Any
 
     import agents
 
 
 def execute_tool_span(
     tool: "agents.Tool", *args: "Any", **kwargs: "Any"
-) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
+) -> "StreamedSpan":
     span = sentry_sdk.traces.start_span(
         name=f"execute_tool {tool.name}",
         attributes={
@@ -33,7 +33,7 @@ def execute_tool_span(
 
 
 def update_execute_tool_span(
-    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
+    span: "StreamedSpan",
     agent: "agents.Agent",
     tool: "agents.Tool",
     result: "Any",

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -18,14 +18,14 @@
 from ..utils import _set_agent_data, _set_usage_data
 
 if TYPE_CHECKING:
-    from typing import Any, Union
+    from typing import Any
 
     import agents
 
 
 def invoke_agent_span(
     context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]"
-) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
+) -> "StreamedSpan":
     client_options = sentry_sdk.get_client().options
     span = sentry_sdk.traces.start_span(
         name=f"invoke_agent {agent.name}",
@@ -95,7 +95,7 @@ def invoke_agent_span(
 
 
 def update_invoke_agent_span(
-    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
+    span: "StreamedSpan",
     context: "agents.RunContextWrapper",
     agent: "agents.Agent",
     output: "Any" = None,