From cb7fa2b91cd67610ea28bc8fa9e882f4f304e7e9 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Tue, 1 Sep 2026 17:34:33 +0800 Subject: [PATCH 1/3] fix(codex): reap subprocesses on early stream close --- .../extensions/experimental/codex/exec.py | 5 +- .../codex/test_codex_exec_thread.py | 53 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index 0001b5700f..2b8a7c1c32 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -207,9 +207,12 @@ async def _read_stdout_line() -> bytes: finally: if cancel_task is not None and not cancel_task.done(): cancel_task.cancel() - await stderr_task if process.returncode is None: process.kill() + try: + await stderr_task + finally: + await process.wait() def _build_env(self, args: CodexExecArgs) -> dict[str, str]: # Respect env overrides when provided; otherwise copy from os.environ. diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index 51c635205e..c8c7ecb548 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -422,6 +422,59 @@ async def fake_create_subprocess_exec(*_args: Any, **kwargs: Any) -> StreamReade assert captured["kwargs"]["limit"] == exec_module._DEFAULT_SUBPROCESS_STREAM_LIMIT_BYTES +@pytest.mark.asyncio +async def test_codex_exec_run_closes_live_process_before_draining_stderr( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stderr_eof = asyncio.Event() + + class BlockingStderr: + async def read(self, _size: int) -> bytes: + await stderr_eof.wait() + return b"" + + class LiveProcess: + def __init__(self) -> None: + self.stdin = FakeStdin() + self.stdout = FakeStdout(["line\n"]) + self.stderr = BlockingStderr() + self.returncode: int | None = None + self.killed = False + self.wait_called = False + + async def wait(self) -> None: + self.wait_called = True + await stderr_eof.wait() + self.returncode = -9 + + def kill(self) -> None: + self.killed = True + stderr_eof.set() + + def terminate(self) -> None: + raise AssertionError("terminate() should not be used when the stream is closed") + + process = LiveProcess() + + async def fake_create_subprocess_exec(*_args: Any, **_kwargs: Any) -> LiveProcess: + return process + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + + exec_client = exec_module.CodexExec(executable_path="/bin/codex") + stream = exec_client.run(exec_module.CodexExecArgs(input="hello")) + + assert await anext(stream) == "line" + close_task = asyncio.create_task(stream.aclose()) + await asyncio.sleep(0) + killed_before_stderr_eof = process.killed + stderr_eof.set() + await close_task + + assert killed_before_stderr_eof is True + assert process.wait_called is True + + @pytest.mark.asyncio @pytest.mark.parametrize( ("enabled", "expected_config"), From 141b020a98c89502e5a2da1083c3b03a9a534cce Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Thu, 10 Sep 2026 10:51:57 +0800 Subject: [PATCH 2/3] fix(codex): drain stdout during stream teardown --- .../extensions/experimental/codex/exec.py | 40 ++++---- .../codex/test_codex_exec_thread.py | 92 ++++++++++++++++++- 2 files changed, 111 insertions(+), 21 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index 2b8a7c1c32..6ebf4ab34f 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -168,22 +168,26 @@ async def _read_stdout_line() -> bytes: return await stdout.readline() read_task: asyncio.Task[bytes] = asyncio.create_task(stdout.readline()) - done, _ = await asyncio.wait( - {read_task}, timeout=args.idle_timeout_seconds, return_when=asyncio.FIRST_COMPLETED - ) - if read_task in done: - return read_task.result() - - if args.signal is not None: - args.signal.set() - if process.returncode is None: - process.terminate() + try: + done, _ = await asyncio.wait( + {read_task}, + timeout=args.idle_timeout_seconds, + return_when=asyncio.FIRST_COMPLETED, + ) + if read_task in done: + return read_task.result() - read_task.cancel() - with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError): - await asyncio.wait_for(read_task, timeout=1) + if args.signal is not None: + args.signal.set() + if process.returncode is None: + process.terminate() - raise RuntimeError(f"Codex stream idle for {args.idle_timeout_seconds} seconds.") + raise RuntimeError(f"Codex stream idle for {args.idle_timeout_seconds} seconds.") + finally: + if not read_task.done(): + read_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await read_task try: while True: @@ -210,9 +214,13 @@ async def _read_stdout_line() -> bytes: if process.returncode is None: process.kill() try: - await stderr_task + while await stdout.read(1024): + pass finally: - await process.wait() + try: + await stderr_task + finally: + await process.wait() def _build_env(self, args: CodexExecArgs) -> dict[str, str]: # Respect env overrides when provided; otherwise copy from os.environ. diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index c8c7ecb548..4630af7549 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -55,6 +55,11 @@ async def readline(self) -> bytes: return b"" return self._lines.pop(0) + async def read(self, _size: int) -> bytes: + if not self._lines: + return b"" + return self._lines.pop(0) + class FakeStderr: def __init__(self, chunks: list[bytes]) -> None: @@ -423,20 +428,38 @@ async def fake_create_subprocess_exec(*_args: Any, **kwargs: Any) -> StreamReade @pytest.mark.asyncio -async def test_codex_exec_run_closes_live_process_before_draining_stderr( +async def test_codex_exec_run_drains_stdout_before_waiting_for_live_process( monkeypatch: pytest.MonkeyPatch, ) -> None: stderr_eof = asyncio.Event() + stdout_drained = asyncio.Event() class BlockingStderr: async def read(self, _size: int) -> bytes: await stderr_eof.wait() return b"" + class BackpressuredStdout: + def __init__(self) -> None: + self._line_read = False + self._unread_chunks = [b"discarded output", b""] + + async def readline(self) -> bytes: + if self._line_read: + return b"" + self._line_read = True + return b"line\n" + + async def read(self, _size: int) -> bytes: + chunk = self._unread_chunks.pop(0) + if not chunk: + stdout_drained.set() + return chunk + class LiveProcess: def __init__(self) -> None: self.stdin = FakeStdin() - self.stdout = FakeStdout(["line\n"]) + self.stdout = BackpressuredStdout() self.stderr = BlockingStderr() self.returncode: int | None = None self.killed = False @@ -444,12 +467,11 @@ def __init__(self) -> None: async def wait(self) -> None: self.wait_called = True - await stderr_eof.wait() + await stdout_drained.wait() self.returncode = -9 def kill(self) -> None: self.killed = True - stderr_eof.set() def terminate(self) -> None: raise AssertionError("terminate() should not be used when the stream is closed") @@ -469,10 +491,70 @@ async def fake_create_subprocess_exec(*_args: Any, **_kwargs: Any) -> LiveProces await asyncio.sleep(0) killed_before_stderr_eof = process.killed stderr_eof.set() - await close_task + await asyncio.wait_for(close_task, timeout=1) assert killed_before_stderr_eof is True + assert stdout_drained.is_set() assert process.wait_called is True + assert process.returncode == -9 + + +@pytest.mark.asyncio +async def test_codex_exec_run_cancels_pending_timeout_read_before_draining_stdout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + read_started = asyncio.Event() + readline_cancelled = asyncio.Event() + + class BlockingStdout: + async def readline(self) -> bytes: + read_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + readline_cancelled.set() + raise + + async def read(self, _size: int) -> bytes: + assert readline_cancelled.is_set() + return b"" + + class LiveProcess: + def __init__(self) -> None: + self.stdin = FakeStdin() + self.stdout = BlockingStdout() + self.stderr = FakeStderr([]) + self.returncode: int | None = None + self.killed = False + + async def wait(self) -> None: + self.returncode = -9 + + def kill(self) -> None: + self.killed = True + + def terminate(self) -> None: + raise AssertionError("terminate() should not be used before the idle timeout") + + process = LiveProcess() + + async def fake_create_subprocess_exec(*_args: Any, **_kwargs: Any) -> LiveProcess: + return process + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + + exec_client = exec_module.CodexExec(executable_path="/bin/codex") + stream = exec_client.run(exec_module.CodexExecArgs(input="hello", idle_timeout_seconds=60)) + read_task = asyncio.create_task(anext(stream)) + await read_started.wait() + read_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await read_task + + assert readline_cancelled.is_set() + assert process.killed is True + assert process.returncode == -9 @pytest.mark.asyncio From fe03e94a87f2a63a94d954703ba3929a64f72f36 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Thu, 10 Sep 2026 10:54:11 +0800 Subject: [PATCH 3/3] test(codex): make blocked stdout read explicit --- tests/extensions/experiemental/codex/test_codex_exec_thread.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index 4630af7549..dc305c2ded 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -514,6 +514,7 @@ async def readline(self) -> bytes: except asyncio.CancelledError: readline_cancelled.set() raise + raise AssertionError("readline unexpectedly completed") async def read(self, _size: int) -> bytes: assert readline_cancelled.is_set()