feat(integrations): openai-agents streaming support

Author: constantinius

sentry_sdk/integrations/mcp.py

@@ -7,6 +7,7 @@
 Supports the low-level `mcp.server.lowlevel.Server` API.
 """
 
+from contextlib import contextmanager
 import inspect
 from functools import wraps
 from typing import TYPE_CHECKING
@@ -352,6 +353,30 @@ def _prepare_handler_data(
     )
 
 
+@contextmanager
+def ensure_span(*args, **kwargs):
+    """Ensure a span is created for the current context."""
+
+    current_span = sentry_sdk.get_current_span()
+    transaction_exists = (
+        current_span is not None and current_span.containing_transaction is not None
+    )
+
+    if transaction_exists:
+        with sentry_sdk.start_span(*args, **kwargs) as span:
+            yield span
+    else:
+        with sentry_sdk.start_transaction(*args, **kwargs):
+            with sentry_sdk.start_span(*args, **kwargs) as span:
+                yield span
+    # with get_start_span_function()(
+    #     op=OP.MCP_SERVER,
+    #     name=span_name,
+    #     origin=MCPIntegration.origin,
+    # ) as span:
+    #     yield span
+
+
 async def _async_handler_wrapper(
     handler_type: str,
     func: "Callable[..., Any]",
@@ -382,7 +407,7 @@ async def _async_handler_wrapper(
     ) = _prepare_handler_data(handler_type, original_args, original_kwargs)
 
     # Start span and execute
-    with get_start_span_function()(
+    with ensure_span(
         op=OP.MCP_SERVER,
         name=span_name,
         origin=MCPIntegration.origin,
@@ -454,7 +479,7 @@ def _sync_handler_wrapper(
     ) = _prepare_handler_data(handler_type, original_args)
 
     # Start span and execute
-    with get_start_span_function()(
+    with ensure_span(
         op=OP.MCP_SERVER,
         name=span_name,
         origin=MCPIntegration.origin,

sentry_sdk/integrations/openai_agents/__init__.py

@@ -4,6 +4,7 @@
     _create_get_model_wrapper,
     _create_get_all_tools_wrapper,
     _create_run_wrapper,
+    _create_run_streamed_wrapper,
     _patch_agent_run,
     _patch_error_tracing,
 )
@@ -25,12 +26,16 @@ def _patch_runner() -> None:
     # Create the root span for one full agent run (including eventual handoffs)
     # Note agents.run.DEFAULT_AGENT_RUNNER.run_sync is a wrapper around
     # agents.run.DEFAULT_AGENT_RUNNER.run. It does not need to be wrapped separately.
-    # TODO-anton: Also patch streaming runner: agents.Runner.run_streamed
     agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
         agents.run.DEFAULT_AGENT_RUNNER.run
     )
 
-    # Creating the actual spans for each agent run.
+    # Patch streaming runner
+    agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper(
+        agents.run.DEFAULT_AGENT_RUNNER.run_streamed
+    )
+
+    # Creating the actual spans for each agent run (works for both streaming and non-streaming).
     _patch_agent_run()
 
 

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -1,5 +1,5 @@
 from .models import _create_get_model_wrapper  # noqa: F401
 from .tools import _create_get_all_tools_wrapper  # noqa: F401
-from .runner import _create_run_wrapper  # noqa: F401
+from .runner import _create_run_wrapper, _create_run_streamed_wrapper  # noqa: F401
 from .agent_run import _patch_agent_run  # noqa: F401
 from .error_tracing import _patch_error_tracing  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -31,6 +31,7 @@ def _patch_agent_run() -> None:
 
     # Store original methods
     original_run_single_turn = agents.run.AgentRunner._run_single_turn
+    original_run_single_turn_streamed = agents.run.AgentRunner._run_single_turn_streamed
     original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
     original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
 
@@ -149,8 +150,48 @@ async def patched_execute_final_output(
 
         return result
 
+    @wraps(
+        original_run_single_turn_streamed.__func__
+        if hasattr(original_run_single_turn_streamed, "__func__")
+        else original_run_single_turn_streamed
+    )
+    async def patched_run_single_turn_streamed(
+        cls: "agents.Runner", *args: "Any", **kwargs: "Any"
+    ) -> "Any":
+        """Patched _run_single_turn_streamed that creates agent invocation spans for streaming"""
+        agent = kwargs.get("agent")
+        context_wrapper = kwargs.get("context_wrapper")
+        should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks")
+
+        span = getattr(context_wrapper, "_sentry_agent_span", None)
+        # Start agent span when agent starts (but only once per agent)
+        if should_run_agent_start_hooks and agent and context_wrapper:
+            # End any existing span for a different agent
+            if _has_active_agent_span(context_wrapper):
+                current_agent = _get_current_agent(context_wrapper)
+                if current_agent and current_agent != agent:
+                    end_invoke_agent_span(context_wrapper, current_agent)
+
+            span = _start_invoke_agent_span(context_wrapper, agent, kwargs)
+            agent._sentry_agent_span = span
+
+        # Call original streaming method
+        try:
+            result = await original_run_single_turn_streamed(*args, **kwargs)
+        except Exception as exc:
+            if span is not None and span.timestamp is None:
+                _record_exception_on_span(span, exc)
+                end_invoke_agent_span(context_wrapper, agent)
+
+            reraise(*sys.exc_info())
+
+        return result
+
     # Apply patches
     agents.run.AgentRunner._run_single_turn = classmethod(patched_run_single_turn)
+    agents.run.AgentRunner._run_single_turn_streamed = classmethod(
+        patched_run_single_turn_streamed
+    )
     agents._run_impl.RunImpl.execute_handoffs = classmethod(patched_execute_handoffs)
     agents._run_impl.RunImpl.execute_final_output = classmethod(
         patched_execute_final_output

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -1,3 +1,4 @@
+import sys
 from functools import wraps
 
 import sentry_sdk
@@ -64,3 +65,74 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
                 return run_result
 
     return wrapper
+
+
+def _create_run_streamed_wrapper(
+    original_func: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Wraps the agents.Runner.run_streamed method to create a root span for streaming agent workflow runs.
+
+    Unlike run(), run_streamed() returns immediately with a RunResultStreaming object
+    while execution continues in a background task. The workflow span must stay open
+    throughout the streaming operation and close when streaming completes or is abandoned.
+    """
+
+    @wraps(original_func)
+    def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
+        # Isolate each workflow so that when agents are run in asyncio tasks they
+        # don't touch each other's scopes
+        isolation_scope = sentry_sdk.isolation_scope()
+        isolation_scope.__enter__()
+
+        # Clone agent because agent invocation spans are attached per run.
+        agent = args[0].clone()
+
+        # Start workflow span immediately (before run_streamed returns)
+        workflow_span = agent_workflow_span(agent)
+        workflow_span.__enter__()
+
+        # Store span and scope on agent for cleanup
+        agent._sentry_workflow_span = workflow_span
+        agent._sentry_isolation_scope = isolation_scope
+
+        args = (agent, *args[1:])
+
+        try:
+            # Call original function to get RunResultStreaming
+            run_result = original_func(*args, **kwargs)
+        except Exception as exc:
+            # If run_streamed itself fails (not the background task), clean up immediately
+            workflow_span.__exit__(*sys.exc_info())
+            isolation_scope.__exit__(None, None, None)
+            _capture_exception(exc)
+            raise exc from None
+
+        # Wrap the result to ensure cleanup when streaming completes
+        original_aclose = getattr(run_result, "aclose", None)
+
+        async def wrapped_aclose() -> None:
+            """Close streaming result and clean up Sentry spans"""
+            try:
+                if original_aclose is not None:
+                    await original_aclose()
+            finally:
+                # End any remaining agent span
+                if hasattr(run_result, "context_wrapper"):
+                    end_invoke_agent_span(run_result.context_wrapper, agent)
+
+                # End workflow span
+                if hasattr(agent, "_sentry_workflow_span"):
+                    workflow_span.__exit__(None, None, None)
+                    delattr(agent, "_sentry_workflow_span")
+
+                # Exit isolation scope
+                if hasattr(agent, "_sentry_isolation_scope"):
+                    isolation_scope.__exit__(None, None, None)
+                    delattr(agent, "_sentry_isolation_scope")
+
+        run_result.aclose = wrapped_aclose
+
+        return run_result
+
+    return wrapper

tests/integrations/openai_agents/test_openai_agents.py

@@ -1998,3 +1998,24 @@ def test_openai_agents_message_truncation(sentry_init, capture_events):
         assert len(parsed_messages) == 2
         assert "small message 4" in str(parsed_messages[0])
         assert "small message 5" in str(parsed_messages[1])
+
+
+def test_streaming_patches_applied(sentry_init):
+    """
+    Test that the streaming patches are applied correctly.
+    """
+    sentry_init(
+        integrations=[OpenAIAgentsIntegration()],
+        traces_sample_rate=1.0,
+    )
+
+    # Verify that run_streamed is patched (will have __wrapped__ attribute if patched)
+    import agents
+
+    # Check that the method exists and has been modified
+    assert hasattr(agents.run.DEFAULT_AGENT_RUNNER, "run_streamed")
+    assert hasattr(agents.run.AgentRunner, "_run_single_turn_streamed")
+
+    # Verify the patches were applied by checking for our wrapper
+    run_streamed_func = agents.run.DEFAULT_AGENT_RUNNER.run_streamed
+    assert run_streamed_func is not None