From 8e7af5931a62075d4ed5294eea8617f8354d683e Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Thu, 23 Jul 2026 10:36:38 +0800 Subject: [PATCH 1/7] fix(codex): clean up subprocess on cancellation --- .../extensions/experimental/codex/exec.py | 104 ++++++++------- .../codex/test_codex_exec_thread.py | 120 +++++++++++++++++- 2 files changed, 177 insertions(+), 47 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index c83a0e98cd..9579b67792 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -20,6 +20,7 @@ _DEFAULT_SUBPROCESS_STREAM_LIMIT_BYTES = 8 * 1024 * 1024 _MIN_SUBPROCESS_STREAM_LIMIT_BYTES = 64 * 1024 _MAX_SUBPROCESS_STREAM_LIMIT_BYTES = 64 * 1024 * 1024 +_SUBPROCESS_DRAIN_CHUNK_BYTES = 64 * 1024 @dataclass(frozen=True) @@ -140,49 +141,58 @@ async def _drain_stderr() -> None: break stderr_chunks.append(chunk) - stderr_task = asyncio.create_task(_drain_stderr()) - - if process.stdin is None: - process.kill() - raise RuntimeError("Codex subprocess has no stdin") - - process.stdin.write(args.input.encode("utf-8")) - await process.stdin.drain() - process.stdin.close() - - if process.stdout is None: - process.kill() - raise RuntimeError("Codex subprocess has no stdout") - stdout = process.stdout + async def _drain_remaining_stdout() -> None: + if process.stdout is None: + return + while await process.stdout.read(_SUBPROCESS_DRAIN_CHUNK_BYTES): + pass + stderr_task = asyncio.create_task(_drain_stderr()) cancel_task: asyncio.Task[None] | None = None - if args.signal is not None: - # Mirror AbortSignal semantics by terminating the subprocess. - cancel_task = asyncio.create_task(_watch_signal(args.signal, process)) + try: + if process.stdin is None: + raise RuntimeError("Codex subprocess has no stdin") - async def _read_stdout_line() -> bytes: - if args.idle_timeout_seconds is None: - return await stdout.readline() + process.stdin.write(args.input.encode("utf-8")) + await process.stdin.drain() + process.stdin.close() - 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 process.stdout is None: + raise RuntimeError("Codex subprocess has no stdout") + stdout = process.stdout if args.signal is not None: - args.signal.set() - if process.returncode is None: - process.terminate() - - read_task.cancel() - with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError): - await asyncio.wait_for(read_task, timeout=1) - - raise RuntimeError(f"Codex stream idle for {args.idle_timeout_seconds} seconds.") + # Mirror AbortSignal semantics by terminating the subprocess. + cancel_task = asyncio.create_task(_watch_signal(args.signal, process)) + + async def _read_stdout_line() -> bytes: + if args.idle_timeout_seconds is None: + return await stdout.readline() + + read_task: asyncio.Task[bytes] = asyncio.create_task(stdout.readline()) + 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() + + 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." + ) + finally: + if not read_task.done(): + read_task.cancel() + with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError): + await asyncio.wait_for(read_task, timeout=1) - try: while True: line = await _read_stdout_line() if not line: @@ -190,11 +200,6 @@ async def _read_stdout_line() -> bytes: yield line.decode("utf-8").rstrip("\n") await process.wait() - if cancel_task is not None: - cancel_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await cancel_task - if process.returncode not in (0, None): await stderr_task stderr_text = b"".join(stderr_chunks).decode("utf-8") @@ -202,11 +207,20 @@ async def _read_stdout_line() -> bytes: f"Codex exec exited with code {process.returncode}: {stderr_text}" ) finally: - if cancel_task is not None and not cancel_task.done(): - cancel_task.cancel() - await stderr_task + if cancel_task is not None: + if not cancel_task.done(): + cancel_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await cancel_task if process.returncode is None: - process.kill() + with contextlib.suppress(ProcessLookupError): + process.kill() + await _drain_remaining_stdout() + await process.wait() + if not stderr_task.done(): + stderr_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stderr_task 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 c1012c51b3..d4fff72c72 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -5,6 +5,7 @@ import inspect import json import os +import sys from dataclasses import fields from pathlib import Path from typing import Any, cast @@ -55,6 +56,13 @@ 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"" + data = b"".join(self._lines) + self._lines = [data[size:]] if len(data) > size else [] + return data[:size] + class FakeStderr: def __init__(self, chunks: list[bytes]) -> None: @@ -442,9 +450,117 @@ async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: pass +@pytest.mark.asyncio +@pytest.mark.parametrize("idle_timeout_seconds", [None, 30.0]) +async def test_codex_exec_run_cancellation_reaps_process_and_stderr_task( + monkeypatch: pytest.MonkeyPatch, idle_timeout_seconds: float | None +) -> None: + stdout_started = asyncio.Event() + stderr_started = asyncio.Event() + stderr_cancelled = asyncio.Event() + + class BlockingStdout: + async def readline(self) -> bytes: + stdout_started.set() + await asyncio.Future[None]() + return b"" + + async def read(self, _size: int) -> bytes: + return b"" + + class BlockingStderr: + async def read(self, _size: int) -> bytes: + stderr_started.set() + try: + await asyncio.Future[None]() + except asyncio.CancelledError: + stderr_cancelled.set() + raise + return b"" + + process = FakeProcess(stdout_lines=[], returncode=None) + process.stdout = cast(Any, BlockingStdout()) + process.stderr = cast(Any, BlockingStderr()) + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + 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=idle_timeout_seconds) + ) + next_line = asyncio.create_task(anext(stream)) + await asyncio.wait_for(stdout_started.wait(), timeout=1) + await asyncio.wait_for(stderr_started.wait(), timeout=1) + + next_line.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(next_line, timeout=1) + + assert process.killed is True + assert process.returncode == 0 + assert stderr_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_codex_exec_run_close_drains_stdout_before_wait( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_create_subprocess_exec = asyncio.create_subprocess_exec + process: asyncio.subprocess.Process | None = None + child_code = ( + "import os\n" + "os.write(1, b'ready\\n')\n" + "chunk = b'x' * 65536\n" + "while True:\n" + " os.write(1, chunk)\n" + ) + + async def create_output_heavy_subprocess( + *_args: Any, **kwargs: Any + ) -> asyncio.subprocess.Process: + nonlocal process + process = await original_create_subprocess_exec(sys.executable, "-c", child_code, **kwargs) + return process + + monkeypatch.setattr( + exec_module.asyncio, "create_subprocess_exec", create_output_heavy_subprocess + ) + + exec_client = exec_module.CodexExec( + executable_path="unused", + subprocess_stream_limit_bytes=exec_module._MIN_SUBPROCESS_STREAM_LIMIT_BYTES, + ) + stream = exec_client.run(exec_module.CodexExecArgs(input="hello")) + + try: + assert await asyncio.wait_for(anext(stream), timeout=5) == "ready" + assert process is not None + assert process.stdout is not None + stdout = cast(Any, process.stdout) + for _ in range(500): + if stdout._paused: + break + await asyncio.sleep(0.01) + assert stdout._paused + + await asyncio.wait_for(stream.aclose(), timeout=5) + assert process.returncode is not None + finally: + if process is not None: + if process.returncode is None: + process.kill() + if process.stdout is not None: + while await process.stdout.read(exec_module._SUBPROCESS_DRAIN_CHUNK_BYTES): + pass + await process.wait() + + @pytest.mark.asyncio async def test_codex_exec_run_raises_without_stdin(monkeypatch: pytest.MonkeyPatch) -> None: - process = FakeProcess(stdout_lines=[], stdin_present=False) + process = FakeProcess(stdout_lines=[], returncode=None, stdin_present=False) async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: return process @@ -462,7 +578,7 @@ async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: @pytest.mark.asyncio async def test_codex_exec_run_raises_without_stdout(monkeypatch: pytest.MonkeyPatch) -> None: - process = FakeProcess(stdout_lines=[], stdout_present=False) + process = FakeProcess(stdout_lines=[], returncode=None, stdout_present=False) async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: return process From da45e04d40665a25ce8617c6fd445643503f5981 Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Thu, 23 Jul 2026 14:15:16 +0800 Subject: [PATCH 2/7] fix(codex): bound subprocess cleanup --- .../extensions/experimental/codex/exec.py | 26 ++++++- .../codex/test_codex_exec_thread.py | 67 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index 9579b67792..4383242108 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -21,6 +21,7 @@ _MIN_SUBPROCESS_STREAM_LIMIT_BYTES = 64 * 1024 _MAX_SUBPROCESS_STREAM_LIMIT_BYTES = 64 * 1024 * 1024 _SUBPROCESS_DRAIN_CHUNK_BYTES = 64 * 1024 +_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS = 1.0 @dataclass(frozen=True) @@ -147,6 +148,28 @@ async def _drain_remaining_stdout() -> None: while await process.stdout.read(_SUBPROCESS_DRAIN_CHUNK_BYTES): pass + async def _drain_stdout_and_wait() -> None: + await _drain_remaining_stdout() + await process.wait() + + async def _cleanup_process() -> None: + try: + await asyncio.wait_for( + _drain_stdout_and_wait(), + timeout=_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + # A descendant may outlive the direct child while holding a pipe open. + # Close the transport so cleanup does not wait indefinitely for EOF. + transport = getattr(process, "_transport", None) + if transport is not None: + transport.close() + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for( + process.wait(), + timeout=_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS, + ) + stderr_task = asyncio.create_task(_drain_stderr()) cancel_task: asyncio.Task[None] | None = None try: @@ -215,8 +238,7 @@ async def _read_stdout_line() -> bytes: if process.returncode is None: with contextlib.suppress(ProcessLookupError): process.kill() - await _drain_remaining_stdout() - await process.wait() + await _cleanup_process() if not stderr_task.done(): stderr_task.cancel() with contextlib.suppress(asyncio.CancelledError): diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index d4fff72c72..ade5826fa7 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -1,10 +1,12 @@ from __future__ import annotations import asyncio +import contextlib import importlib import inspect import json import os +import signal import sys from dataclasses import fields from pathlib import Path @@ -558,6 +560,71 @@ async def create_output_heavy_subprocess( await process.wait() +@pytest.mark.asyncio +async def test_codex_exec_run_close_is_bounded_when_descendant_holds_stdout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_create_subprocess_exec = asyncio.create_subprocess_exec + process: asyncio.subprocess.Process | None = None + helper_pid: int | None = None + helper_code = "import time; time.sleep(10)" + child_code = ( + "import os\n" + "import subprocess\n" + "import sys\n" + f"helper = subprocess.Popen([sys.executable, '-c', {helper_code!r}], " + "stdin=subprocess.DEVNULL, stdout=sys.stdout, stderr=subprocess.DEVNULL)\n" + "os.write(1, f'ready {helper.pid}\\n'.encode())\n" + "chunk = b'x' * 65536\n" + "while True:\n" + " os.write(1, chunk)\n" + ) + + async def create_output_heavy_subprocess( + *_args: Any, **kwargs: Any + ) -> asyncio.subprocess.Process: + nonlocal process + process = await original_create_subprocess_exec(sys.executable, "-c", child_code, **kwargs) + return process + + monkeypatch.setattr( + exec_module.asyncio, "create_subprocess_exec", create_output_heavy_subprocess + ) + monkeypatch.setattr(exec_module, "_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS", 0.05, raising=False) + + exec_client = exec_module.CodexExec( + executable_path="unused", + subprocess_stream_limit_bytes=exec_module._MIN_SUBPROCESS_STREAM_LIMIT_BYTES, + ) + stream = exec_client.run(exec_module.CodexExecArgs(input="hello")) + + try: + ready_line = await asyncio.wait_for(anext(stream), timeout=5) + helper_pid = int(ready_line.removeprefix("ready ")) + assert process is not None + assert process.stdout is not None + stdout = cast(Any, process.stdout) + for _ in range(500): + if stdout._paused: + break + await asyncio.sleep(0.01) + assert stdout._paused + + await asyncio.wait_for(stream.aclose(), timeout=1) + assert process.returncode is not None + finally: + if helper_pid is not None: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.kill(helper_pid, signal.SIGTERM) + if process is not None: + if process.returncode is None: + process.kill() + if process.stdout is not None: + while await process.stdout.read(exec_module._SUBPROCESS_DRAIN_CHUNK_BYTES): + pass + await process.wait() + + @pytest.mark.asyncio async def test_codex_exec_run_raises_without_stdin(monkeypatch: pytest.MonkeyPatch) -> None: process = FakeProcess(stdout_lines=[], returncode=None, stdin_present=False) From c0048076d56e8501ddb930f7b1578e7c123df7f5 Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Thu, 23 Jul 2026 16:22:51 +0800 Subject: [PATCH 3/7] fix(codex): terminate subprocess trees on cancellation --- .../extensions/experimental/codex/exec.py | 129 +++++++++-- .../codex/test_codex_exec_thread.py | 213 +++++++++++++++++- 2 files changed, 322 insertions(+), 20 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index 4383242108..deeb02c2c2 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -5,6 +5,8 @@ import os import platform import shutil +import signal +import subprocess import sys from collections.abc import AsyncGenerator from dataclasses import dataclass @@ -128,6 +130,8 @@ async def run(self, args: CodexExecArgs) -> AsyncGenerator[str, None]: # default 64 KiB readline limit. limit=self._subprocess_stream_limit_bytes, env=env, + # Give POSIX descendants a process-group boundary that this execution owns. + start_new_session=os.name != "nt", ) stderr_chunks: list[bytes] = [] @@ -204,8 +208,6 @@ async def _read_stdout_line() -> bytes: 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." @@ -230,19 +232,44 @@ async def _read_stdout_line() -> bytes: f"Codex exec exited with code {process.returncode}: {stderr_text}" ) finally: - if cancel_task is not None: - if not cancel_task.done(): - cancel_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await cancel_task - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - process.kill() - await _cleanup_process() - if not stderr_task.done(): - stderr_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await stderr_task + original_error = sys.exc_info()[1] + + async def _cleanup_resources() -> None: + cleanup_error: BaseException | None = None + + if cancel_task is not None: + try: + await _cancel_and_wait(cancel_task) + except BaseException as exc: + cleanup_error = exc + + if original_error is not None or process.returncode is None: + try: + await _terminate_process_tree(process) + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc + + try: + await _cleanup_process() + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc + + try: + await _cancel_and_wait(stderr_task) + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc + + if cleanup_error is not None: + raise cleanup_error + + cleanup_task = asyncio.create_task(_cleanup_resources()) + await _await_cleanup_task( + cleanup_task, + preserve_exception=original_error is not None, + ) def _build_env(self, args: CodexExecArgs) -> dict[str, str]: # Respect env overrides when provided; otherwise copy from os.environ. @@ -264,10 +291,74 @@ def _build_env(self, args: CodexExecArgs) -> dict[str, str]: return env -async def _watch_signal(signal: asyncio.Event, process: asyncio.subprocess.Process) -> None: - await signal.wait() - if process.returncode is None: - process.terminate() +async def _cancel_and_wait(task: asyncio.Task[object]) -> None: + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +async def _await_cleanup_task( + cleanup_task: asyncio.Task[None], + *, + preserve_exception: bool, +) -> None: + cancellation_error: asyncio.CancelledError | None = None + + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError as exc: + if cleanup_task.cancelled(): + break + cancellation_error = exc + except BaseException: + break + + cleanup_error: BaseException | None = None + try: + cleanup_task.result() + except BaseException as exc: + cleanup_error = exc + + if preserve_exception: + return + if cancellation_error is not None: + raise cancellation_error + if cleanup_error is not None: + raise cleanup_error + + +def _terminate_windows_process_tree(pid: int) -> None: + with contextlib.suppress(OSError, subprocess.TimeoutExpired): + subprocess.run( + ["taskkill", "/PID", str(pid), "/T", "/F"], + check=False, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS, + ) + + +async def _terminate_process_tree(process: asyncio.subprocess.Process) -> None: + pid = getattr(process, "pid", None) + try: + if pid is not None: + if os.name == "nt": + await asyncio.to_thread(_terminate_windows_process_tree, pid) + else: + with contextlib.suppress(ProcessLookupError): + os.killpg(pid, signal.SIGKILL) + finally: + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + + +async def _watch_signal(signal_event: asyncio.Event, process: asyncio.subprocess.Process) -> None: + await signal_event.wait() + await _terminate_process_tree(process) def _platform_target_triple() -> str: diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index ade5826fa7..cce9eda210 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -7,6 +7,7 @@ import json import os import signal +import subprocess import sys from dataclasses import fields from pathlib import Path @@ -361,6 +362,7 @@ async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: assert env[exec_module._INTERNAL_ORIGINATOR_ENV] == exec_module._TYPESCRIPT_SDK_ORIGINATOR assert env["OPENAI_BASE_URL"] == "https://example.com" assert env["CODEX_API_KEY"] == "api-key" + assert captured["kwargs"]["start_new_session"] is (os.name != "nt") @pytest.mark.asyncio @@ -506,6 +508,154 @@ async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: assert stderr_cancelled.is_set() +@pytest.mark.asyncio +async def test_codex_exec_run_cancellation_preserves_error_when_wait_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stdout_started = asyncio.Event() + stderr_cancelled = asyncio.Event() + + class BlockingStdout: + async def readline(self) -> bytes: + stdout_started.set() + await asyncio.Future[None]() + return b"" + + async def read(self, _size: int) -> bytes: + return b"" + + class BlockingStderr: + async def read(self, _size: int) -> bytes: + try: + await asyncio.Future[None]() + except asyncio.CancelledError: + stderr_cancelled.set() + raise + return b"" + + class FailingWaitProcess(FakeProcess): + async def wait(self) -> None: + raise RuntimeError("wait failed") + + process = FailingWaitProcess(stdout_lines=[], returncode=None) + process.stdout = cast(Any, BlockingStdout()) + process.stderr = cast(Any, BlockingStderr()) + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + + stream = exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello") + ) + next_line = asyncio.create_task(anext(stream)) + await asyncio.wait_for(stdout_started.wait(), timeout=1) + + next_line.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(next_line, timeout=1) + + assert stderr_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_codex_exec_run_preserves_execution_error_when_stdout_drain_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stderr_cancelled = asyncio.Event() + + class FailingDrainStdout(FakeStdout): + async def read(self, _size: int) -> bytes: + raise RuntimeError("stdout drain failed") + + class BlockingStderr: + async def read(self, _size: int) -> bytes: + try: + await asyncio.Future[None]() + except asyncio.CancelledError: + stderr_cancelled.set() + raise + return b"" + + process = FakeProcess(stdout_lines=[], returncode=None, stdin_present=False) + process.stdout = FailingDrainStdout([]) + process.stderr = cast(Any, BlockingStderr()) + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + + with pytest.raises(RuntimeError, match="no stdin"): + async for _ in exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello") + ): + pass + + assert stderr_cancelled.is_set() + + +@pytest.mark.asyncio +async def test_codex_exec_run_repeated_cancellation_finishes_stderr_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stdout_started = asyncio.Event() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + stderr_cancelled = asyncio.Event() + + class BlockingStdout: + async def readline(self) -> bytes: + stdout_started.set() + await asyncio.Future[None]() + return b"" + + async def read(self, _size: int) -> bytes: + return b"" + + class BlockingStderr: + async def read(self, _size: int) -> bytes: + try: + await asyncio.Future[None]() + except asyncio.CancelledError: + stderr_cancelled.set() + raise + return b"" + + process = FakeProcess(stdout_lines=[], returncode=None) + process.stdout = cast(Any, BlockingStdout()) + process.stderr = cast(Any, BlockingStderr()) + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + async def blocking_terminate_process_tree(_process: FakeProcess) -> None: + cleanup_started.set() + await release_cleanup.wait() + process.kill() + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + monkeypatch.setattr(exec_module, "_terminate_process_tree", blocking_terminate_process_tree) + + stream = exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello") + ) + next_line = asyncio.create_task(anext(stream)) + await asyncio.wait_for(stdout_started.wait(), timeout=1) + + next_line.cancel() + await asyncio.wait_for(cleanup_started.wait(), timeout=1) + next_line.cancel() + release_cleanup.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(next_line, timeout=1) + + assert stderr_cancelled.is_set() + + +@pytest.mark.skipif(sys.platform == "win32", reason="SelectorEventLoop lacks subprocess support") @pytest.mark.asyncio async def test_codex_exec_run_close_drains_stdout_before_wait( monkeypatch: pytest.MonkeyPatch, @@ -560,6 +710,7 @@ async def create_output_heavy_subprocess( await process.wait() +@pytest.mark.skipif(sys.platform == "win32", reason="SelectorEventLoop lacks subprocess support") @pytest.mark.asyncio async def test_codex_exec_run_close_is_bounded_when_descendant_holds_stdout( monkeypatch: pytest.MonkeyPatch, @@ -612,6 +763,14 @@ async def create_output_heavy_subprocess( await asyncio.wait_for(stream.aclose(), timeout=1) assert process.returncode is not None + for _ in range(500): + try: + os.kill(helper_pid, 0) + except ProcessLookupError: + break + await asyncio.sleep(0.01) + else: + pytest.fail("Codex subprocess descendant remained alive after stream closure") finally: if helper_pid is not None: with contextlib.suppress(ProcessLookupError, PermissionError): @@ -625,6 +784,58 @@ async def create_output_heavy_subprocess( await process.wait() +@pytest.mark.skipif(sys.platform != "win32", reason="Windows process-tree behavior") +def test_terminate_windows_process_tree_stops_descendant() -> None: + import ctypes + + helper_code = "import time; time.sleep(60)" + child_code = ( + "import subprocess\n" + "import sys\n" + "import time\n" + f"helper = subprocess.Popen([sys.executable, '-c', {helper_code!r}])\n" + "print(helper.pid, flush=True)\n" + "time.sleep(60)\n" + ) + process = subprocess.Popen( + [sys.executable, "-c", child_code], + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + helper_pid: int | None = None + + try: + assert process.stdout is not None + helper_pid = int(process.stdout.readline()) + kernel32 = cast(Any, ctypes).windll.kernel32 + helper_handle = kernel32.OpenProcess(0x00100000, False, helper_pid) + assert helper_handle + try: + exec_module._terminate_windows_process_tree(process.pid) + process.wait(timeout=5) + assert kernel32.WaitForSingleObject(helper_handle, 5000) == 0 + finally: + kernel32.CloseHandle(helper_handle) + finally: + if process.poll() is None: + subprocess.run( + ["taskkill", "/PID", str(process.pid), "/T", "/F"], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + process.wait(timeout=5) + if helper_pid is not None: + subprocess.run( + ["taskkill", "/PID", str(helper_pid), "/F"], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + @pytest.mark.asyncio async def test_codex_exec_run_raises_without_stdin(monkeypatch: pytest.MonkeyPatch) -> None: process = FakeProcess(stdout_lines=[], returncode=None, stdin_present=False) @@ -670,7 +881,7 @@ async def test_watch_signal_terminates_process() -> None: signal.set() await task - assert process.terminated is True + assert process.killed is True @pytest.mark.parametrize( From 378c14f764c1f3de2263b469e23210a8d2a84bc6 Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Thu, 23 Jul 2026 17:15:37 +0800 Subject: [PATCH 4/7] test(codex): tolerate defunct descendants --- .../codex/test_codex_exec_thread.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index cce9eda210..af4a8710bb 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -35,6 +35,25 @@ ) +def _process_is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + + if sys.platform.startswith("linux"): + try: + status = Path(f"/proc/{pid}/status").read_text() + except FileNotFoundError: + return False + for line in status.splitlines(): + if line.startswith("State:"): + # A zombie cannot execute but still responds to signal 0 until it is reaped. + return line.split()[1] != "Z" + + return True + + class FakeStdin: def __init__(self) -> None: self.buffer = b"" @@ -764,9 +783,7 @@ async def create_output_heavy_subprocess( await asyncio.wait_for(stream.aclose(), timeout=1) assert process.returncode is not None for _ in range(500): - try: - os.kill(helper_pid, 0) - except ProcessLookupError: + if not _process_is_running(helper_pid): break await asyncio.sleep(0.01) else: From b6bd3027997a6713e58fa8714b321d72bf53a88d Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Thu, 23 Jul 2026 17:34:00 +0800 Subject: [PATCH 5/7] fix(codex): terminate descendants after clean exit --- .../extensions/experimental/codex/exec.py | 11 ++-- .../codex/test_codex_exec_thread.py | 59 +++++++++++++++++++ 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index deeb02c2c2..c789282cb9 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -243,12 +243,11 @@ async def _cleanup_resources() -> None: except BaseException as exc: cleanup_error = exc - if original_error is not None or process.returncode is None: - try: - await _terminate_process_tree(process) - except BaseException as exc: - if cleanup_error is None: - cleanup_error = exc + try: + await _terminate_process_tree(process) + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc try: await _cleanup_process() diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index af4a8710bb..417909642b 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -9,6 +9,7 @@ import signal import subprocess import sys +from collections.abc import AsyncGenerator from dataclasses import fields from pathlib import Path from typing import Any, cast @@ -801,6 +802,64 @@ async def create_output_heavy_subprocess( await process.wait() +@pytest.mark.skipif(sys.platform == "win32", reason="SelectorEventLoop lacks subprocess support") +@pytest.mark.asyncio +async def test_codex_exec_run_terminates_descendant_after_clean_parent_exit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_create_subprocess_exec = asyncio.create_subprocess_exec + process: asyncio.subprocess.Process | None = None + helper_pid: int | None = None + helper_code = "import time; time.sleep(10)" + child_code = ( + "import os\n" + "import subprocess\n" + "import sys\n" + "sys.stdin.buffer.read()\n" + f"helper = subprocess.Popen([sys.executable, '-c', {helper_code!r}], " + "stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=sys.stderr)\n" + "os.write(1, f'ready {helper.pid}\\n'.encode())\n" + ) + + async def create_process_with_descendant( + *_args: Any, **kwargs: Any + ) -> asyncio.subprocess.Process: + nonlocal process + process = await original_create_subprocess_exec(sys.executable, "-c", child_code, **kwargs) + return process + + async def collect_output(stream: AsyncGenerator[str, None]) -> list[str]: + return [line async for line in stream] + + monkeypatch.setattr( + exec_module.asyncio, "create_subprocess_exec", create_process_with_descendant + ) + + stream = exec_module.CodexExec(executable_path="unused").run( + exec_module.CodexExecArgs(input="hello") + ) + + try: + output = await asyncio.wait_for(collect_output(stream), timeout=5) + helper_pid = int(output[0].removeprefix("ready ")) + assert process is not None + assert process.returncode == 0 + for _ in range(500): + if not _process_is_running(helper_pid): + break + await asyncio.sleep(0.01) + else: + pytest.fail("Codex subprocess descendant remained alive after clean parent exit") + finally: + if helper_pid is not None: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.kill(helper_pid, signal.SIGTERM) + if process is not None: + if process.returncode is None: + process.kill() + await process.wait() + + @pytest.mark.skipif(sys.platform != "win32", reason="Windows process-tree behavior") def test_terminate_windows_process_tree_stops_descendant() -> None: import ctypes From 7552c7204cd9da0781df5a79d7b1651b4b5cd4f1 Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Fri, 24 Jul 2026 10:04:18 +0800 Subject: [PATCH 6/7] fix(codex): observe direct process exit before cleanup --- .../extensions/experimental/codex/exec.py | 60 ++++++++-- .../codex/test_codex_exec_thread.py | 107 ++++++++++++++++++ 2 files changed, 158 insertions(+), 9 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index c789282cb9..325d095e58 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -133,6 +133,7 @@ async def run(self, args: CodexExecArgs) -> AsyncGenerator[str, None]: # Give POSIX descendants a process-group boundary that this execution owns. start_new_session=os.name != "nt", ) + direct_process_exit = _create_process_exit_event(process) stderr_chunks: list[bytes] = [] @@ -174,7 +175,23 @@ async def _cleanup_process() -> None: timeout=_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS, ) + async def _wait_for_direct_process() -> None: + if direct_process_exit is None: + await process.wait() + else: + await direct_process_exit.wait() + stderr_task = asyncio.create_task(_drain_stderr()) + + async def _finish_stderr() -> None: + try: + await asyncio.wait_for( + asyncio.shield(stderr_task), + timeout=_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + await _cancel_and_wait(stderr_task) + cancel_task: asyncio.Task[None] | None = None try: if process.stdin is None: @@ -224,13 +241,9 @@ async def _read_stdout_line() -> bytes: break yield line.decode("utf-8").rstrip("\n") - await process.wait() - if process.returncode not in (0, None): - await stderr_task - stderr_text = b"".join(stderr_chunks).decode("utf-8") - raise RuntimeError( - f"Codex exec exited with code {process.returncode}: {stderr_text}" - ) + # Wait for the direct process independently of pipes inherited by descendants. + await _wait_for_direct_process() + finally: original_error = sys.exc_info()[1] @@ -256,12 +269,15 @@ async def _cleanup_resources() -> None: cleanup_error = exc try: - await _cancel_and_wait(stderr_task) + if original_error is None: + await _finish_stderr() + else: + await _cancel_and_wait(stderr_task) except BaseException as exc: if cleanup_error is None: cleanup_error = exc - if cleanup_error is not None: + if cleanup_error is not None and process.returncode in (0, None): raise cleanup_error cleanup_task = asyncio.create_task(_cleanup_resources()) @@ -270,6 +286,10 @@ async def _cleanup_resources() -> None: preserve_exception=original_error is not None, ) + if process.returncode not in (0, None): + stderr_text = b"".join(stderr_chunks).decode("utf-8") + raise RuntimeError(f"Codex exec exited with code {process.returncode}: {stderr_text}") + def _build_env(self, args: CodexExecArgs) -> dict[str, str]: # Respect env overrides when provided; otherwise copy from os.environ. env: dict[str, str] = {} @@ -297,6 +317,28 @@ async def _cancel_and_wait(task: asyncio.Task[object]) -> None: await task +def _create_process_exit_event( + process: asyncio.subprocess.Process, +) -> asyncio.Event | None: + protocol = getattr(process, "_protocol", None) + process_exited = getattr(protocol, "process_exited", None) + if protocol is None or not callable(process_exited): + return None + + process_exit_event = asyncio.Event() + + def _process_exited() -> None: + try: + process_exited() + finally: + process_exit_event.set() + + protocol.process_exited = _process_exited + if process.returncode is not None: + process_exit_event.set() + return process_exit_event + + async def _await_cleanup_task( cleanup_task: asyncio.Task[None], *, diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index 417909642b..55e34e498c 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -97,6 +97,14 @@ async def read(self, _size: int) -> bytes: return self._chunks.pop(0) +class FakeProcessProtocol: + def __init__(self) -> None: + self.process_exited_calls = 0 + + def process_exited(self) -> None: + self.process_exited_calls += 1 + + class FakeProcess: def __init__( self, @@ -114,6 +122,7 @@ def __init__( self.returncode = returncode self.killed = False self.terminated = False + self._protocol: FakeProcessProtocol | None = None async def wait(self) -> None: if self.returncode is None: @@ -675,6 +684,104 @@ async def blocking_terminate_process_tree(_process: FakeProcess) -> None: assert stderr_cancelled.is_set() +@pytest.mark.asyncio +async def test_codex_exec_run_waits_for_direct_exit_after_stdout_eof( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stdout_eof = asyncio.Event() + tree_terminated = asyncio.Event() + + class EofStdout(FakeStdout): + async def readline(self) -> bytes: + stdout_eof.set() + return b"" + + process = FakeProcess(stdout_lines=[], returncode=None) + process.stdout = EofStdout([]) + process._protocol = FakeProcessProtocol() + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + async def terminate_process_tree(_process: FakeProcess) -> None: + tree_terminated.set() + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + monkeypatch.setattr(exec_module, "_terminate_process_tree", terminate_process_tree) + + stream = exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello") + ) + next_line = asyncio.create_task(anext(stream)) + await asyncio.wait_for(stdout_eof.wait(), timeout=1) + await asyncio.sleep(0) + + assert tree_terminated.is_set() is False + + process.returncode = 0 + assert process._protocol is not None + process._protocol.process_exited() + + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(next_line, timeout=1) + + assert tree_terminated.is_set() + assert process.returncode == 0 + assert process.killed is False + assert process._protocol.process_exited_calls == 1 + + +@pytest.mark.asyncio +async def test_create_process_exit_event_observes_already_exited_process() -> None: + process = FakeProcess(stdout_lines=[], returncode=0) + process._protocol = FakeProcessProtocol() + + process_exit_event = exec_module._create_process_exit_event(cast(Any, process)) + + assert process_exit_event is not None + assert process_exit_event.is_set() + + +@pytest.mark.asyncio +async def test_codex_exec_run_uses_direct_exit_when_descendant_holds_inherited_pipes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tree_terminated = asyncio.Event() + + class ExitOnEofStdout(FakeStdout): + async def readline(self) -> bytes: + process.returncode = 0 + assert process._protocol is not None + process._protocol.process_exited() + return b"" + + class PipeHoldingProcess(FakeProcess): + async def wait(self) -> None: + await tree_terminated.wait() + + process = PipeHoldingProcess(stdout_lines=[], returncode=None) + process.stdout = ExitOnEofStdout([]) + process._protocol = FakeProcessProtocol() + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + async def terminate_process_tree(_process: FakeProcess) -> None: + tree_terminated.set() + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + monkeypatch.setattr(exec_module, "_terminate_process_tree", terminate_process_tree) + + stream = exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello") + ) + + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(stream), timeout=1) + + assert tree_terminated.is_set() + + @pytest.mark.skipif(sys.platform == "win32", reason="SelectorEventLoop lacks subprocess support") @pytest.mark.asyncio async def test_codex_exec_run_close_drains_stdout_before_wait( From 8b1c969db29f988c3a07ad827a72f8a599eb9797 Mon Sep 17 00:00:00 2001 From: shiwenbin Date: Fri, 24 Jul 2026 10:58:47 +0800 Subject: [PATCH 7/7] fix(codex): harden subprocess tree cleanup --- .../extensions/experimental/codex/exec.py | 57 +++++-- .../codex/test_codex_exec_thread.py | 148 ++++++++++++++++-- 2 files changed, 179 insertions(+), 26 deletions(-) diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index 325d095e58..a6509a8bf2 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -10,7 +10,7 @@ import sys from collections.abc import AsyncGenerator from dataclasses import dataclass -from pathlib import Path +from pathlib import Path, PureWindowsPath from agents.exceptions import UserError @@ -134,6 +134,25 @@ async def run(self, args: CodexExecArgs) -> AsyncGenerator[str, None]: start_new_session=os.name != "nt", ) direct_process_exit = _create_process_exit_event(process) + process_tree_termination_task: asyncio.Task[None] | None = None + + async def _terminate_process_tree_once() -> None: + nonlocal process_tree_termination_task + if process_tree_termination_task is None: + process_tree_termination_task = asyncio.create_task( + _terminate_process_tree(process) + ) + await asyncio.shield(process_tree_termination_task) + + async def _watch_termination_event(event: asyncio.Event) -> None: + await event.wait() + await _terminate_process_tree_once() + + process_exit_cleanup_task = ( + asyncio.create_task(_watch_termination_event(direct_process_exit)) + if direct_process_exit is not None + else None + ) stderr_chunks: list[bytes] = [] @@ -207,7 +226,7 @@ async def _finish_stderr() -> None: if args.signal is not None: # Mirror AbortSignal semantics by terminating the subprocess. - cancel_task = asyncio.create_task(_watch_signal(args.signal, process)) + cancel_task = asyncio.create_task(_watch_termination_event(args.signal)) async def _read_stdout_line() -> bytes: if args.idle_timeout_seconds is None: @@ -256,8 +275,15 @@ async def _cleanup_resources() -> None: except BaseException as exc: cleanup_error = exc + if process_exit_cleanup_task is not None: + try: + await _cancel_and_wait(process_exit_cleanup_task) + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc + try: - await _terminate_process_tree(process) + await _terminate_process_tree_once() except BaseException as exc: if cleanup_error is None: cleanup_error = exc @@ -370,10 +396,28 @@ async def _await_cleanup_task( raise cleanup_error +def _windows_system_taskkill_path() -> str | None: + import ctypes + + windll = getattr(ctypes, "windll", None) + if windll is None: + return None + + buffer = ctypes.create_unicode_buffer(32768) + length = windll.kernel32.GetSystemDirectoryW(buffer, len(buffer)) + if length == 0 or length >= len(buffer): + return None + return str(PureWindowsPath(buffer.value) / "taskkill.exe") + + def _terminate_windows_process_tree(pid: int) -> None: + taskkill_path = _windows_system_taskkill_path() + if taskkill_path is None: + return + with contextlib.suppress(OSError, subprocess.TimeoutExpired): subprocess.run( - ["taskkill", "/PID", str(pid), "/T", "/F"], + [taskkill_path, "/PID", str(pid), "/T", "/F"], check=False, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, @@ -397,11 +441,6 @@ async def _terminate_process_tree(process: asyncio.subprocess.Process) -> None: process.kill() -async def _watch_signal(signal_event: asyncio.Event, process: asyncio.subprocess.Process) -> None: - await signal_event.wait() - await _terminate_process_tree(process) - - def _platform_target_triple() -> str: # Map the running platform to the vendor layout used in Codex releases. system = sys.platform diff --git a/tests/extensions/experiemental/codex/test_codex_exec_thread.py b/tests/extensions/experiemental/codex/test_codex_exec_thread.py index 55e34e498c..920b4fd179 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -743,16 +743,19 @@ async def test_create_process_exit_event_observes_already_exited_process() -> No @pytest.mark.asyncio -async def test_codex_exec_run_uses_direct_exit_when_descendant_holds_inherited_pipes( +async def test_codex_exec_run_terminates_tree_when_descendant_holds_stdout_open( monkeypatch: pytest.MonkeyPatch, ) -> None: tree_terminated = asyncio.Event() - class ExitOnEofStdout(FakeStdout): + class PipeHoldingStdout(FakeStdout): async def readline(self) -> bytes: - process.returncode = 0 - assert process._protocol is not None - process._protocol.process_exited() + if process.returncode is None: + process.returncode = 0 + assert process._protocol is not None + process._protocol.process_exited() + return b"ready\n" + await tree_terminated.wait() return b"" class PipeHoldingProcess(FakeProcess): @@ -760,7 +763,7 @@ async def wait(self) -> None: await tree_terminated.wait() process = PipeHoldingProcess(stdout_lines=[], returncode=None) - process.stdout = ExitOnEofStdout([]) + process.stdout = PipeHoldingStdout([]) process._protocol = FakeProcessProtocol() async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: @@ -776,9 +779,12 @@ async def terminate_process_tree(_process: FakeProcess) -> None: exec_module.CodexExecArgs(input="hello") ) - with pytest.raises(StopAsyncIteration): - await asyncio.wait_for(anext(stream), timeout=1) + async def collect_output() -> list[str]: + return [line async for line in stream] + + output = await asyncio.wait_for(collect_output(), timeout=1) + assert output == ["ready"] assert tree_terminated.is_set() @@ -911,20 +917,24 @@ async def create_output_heavy_subprocess( @pytest.mark.skipif(sys.platform == "win32", reason="SelectorEventLoop lacks subprocess support") @pytest.mark.asyncio +@pytest.mark.parametrize("inherited_stream", ["stdout", "stderr"]) async def test_codex_exec_run_terminates_descendant_after_clean_parent_exit( monkeypatch: pytest.MonkeyPatch, + inherited_stream: str, ) -> None: original_create_subprocess_exec = asyncio.create_subprocess_exec process: asyncio.subprocess.Process | None = None helper_pid: int | None = None helper_code = "import time; time.sleep(10)" + helper_stdout = "sys.stdout" if inherited_stream == "stdout" else "subprocess.DEVNULL" + helper_stderr = "sys.stderr" if inherited_stream == "stderr" else "subprocess.DEVNULL" child_code = ( "import os\n" "import subprocess\n" "import sys\n" "sys.stdin.buffer.read()\n" f"helper = subprocess.Popen([sys.executable, '-c', {helper_code!r}], " - "stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=sys.stderr)\n" + f"stdin=subprocess.DEVNULL, stdout={helper_stdout}, stderr={helper_stderr})\n" "os.write(1, f'ready {helper.pid}\\n'.encode())\n" ) @@ -967,10 +977,74 @@ async def collect_output(stream: AsyncGenerator[str, None]) -> list[str]: await process.wait() +def test_terminate_windows_process_tree_uses_trusted_taskkill_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import ctypes + + class FakeKernel32: + def GetSystemDirectoryW(self, buffer: Any, _size: int) -> int: + buffer.value = r"C:\Windows\System32" + return len(buffer.value) + + class FakeWindll: + kernel32 = FakeKernel32() + + trusted_path = r"C:\Windows\System32\taskkill.exe" + captured_command: list[str] | None = None + + def fake_run(command: list[str], **_kwargs: Any) -> subprocess.CompletedProcess[bytes]: + nonlocal captured_command + captured_command = command + return subprocess.CompletedProcess(command, 0) + + monkeypatch.setattr(ctypes, "windll", FakeWindll(), raising=False) + monkeypatch.setattr(exec_module.subprocess, "run", fake_run) + + exec_module._terminate_windows_process_tree(123) + + assert captured_command == [trusted_path, "/PID", "123", "/T", "/F"] + + +def test_terminate_windows_process_tree_does_not_fallback_to_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import ctypes + + def unexpected_run(*_args: Any, **_kwargs: Any) -> None: + pytest.fail("taskkill must not be resolved through PATH") + + monkeypatch.delattr(ctypes, "windll", raising=False) + monkeypatch.setattr(exec_module.subprocess, "run", unexpected_run) + + exec_module._terminate_windows_process_tree(123) + + +@pytest.mark.parametrize("length", [0, 32768]) +def test_windows_system_taskkill_path_rejects_invalid_length( + monkeypatch: pytest.MonkeyPatch, + length: int, +) -> None: + import ctypes + + class FakeKernel32: + def GetSystemDirectoryW(self, _buffer: Any, _size: int) -> int: + return length + + class FakeWindll: + kernel32 = FakeKernel32() + + monkeypatch.setattr(ctypes, "windll", FakeWindll(), raising=False) + + assert exec_module._windows_system_taskkill_path() is None + + @pytest.mark.skipif(sys.platform != "win32", reason="Windows process-tree behavior") def test_terminate_windows_process_tree_stops_descendant() -> None: import ctypes + taskkill_path = exec_module._windows_system_taskkill_path() + assert taskkill_path is not None helper_code = "import time; time.sleep(60)" child_code = ( "import subprocess\n" @@ -1004,7 +1078,7 @@ def test_terminate_windows_process_tree_stops_descendant() -> None: finally: if process.poll() is None: subprocess.run( - ["taskkill", "/PID", str(process.pid), "/T", "/F"], + [taskkill_path, "/PID", str(process.pid), "/T", "/F"], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, @@ -1012,7 +1086,7 @@ def test_terminate_windows_process_tree_stops_descendant() -> None: process.wait(timeout=5) if helper_pid is not None: subprocess.run( - ["taskkill", "/PID", str(helper_pid), "/F"], + [taskkill_path, "/PID", str(helper_pid), "/F"], check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, @@ -1056,15 +1130,55 @@ async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: @pytest.mark.asyncio -async def test_watch_signal_terminates_process() -> None: - signal = asyncio.Event() +async def test_codex_exec_run_terminates_tree_once_when_signal_and_exit_overlap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stdout_started = asyncio.Event() + tree_terminated = asyncio.Event() + termination_calls = 0 + + class BlockingStdout(FakeStdout): + async def readline(self) -> bytes: + stdout_started.set() + await tree_terminated.wait() + return b"" + process = FakeProcess(stdout_lines=[], returncode=None) + process.stdout = BlockingStdout([]) + process._protocol = FakeProcessProtocol() - task = asyncio.create_task(exec_module._watch_signal(signal, process)) - signal.set() - await task + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process - assert process.killed is True + async def terminate_process_tree(_process: FakeProcess) -> None: + nonlocal termination_calls + termination_calls += 1 + if termination_calls == 1: + process.returncode = -9 + assert process._protocol is not None + process._protocol.process_exited() + await asyncio.sleep(0) + tree_terminated.set() + + monkeypatch.setattr(exec_module.asyncio, "create_subprocess_exec", fake_create_subprocess_exec) + monkeypatch.setattr(exec_module, "_terminate_process_tree", terminate_process_tree) + + signal_event = asyncio.Event() + stream = exec_module.CodexExec(executable_path="/bin/codex").run( + exec_module.CodexExecArgs(input="hello", signal=signal_event) + ) + + async def collect_output() -> list[str]: + return [line async for line in stream] + + output_task = asyncio.create_task(collect_output()) + await asyncio.wait_for(stdout_started.wait(), timeout=1) + signal_event.set() + + with pytest.raises(RuntimeError, match="exited with code -9"): + await asyncio.wait_for(output_task, timeout=1) + assert termination_calls == 1 + assert process._protocol.process_exited_calls == 1 @pytest.mark.parametrize(