fix: stop waiting on Queue.join() when the event consumer can die

src/agents/agent.py

@@ -942,7 +942,13 @@ async def dispatch_stream_events() -> None:
                                 pass
                         else:
                             await event_queue.put(None)
-                            await event_queue.join()
+                            # Wait on the dispatcher itself rather than on the queue draining.
+                            # The sentinel is the last item queued and the dispatcher only
+                            # returns after consuming it, so awaiting the task already implies
+                            # every event was handled. Waiting on `join()` instead would hang
+                            # forever when a handler raised a BaseException such as
+                            # CancelledError, because the dispatcher is then gone and the
+                            # outstanding `task_done()` calls can never arrive.
                             await dispatch_task
                     run_result = run_result_streaming
                 else:

src/agents/extensions/experimental/codex/codex_tool.py

@@ -1115,8 +1115,13 @@ async def _dispatch() -> None:
     finally:
         if event_queue is not None:
             await event_queue.put(None)
-            await event_queue.join()
         if dispatch_task is not None:
+            # Wait on the dispatcher itself rather than on the queue draining. The sentinel is
+            # the last item queued and the dispatcher only returns after consuming it, so
+            # awaiting the task already implies every event was handled. Waiting on `join()`
+            # instead would hang forever when a handler raised a BaseException such as
+            # CancelledError, because the dispatcher is then gone and the outstanding
+            # `task_done()` calls can never arrive -- stranding the open spans below too.
             await dispatch_task
 
         # Ensure any open spans are closed even on failure.

src/agents/sandbox/memory/manager.py

@@ -131,8 +131,12 @@ async def flush(self) -> None:
                 self._ensure_worker()
                 for rollout_file in rollout_files:
                     self._queue.put_nowait(rollout_file)
-                await self._queue.join()
                 if self._worker_task is not None:
+                    # Wait on the worker rather than on the queue draining. The worker consumes
+                    # FIFO and only returns on _STOP, which is queued behind every rollout, so
+                    # awaiting it already implies each one was processed. Waiting on `join()`
+                    # instead would hang forever if the worker died on a BaseException, because
+                    # the outstanding `task_done()` calls can never arrive.
                     self._queue.put_nowait(_STOP)
                     await self._worker_task
                     self._worker_task = None

tests/extensions/experiemental/codex/test_codex_tool.py

@@ -2108,3 +2108,69 @@ async def test_codex_tool_argument_errors_respect_tool_data_redaction(
     else:
         assert _CODEX_TOOL_ARGUMENT_SECRET in str(error)
         assert isinstance(error.__cause__, cause_type)
+
+
+class _FatalCodexHandlerError(BaseException):
+    """A BaseException, so `_run_handler`'s `except Exception` does not catch it."""
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "raised",
+    [_FatalCodexHandlerError("fatal"), asyncio.CancelledError()],
+    ids=["base_exception", "cancelled_error"],
+)
+async def test_codex_tool_consume_events_does_not_hang_on_base_exception(
+    raised: BaseException,
+) -> None:
+    """A handler raising a BaseException must not strand `_consume_events`.
+
+    The dispatcher only catches `Exception`, so a BaseException such as `CancelledError`
+    terminates it. Waiting on `event_queue.join()` then waited for `task_done()` calls that
+    could never arrive, deadlocking the tool and leaving the active tracing spans unfinished.
+    """
+    events = [
+        {
+            "type": "item.completed",
+            "item": {"id": "agent-1", "type": "agent_message", "text": "done"},
+        },
+        {
+            "type": "turn.completed",
+            "usage": {"input_tokens": 1, "cached_input_tokens": 0, "output_tokens": 1},
+        },
+    ]
+
+    async def event_stream():
+        for event in events:
+            yield event
+
+    def on_stream(payload: CodexToolStreamEvent) -> None:
+        del payload
+        raise raised
+
+    context = ToolContext(
+        context=None,
+        tool_name="codex",
+        tool_call_id="call-1",
+        tool_arguments="{}",
+    )
+
+    async def invoke():
+        with trace("codex-test"):
+            return await codex_tool_module._consume_events(
+                event_stream(),
+                {"inputs": [{"type": "text", "text": "hello"}]},
+                context,
+                SimpleNamespace(id="thread-1"),
+                on_stream,
+                64,
+            )
+
+    try:
+        await asyncio.wait_for(invoke(), timeout=5)
+    except asyncio.TimeoutError:
+        pytest.fail("_consume_events deadlocked after the on_stream handler raised")
+    except _FatalCodexHandlerError:
+        pass
+    except asyncio.CancelledError:
+        pass

tests/test_agent_as_tool.py

@@ -3047,3 +3047,140 @@ def test_replaced_agent_as_tool_preserves_agent_markers_for_build_agent_map() ->
     agent_map = _build_agent_map(parent_agent)
 
     assert agent_map["nested_agent"] is nested_agent
+
+
+class _FatalHandlerError(BaseException):
+    """A BaseException, so it is not caught by the on_stream handler's `except Exception`."""
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "raised",
+    [_FatalHandlerError("fatal"), asyncio.CancelledError()],
+    ids=["base_exception", "cancelled_error"],
+)
+async def test_agent_as_tool_streaming_does_not_hang_when_handler_raises_base_exception(
+    raised: BaseException,
+) -> None:
+    """A handler raising a BaseException must not strand the run.
+
+    The dispatcher only catches `Exception`, so a BaseException such as `CancelledError`
+    terminates it. Waiting on `event_queue.join()` then waited for `task_done()` calls that
+    could never arrive, deadlocking the tool invocation with no error and no timeout.
+    """
+    agent = Agent(
+        name="streamer",
+        model=FakeModel(
+            initial_output=[
+                ResponseOutputMessage(
+                    id="msg-fatal",
+                    role="assistant",
+                    status="completed",
+                    type="message",
+                    content=[
+                        ResponseOutputText(
+                            annotations=[],
+                            text="streamed",
+                            type="output_text",
+                            logprobs=[],
+                        )
+                    ],
+                )
+            ]
+        ),
+    )
+
+    async def on_stream(payload: AgentToolStreamEvent) -> None:
+        del payload
+        raise raised
+
+    tool_call = ResponseFunctionToolCall(
+        id="call_fatal",
+        arguments='{"input": "go"}',
+        call_id="call-fatal",
+        name="stream_tool",
+        type="function_call",
+    )
+    tool = agent.as_tool(
+        tool_name="stream_tool",
+        tool_description="Streams events",
+        on_stream=on_stream,
+    )
+    tool_context = ToolContext(
+        context=None,
+        tool_name="stream_tool",
+        tool_call_id=tool_call.call_id,
+        tool_arguments=tool_call.arguments,
+        tool_call=tool_call,
+    )
+
+    # The call must settle rather than hang; whether it returns or propagates is secondary.
+    try:
+        await asyncio.wait_for(
+            tool.on_invoke_tool(tool_context, '{"input": "go"}'),
+            timeout=5,
+        )
+    except asyncio.TimeoutError:
+        pytest.fail("on_invoke_tool deadlocked after the on_stream handler raised")
+    except _FatalHandlerError:
+        pass
+    except asyncio.CancelledError:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_agent_as_tool_streaming_reports_every_event_before_returning() -> None:
+    """Dropping `join()` must not let the tool return before handlers have run."""
+    agent = Agent(
+        name="streamer",
+        model=FakeModel(
+            initial_output=[
+                ResponseOutputMessage(
+                    id="msg-ordered",
+                    role="assistant",
+                    status="completed",
+                    type="message",
+                    content=[
+                        ResponseOutputText(
+                            annotations=[],
+                            text="ordered",
+                            type="output_text",
+                            logprobs=[],
+                        )
+                    ],
+                )
+            ]
+        ),
+    )
+
+    handled: list[str] = []
+
+    async def on_stream(payload: AgentToolStreamEvent) -> None:
+        # Yield control so a missing wait would let the tool return first.
+        await asyncio.sleep(0)
+        handled.append(payload["event"].type)
+
+    tool_call = ResponseFunctionToolCall(
+        id="call_ordered",
+        arguments='{"input": "go"}',
+        call_id="call-ordered",
+        name="stream_tool",
+        type="function_call",
+    )
+    tool = agent.as_tool(
+        tool_name="stream_tool",
+        tool_description="Streams events",
+        on_stream=on_stream,
+    )
+    tool_context = ToolContext(
+        context=None,
+        tool_name="stream_tool",
+        tool_call_id=tool_call.call_id,
+        tool_arguments=tool_call.arguments,
+        tool_call=tool_call,
+    )
+
+    output = await tool.on_invoke_tool(tool_context, '{"input": "go"}')
+
+    assert output == "ordered"
+    assert handled, "every streamed event must reach the handler before the tool returns"