diff --git a/src/agents/agent.py b/src/agents/agent.py index 1d42624f2b..7768372cee 100644 --- a/src/agents/agent.py +++ b/src/agents/agent.py @@ -58,7 +58,7 @@ ) from .tool_context import ToolContext from .util import _transforms -from .util._asyncio_tasks import gather_with_cancel +from .util._asyncio_tasks import gather_with_cancel, run_producer_consumer from .util._types import MaybeAwaitable if TYPE_CHECKING: @@ -912,10 +912,7 @@ async def dispatch_stream_events() -> None: if is_sentinel: break - dispatch_task = asyncio.create_task(dispatch_stream_events()) - stream_iteration_cancelled = False - - try: + async def enqueue_stream_events() -> None: from .stream_events import AgentUpdatedStreamEvent current_agent = run_result_streaming.current_agent @@ -930,20 +927,10 @@ async def dispatch_stream_events() -> None: "tool_call": context.tool_call, } await event_queue.put(payload) - except asyncio.CancelledError: - stream_iteration_cancelled = True - raise - finally: - if stream_iteration_cancelled: - dispatch_task.cancel() - try: - await dispatch_task - except asyncio.CancelledError: - pass - else: + finally: await event_queue.put(None) - await event_queue.join() - await dispatch_task + + await run_producer_consumer(enqueue_stream_events(), dispatch_stream_events()) run_result = run_result_streaming else: run_result = await Runner.run( diff --git a/src/agents/extensions/experimental/codex/codex_tool.py b/src/agents/extensions/experimental/codex/codex_tool.py index 7138286dfe..62c9eea334 100644 --- a/src/agents/extensions/experimental/codex/codex_tool.py +++ b/src/agents/extensions/experimental/codex/codex_tool.py @@ -31,6 +31,7 @@ from agents.tool_context import ToolContext from agents.tracing import SpanError, custom_span from agents.usage import Usage as AgentsUsage, _make_input_tokens_details +from agents.util._asyncio_tasks import run_producer_consumer from agents.util._types import MaybeAwaitable from .codex import Codex @@ -1047,78 +1048,81 @@ async def _consume_events( resolved_thread_id_holder["thread_id"] = resolved_thread_id event_queue: asyncio.Queue[CodexToolStreamEvent | None] | None = None - dispatch_task: asyncio.Task[None] | None = None - if on_stream is not None: # Buffer events so user callbacks cannot block the Codex stream loop. event_queue = asyncio.Queue() - async def _run_handler(payload: CodexToolStreamEvent) -> None: - # Dispatch user callbacks asynchronously to avoid blocking the stream. + async def _run_handler(payload: CodexToolStreamEvent) -> None: + # Dispatch user callbacks asynchronously to avoid blocking the stream. + assert on_stream is not None + try: + maybe_result = on_stream(payload) + if inspect.isawaitable(maybe_result): + await maybe_result + except Exception as exc: + log_model_and_tool_action_error( + logger, + "Error while handling Codex on_stream event", + exc, + ) + + async def _dispatch() -> None: + assert event_queue is not None + while True: + payload = await event_queue.get() + is_sentinel = payload is None try: - maybe_result = on_stream(payload) - if inspect.isawaitable(maybe_result): - await maybe_result - except Exception as exc: - log_model_and_tool_action_error( - logger, - "Error while handling Codex on_stream event", - exc, - ) + if payload is not None: + await _run_handler(payload) + finally: + event_queue.task_done() + if is_sentinel: + break - async def _dispatch() -> None: - assert event_queue is not None - while True: - payload = await event_queue.get() - is_sentinel = payload is None - try: - if payload is not None: - await _run_handler(payload) - finally: - event_queue.task_done() - if is_sentinel: - break + async def _process_events() -> None: + nonlocal final_response, resolved_thread_id, usage - dispatch_task = asyncio.create_task(_dispatch()) + try: + async for raw_event in events: + event = coerce_thread_event(raw_event) + if event_queue is not None: + await event_queue.put( + CodexToolStreamEvent( + event=event, + thread=thread, + tool_call=ctx.tool_call, + ) + ) - try: - async for raw_event in events: - event = coerce_thread_event(raw_event) + if isinstance(event, ItemStartedEvent): + _handle_item_started(event.item, active_spans, span_data_max_chars) + elif isinstance(event, ItemUpdatedEvent): + _handle_item_updated(event.item, active_spans, span_data_max_chars) + elif isinstance(event, ItemCompletedEvent): + _handle_item_completed(event.item, active_spans, span_data_max_chars) + if is_agent_message_item(event.item): + final_response = event.item.text + elif isinstance(event, TurnCompletedEvent): + usage = event.usage + elif isinstance(event, ThreadStartedEvent): + resolved_thread_id = event.thread_id + if resolved_thread_id_holder is not None: + resolved_thread_id_holder["thread_id"] = resolved_thread_id + elif isinstance(event, TurnFailedEvent): + error = event.error.message + raise UserError(f"Codex turn failed{(': ' + error) if error else ''}") + elif isinstance(event, ThreadErrorEvent): + raise UserError(f"Codex stream error: {event.message}") + finally: if event_queue is not None: - await event_queue.put( - CodexToolStreamEvent( - event=event, - thread=thread, - tool_call=ctx.tool_call, - ) - ) + await event_queue.put(None) - if isinstance(event, ItemStartedEvent): - _handle_item_started(event.item, active_spans, span_data_max_chars) - elif isinstance(event, ItemUpdatedEvent): - _handle_item_updated(event.item, active_spans, span_data_max_chars) - elif isinstance(event, ItemCompletedEvent): - _handle_item_completed(event.item, active_spans, span_data_max_chars) - if is_agent_message_item(event.item): - final_response = event.item.text - elif isinstance(event, TurnCompletedEvent): - usage = event.usage - elif isinstance(event, ThreadStartedEvent): - resolved_thread_id = event.thread_id - if resolved_thread_id_holder is not None: - resolved_thread_id_holder["thread_id"] = resolved_thread_id - elif isinstance(event, TurnFailedEvent): - error = event.error.message - raise UserError(f"Codex turn failed{(': ' + error) if error else ''}") - elif isinstance(event, ThreadErrorEvent): - raise UserError(f"Codex stream error: {event.message}") + try: + if on_stream is None: + await _process_events() + else: + await run_producer_consumer(_process_events(), _dispatch()) finally: - if event_queue is not None: - await event_queue.put(None) - await event_queue.join() - if dispatch_task is not None: - await dispatch_task - # Ensure any open spans are closed even on failure. for span in active_spans.values(): span.finish() diff --git a/src/agents/sandbox/memory/manager.py b/src/agents/sandbox/memory/manager.py index 9919d8035b..8c16f3534a 100644 --- a/src/agents/sandbox/memory/manager.py +++ b/src/agents/sandbox/memory/manager.py @@ -131,11 +131,15 @@ 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: self._queue.put_nowait(_STOP) - await self._worker_task - self._worker_task = None + worker_task = self._worker_task + try: + # The stop marker follows every rollout, so worker completion implies + # that all preceding rollout files were processed. + await worker_task + finally: + self._worker_task = None await self._run_phase_two() finally: _unregister_memory_generation_manager(session=self._session, manager=self) diff --git a/src/agents/util/_asyncio_tasks.py b/src/agents/util/_asyncio_tasks.py index 90c3146f18..2974af39f7 100644 --- a/src/agents/util/_asyncio_tasks.py +++ b/src/agents/util/_asyncio_tasks.py @@ -11,6 +11,8 @@ T4 = TypeVar("T4") T5 = TypeVar("T5") T6 = TypeVar("T6") +TProducer = TypeVar("TProducer") +TConsumer = TypeVar("TConsumer") def _consume_future_exception(future: asyncio.Future[Any]) -> None: @@ -110,3 +112,41 @@ async def gather_with_cancel( task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise + + +async def run_producer_consumer( + producer: Awaitable[TProducer], + consumer: Awaitable[TConsumer], + /, +) -> tuple[TProducer, TConsumer]: + """Run a producer and consumer with asymmetric failure handling. + + The producer must signal completion to the consumer in a ``finally`` block. A producer + failure waits for the consumer to drain before propagating, while a consumer failure or + parent cancellation cancels and drains the sibling task. + """ + producer_task = asyncio.ensure_future(producer) + consumer_task = asyncio.ensure_future(consumer) + tasks = (producer_task, consumer_task) + + try: + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if consumer_task in done: + consumer_result = consumer_task.result() + producer_result = await producer_task + return producer_result, consumer_result + + try: + producer_result = producer_task.result() + except BaseException: + await consumer_task + raise + + consumer_result = await consumer_task + return producer_result, consumer_result + except BaseException: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise diff --git a/tests/extensions/experiemental/codex/test_codex_tool.py b/tests/extensions/experiemental/codex/test_codex_tool.py index 36b6a2822a..1fd67e933d 100644 --- a/tests/extensions/experiemental/codex/test_codex_tool.py +++ b/tests/extensions/experiemental/codex/test_codex_tool.py @@ -2108,3 +2108,191 @@ 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 _FatalCodexStreamHandlerError(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "handler_error", + [_FatalCodexStreamHandlerError("fatal"), asyncio.CancelledError()], + ids=["base_exception", "cancelled_error"], +) +async def test_codex_tool_streaming_propagates_base_exception_and_finishes_spans( + monkeypatch: pytest.MonkeyPatch, + handler_error: BaseException, +) -> None: + class RecordingSpan: + def __init__(self) -> None: + self.started = False + self.finished = False + + def start(self) -> None: + self.started = True + + def finish(self) -> None: + self.finished = True + + span = RecordingSpan() + monkeypatch.setattr(codex_tool_module, "custom_span", lambda **_kwargs: span) + source_cancelled = asyncio.Event() + + async def event_stream(): + yield { + "type": "item.started", + "item": { + "id": "cmd-1", + "type": "command_execution", + "command": "pwd", + "status": "in_progress", + }, + } + try: + await asyncio.Event().wait() + finally: + source_cancelled.set() + + def on_stream(payload: CodexToolStreamEvent) -> None: + del payload + raise handler_error + + context = ToolContext( + context=None, + tool_name="codex", + tool_call_id="call-1", + tool_arguments="{}", + ) + + with pytest.raises(type(handler_error)): + await asyncio.wait_for( + codex_tool_module._consume_events( + event_stream(), + {"inputs": [{"type": "text", "text": "hello"}]}, + context, + SimpleNamespace(id="thread-1"), + on_stream, + 64, + ), + timeout=1.0, + ) + + assert span.started + assert span.finished + assert source_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_codex_tool_streaming_parent_cancellation_stops_dispatcher() -> None: + handler_started = asyncio.Event() + handler_cancelled = asyncio.Event() + source_cancelled = asyncio.Event() + + async def event_stream(): + yield { + "type": "turn.completed", + "usage": {"input_tokens": 1, "cached_input_tokens": 0, "output_tokens": 1}, + } + try: + await asyncio.Event().wait() + finally: + source_cancelled.set() + + async def on_stream(payload: CodexToolStreamEvent) -> None: + del payload + handler_started.set() + try: + await asyncio.Event().wait() + finally: + handler_cancelled.set() + + context = ToolContext( + context=None, + tool_name="codex", + tool_call_id="call-1", + tool_arguments="{}", + ) + invoke_task = asyncio.create_task( + codex_tool_module._consume_events( + event_stream(), + {"inputs": [{"type": "text", "text": "hello"}]}, + context, + SimpleNamespace(id="thread-1"), + on_stream, + 64, + ) + ) + + await asyncio.wait_for(handler_started.wait(), timeout=1.0) + invoke_task.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(invoke_task, timeout=1.0) + finally: + if not invoke_task.done(): + invoke_task.cancel() + await asyncio.gather(invoke_task, return_exceptions=True) + + assert source_cancelled.is_set() + assert handler_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_codex_tool_streaming_drains_events_before_stream_error() -> None: + handler_started = asyncio.Event() + terminal_event_emitted = asyncio.Event() + allow_handler_to_finish = asyncio.Event() + handler_cancelled = asyncio.Event() + handled_event_types: list[str] = [] + + async def event_stream(): + yield {"type": "turn.started"} + await handler_started.wait() + terminal_event_emitted.set() + yield {"type": "turn.failed", "error": {"message": "boom"}} + + async def on_stream(payload: CodexToolStreamEvent) -> None: + if not handled_event_types: + handler_started.set() + try: + await allow_handler_to_finish.wait() + except asyncio.CancelledError: + handler_cancelled.set() + raise + handled_event_types.append(payload.event.type) + + context = ToolContext( + context=None, + tool_name="codex", + tool_call_id="call-1", + tool_arguments="{}", + ) + invoke_task = asyncio.create_task( + 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(terminal_event_emitted.wait(), timeout=1.0) + await asyncio.sleep(0) + + assert not invoke_task.done() + assert not handler_cancelled.is_set() + + allow_handler_to_finish.set() + with pytest.raises(UserError, match="Codex turn failed: boom"): + await asyncio.wait_for(invoke_task, timeout=1.0) + finally: + if not invoke_task.done(): + invoke_task.cancel() + await asyncio.gather(invoke_task, return_exceptions=True) + + assert handled_event_types == ["turn.started", "turn.failed"] diff --git a/tests/sandbox/test_memory.py b/tests/sandbox/test_memory.py index 61625e77f4..1a8ed9a560 100644 --- a/tests/sandbox/test_memory.py +++ b/tests/sandbox/test_memory.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import io import json import logging @@ -1497,6 +1498,112 @@ async def test_sandbox_memory_unregisters_manager_on_session_close() -> None: await client.delete(session) +class _FatalMemoryWorkerError(BaseException): + pass + + +@pytest.mark.parametrize( + "worker_error", + [_FatalMemoryWorkerError("fatal"), asyncio.CancelledError()], + ids=["base_exception", "cancelled_error"], +) +@pytest.mark.asyncio +async def test_sandbox_memory_flush_propagates_worker_base_exception_without_hanging( + monkeypatch: pytest.MonkeyPatch, + worker_error: BaseException, +) -> None: + client = UnixLocalSandboxClient() + session = await client.create(manifest=Manifest()) + memory = _memory_config() + manager = get_or_create_memory_generation_manager(session=session, memory=memory) + + async def fail_processing(_rollout_file_name: str) -> None: + raise worker_error + + monkeypatch.setattr(manager, "_process_rollout_file", fail_processing) + + try: + await manager.enqueue_rollout_payload( + { + "updated_at": "2026-08-05T00:00:00+00:00", + "input": [], + "generated_items": [], + "terminal_metadata": { + "terminal_state": "completed", + "has_final_output": False, + }, + }, + rollout_id="fatal-worker", + ) + + with pytest.raises(type(worker_error)): + await asyncio.wait_for(manager.flush(), timeout=1.0) + + assert manager._worker_task is None + assert memory_manager_module._MEMORY_GENERATION_MANAGERS.get(session) is None + finally: + await client.delete(session) + + +@pytest.mark.asyncio +async def test_sandbox_memory_flush_parent_cancellation_stops_worker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = UnixLocalSandboxClient() + session = await client.create(manifest=Manifest()) + memory = _memory_config() + manager = get_or_create_memory_generation_manager(session=session, memory=memory) + worker_started = asyncio.Event() + worker_cancelled = asyncio.Event() + phase_two_called = False + + async def block_processing(_rollout_file_name: str) -> None: + worker_started.set() + try: + await asyncio.Event().wait() + finally: + worker_cancelled.set() + + async def record_phase_two() -> None: + nonlocal phase_two_called + phase_two_called = True + + monkeypatch.setattr(manager, "_process_rollout_file", block_processing) + monkeypatch.setattr(manager, "_run_phase_two", record_phase_two) + + try: + await manager.enqueue_rollout_payload( + { + "updated_at": "2026-08-05T00:00:00+00:00", + "input": [], + "generated_items": [], + "terminal_metadata": { + "terminal_state": "completed", + "has_final_output": False, + }, + }, + rollout_id="cancelled-flush", + ) + flush_task = asyncio.create_task(manager.flush()) + await asyncio.wait_for(worker_started.wait(), timeout=1.0) + flush_task.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(flush_task, timeout=1.0) + finally: + if not flush_task.done(): + flush_task.cancel() + await asyncio.gather(flush_task, return_exceptions=True) + + assert worker_cancelled.is_set() + assert manager._worker_task is None + assert memory_manager_module._MEMORY_GENERATION_MANAGERS.get(session) is None + assert not phase_two_called + finally: + await client.delete(session) + + @pytest.mark.parametrize("streamed", [False, True], ids=["non_streamed", "streamed"]) @pytest.mark.parametrize( ("model_redacted", "tool_redacted"), diff --git a/tests/test_agent_as_tool.py b/tests/test_agent_as_tool.py index ebe53f3315..bb52d5743c 100644 --- a/tests/test_agent_as_tool.py +++ b/tests/test_agent_as_tool.py @@ -3046,3 +3046,158 @@ 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 _FatalAgentToolStreamHandlerError(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "handler_error", + [_FatalAgentToolStreamHandlerError("fatal"), asyncio.CancelledError()], + ids=["base_exception", "cancelled_error"], +) +async def test_agent_as_tool_streaming_propagates_base_exception_without_hanging( + monkeypatch: pytest.MonkeyPatch, + handler_error: BaseException, +) -> None: + agent = Agent(name="streamer") + source_cancelled = asyncio.Event() + stream_event = RawResponsesStreamEvent(data=cast(Any, {"type": "response_started"})) + + class DummyStreamingResult: + def __init__(self) -> None: + self.final_output = "streamed" + self.current_agent = agent + + async def stream_events(self): + yield stream_event + try: + await asyncio.Event().wait() + finally: + source_cancelled.set() + + monkeypatch.setattr( + Runner, + "run_streamed", + classmethod(lambda *args, **kwargs: DummyStreamingResult()), + ) + + async def on_stream(payload: AgentToolStreamEvent) -> None: + del payload + raise handler_error + + 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, + ) + + with pytest.raises(type(handler_error)): + await asyncio.wait_for( + tool.on_invoke_tool(tool_context, '{"input": "go"}'), + timeout=1.0, + ) + + assert source_cancelled.is_set() + + +class _NestedAgentStreamError(Exception): + pass + + +@pytest.mark.asyncio +async def test_agent_as_tool_streaming_drains_emitted_events_before_stream_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = Agent(name="streamer") + stream_event = RawResponsesStreamEvent(data=cast(Any, {"type": "response_started"})) + handler_started = asyncio.Event() + producer_failed = asyncio.Event() + allow_handler_to_finish = asyncio.Event() + handler_cancelled = asyncio.Event() + handled_events: list[RawResponsesStreamEvent] = [] + + class DummyStreamingResult: + def __init__(self) -> None: + self.final_output = "streamed" + self.current_agent = agent + + async def stream_events(self): + yield stream_event + await handler_started.wait() + producer_failed.set() + raise _NestedAgentStreamError("nested stream failed") + + monkeypatch.setattr( + Runner, + "run_streamed", + classmethod(lambda *args, **kwargs: DummyStreamingResult()), + ) + + async def on_stream(payload: AgentToolStreamEvent) -> None: + handler_started.set() + try: + await allow_handler_to_finish.wait() + except asyncio.CancelledError: + handler_cancelled.set() + raise + handled_events.append(cast(RawResponsesStreamEvent, payload["event"])) + + tool_call = ResponseFunctionToolCall( + id="call_stream_error", + arguments='{"input": "go"}', + call_id="call-stream-error", + name="stream_tool", + type="function_call", + ) + tool = agent.as_tool( + tool_name="stream_tool", + tool_description="Streams events", + on_stream=on_stream, + failure_error_function=None, + ) + 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, + ) + + async def invoke() -> Any: + return await tool.on_invoke_tool(tool_context, '{"input": "go"}') + + invoke_task = asyncio.create_task(invoke()) + + try: + await asyncio.wait_for(producer_failed.wait(), timeout=1.0) + await asyncio.sleep(0) + + assert not invoke_task.done() + assert not handler_cancelled.is_set() + + allow_handler_to_finish.set() + with pytest.raises(_NestedAgentStreamError, match="nested stream failed"): + await asyncio.wait_for(invoke_task, timeout=1.0) + finally: + if not invoke_task.done(): + invoke_task.cancel() + await asyncio.gather(invoke_task, return_exceptions=True) + + assert handled_events == [stream_event] diff --git a/tests/test_asyncio_tasks.py b/tests/test_asyncio_tasks.py index 0315a63705..5a77291a47 100644 --- a/tests/test_asyncio_tasks.py +++ b/tests/test_asyncio_tasks.py @@ -4,7 +4,7 @@ import pytest -from agents.util._asyncio_tasks import gather_with_cancel +from agents.util._asyncio_tasks import gather_with_cancel, run_producer_consumer @pytest.mark.asyncio @@ -75,3 +75,57 @@ async def child() -> None: assert not child_failure_reported.is_set() assert loop_errors == [] + + +@pytest.mark.asyncio +async def test_run_producer_consumer_drains_consumer_before_producer_failure() -> None: + class ProducerError(Exception): + pass + + item_ready = asyncio.Event() + allow_consumer_to_finish = asyncio.Event() + consumer_finished = asyncio.Event() + + async def producer() -> None: + item_ready.set() + raise ProducerError("producer failed") + + async def consumer() -> None: + await item_ready.wait() + await allow_consumer_to_finish.wait() + consumer_finished.set() + + task = asyncio.create_task(run_producer_consumer(producer(), consumer())) + await item_ready.wait() + await asyncio.sleep(0) + + assert not task.done() + allow_consumer_to_finish.set() + + with pytest.raises(ProducerError, match="producer failed"): + await task + assert consumer_finished.is_set() + + +@pytest.mark.asyncio +async def test_run_producer_consumer_cancels_producer_after_consumer_failure() -> None: + class ConsumerError(BaseException): + pass + + producer_started = asyncio.Event() + producer_cancelled = asyncio.Event() + + async def producer() -> None: + producer_started.set() + try: + await asyncio.Event().wait() + finally: + producer_cancelled.set() + + async def consumer() -> None: + await producer_started.wait() + raise ConsumerError("consumer failed") + + with pytest.raises(ConsumerError, match="consumer failed"): + await run_producer_consumer(producer(), consumer()) + assert producer_cancelled.is_set()