improve the concurrency of event handling

Author: seratch

src/agents/agent.py

@@ -5,10 +5,10 @@
 import inspect
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypeAlias, cast
 
 from openai.types.responses.response_prompt_param import ResponsePromptParam
-from typing_extensions import NotRequired, TypeAlias, TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from .agent_output import AgentOutputSchemaBase
 from .guardrail import InputGuardrail, OutputGuardrail
@@ -457,12 +457,12 @@ async def run_agent(context: RunContextWrapper, input: str) -> Any:
                     conversation_id=conversation_id,
                     session=session,
                 )
-                async for event in run_result.stream_events():
-                    payload: AgentToolStreamEvent = {
-                        "event": event,
-                        "agent": self,
-                        "tool_call": getattr(context, "tool_call", None),
-                    }
+                # Dispatch callbacks in the background so slow handlers do not block
+                # event consumption.
+                event_queue: asyncio.Queue[AgentToolStreamEvent | None] = asyncio.Queue()
+
+                async def _run_handler(payload: AgentToolStreamEvent) -> None:
+                    """Execute the user callback while capturing exceptions."""
                     try:
                         maybe_result = on_stream(payload)
                         if inspect.isawaitable(maybe_result):
@@ -472,6 +472,34 @@ async def run_agent(context: RunContextWrapper, input: str) -> Any:
                             "Error while handling on_stream event for agent tool %s.",
                             self.name,
                         )
+
+                async def dispatch_stream_events() -> None:
+                    while True:
+                        payload = await event_queue.get()
+                        is_sentinel = payload is None  # None marks the end of the stream.
+                        try:
+                            if payload is not None:
+                                await _run_handler(payload)
+                        finally:
+                            event_queue.task_done()
+
+                        if is_sentinel:
+                            break
+
+                dispatch_task = asyncio.create_task(dispatch_stream_events())
+
+                try:
+                    async for event in run_result.stream_events():
+                        payload: AgentToolStreamEvent = {
+                            "event": event,
+                            "agent": self,
+                            "tool_call": getattr(context, "tool_call", None),
+                        }
+                        await event_queue.put(payload)
+                finally:
+                    await event_queue.put(None)
+                    await event_queue.join()
+                    await dispatch_task
             else:
                 run_result = await Runner.run(
                     starting_agent=self,

tests/test_agent_as_tool.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from typing import Any, cast
 
 import pytest
@@ -612,6 +613,90 @@ def sync_handler(event: AgentToolStreamEvent) -> None:
     assert calls == ["raw_response_event"]
 
 
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_dispatches_without_blocking(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """on_stream handlers should not block streaming iteration."""
+    agent = Agent(name="nonblocking_agent")
+
+    first_handler_started = asyncio.Event()
+    allow_handler_to_continue = asyncio.Event()
+    second_event_yielded = asyncio.Event()
+    second_event_handled = asyncio.Event()
+
+    first_event = RawResponsesStreamEvent(data=cast(Any, {"type": "response_started"}))
+    second_event = RawResponsesStreamEvent(
+        data=cast(Any, {"type": "output_text_delta", "delta": "hi"})
+    )
+
+    class DummyStreamingResult:
+        def __init__(self) -> None:
+            self.final_output = "ok"
+
+        async def stream_events(self):
+            yield first_event
+            second_event_yielded.set()
+            yield second_event
+
+    dummy_result = DummyStreamingResult()
+
+    monkeypatch.setattr(
+        Runner, "run_streamed", classmethod(lambda *args, **kwargs: dummy_result)
+    )
+    monkeypatch.setattr(
+        Runner,
+        "run",
+        classmethod(lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("no run"))),
+    )
+
+    async def on_stream(payload: AgentToolStreamEvent) -> None:
+        if payload["event"] is first_event:
+            first_handler_started.set()
+            await allow_handler_to_continue.wait()
+        else:
+            second_event_handled.set()
+
+    tool_call = ResponseFunctionToolCall(
+        id="call_nonblocking",
+        arguments='{"input": "go"}',
+        call_id="call-nonblocking",
+        name="nonblocking_tool",
+        type="function_call",
+    )
+
+    tool = cast(
+        FunctionTool,
+        agent.as_tool(
+            tool_name="nonblocking_tool",
+            tool_description="Uses non-blocking streaming handler",
+            on_stream=on_stream,
+        ),
+    )
+    tool_context = ToolContext(
+        context=None,
+        tool_name="nonblocking_tool",
+        tool_call_id=tool_call.call_id,
+        tool_arguments=tool_call.arguments,
+        tool_call=tool_call,
+    )
+
+    async def _invoke_tool() -> Any:
+        return await tool.on_invoke_tool(tool_context, '{"input": "go"}')
+
+    invoke_task: asyncio.Task[Any] = asyncio.create_task(_invoke_tool())
+
+    await asyncio.wait_for(first_handler_started.wait(), timeout=1.0)
+    await asyncio.wait_for(second_event_yielded.wait(), timeout=1.0)
+    assert invoke_task.done() is False
+
+    allow_handler_to_continue.set()
+    await asyncio.wait_for(second_event_handled.wait(), timeout=1.0)
+    output = await asyncio.wait_for(invoke_task, timeout=1.0)
+
+    assert output == "ok"
+
+
 @pytest.mark.asyncio
 async def test_agent_as_tool_streaming_handler_exception_does_not_fail_call(
     monkeypatch: pytest.MonkeyPatch,