From 44ccfb2a12dd4b555be290175dfda4adccc55872 Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Wed, 5 Aug 2026 14:19:42 +0530 Subject: [PATCH] fix: stop waiting on Queue.join() when the consumer can die Three streaming dispatchers hand user events to a background task and then wait for asyncio.Queue.join() before returning. Each dispatcher catches only Exception, so a handler raising a BaseException -- CancelledError being the realistic one, from a timeout or cancellation inside the callback -- kills the consumer. The outstanding task_done() calls can then never arrive and join() blocks forever, hanging the run with no error and no timeout. In every case join() is also redundant. The sentinel is queued last and the consumer only returns after consuming it, so awaiting the consumer task already implies every event was handled, and it surfaces a dead consumer instead of waiting on it. - Agent.as_tool(on_stream=...) hung the tool invocation. - Codex tool _consume_events() hung and left its tracing spans unfinished, because the span cleanup runs after the join. - The sandbox memory manager flush hung before its phase-two consolidation. Same failure shape as the STT listener fix in #4170. --- src/agents/agent.py | 8 +- .../experimental/codex/codex_tool.py | 7 +- src/agents/sandbox/memory/manager.py | 6 +- .../experiemental/codex/test_codex_tool.py | 66 +++++++++ tests/test_agent_as_tool.py | 137 ++++++++++++++++++ 5 files changed, 221 insertions(+), 3 deletions(-) diff --git a/src/agents/agent.py b/src/agents/agent.py index 1d42624f2b..572a7f4e97 100644 --- a/src/agents/agent.py +++ b/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: diff --git a/src/agents/extensions/experimental/codex/codex_tool.py b/src/agents/extensions/experimental/codex/codex_tool.py index 7138286dfe..fc2bca0571 100644 --- a/src/agents/extensions/experimental/codex/codex_tool.py +++ b/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. diff --git a/src/agents/sandbox/memory/manager.py b/src/agents/sandbox/memory/manager.py index 9919d8035b..b8a5f71825 100644 --- a/src/agents/sandbox/memory/manager.py +++ b/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 diff --git a/tests/extensions/experiemental/codex/test_codex_tool.py b/tests/extensions/experiemental/codex/test_codex_tool.py index 36b6a2822a..cdd4f0bdb8 100644 --- a/tests/extensions/experiemental/codex/test_codex_tool.py +++ b/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 diff --git a/tests/test_agent_as_tool.py b/tests/test_agent_as_tool.py index ec2c4bbc20..a3a1d89dd0 100644 --- a/tests/test_agent_as_tool.py +++ b/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"