mypy2

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -13,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
 
@@ -51,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.
@@ -79,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

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -17,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
 
@@ -42,7 +42,7 @@ def _set_response_model_on_agent_span(
 
 
 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:

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -142,7 +142,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