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"