diff --git a/src/agents/extensions/experimental/codex/exec.py b/src/agents/extensions/experimental/codex/exec.py index c83a0e98cd..a6509a8bf2 100644 --- a/src/agents/extensions/experimental/codex/exec.py +++ b/src/agents/extensions/experimental/codex/exec.py @@ -5,10 +5,12 @@ import os import platform import shutil +import signal +import subprocess 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 @@ -20,6 +22,8 @@ _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 +_SUBPROCESS_CLEANUP_TIMEOUT_SECONDS = 1.0 @dataclass(frozen=True) @@ -126,6 +130,28 @@ 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", + ) + 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] = [] @@ -140,73 +166,155 @@ async def _drain_stderr() -> None: break stderr_chunks.append(chunk) - stderr_task = asyncio.create_task(_drain_stderr()) + async def _drain_remaining_stdout() -> None: + if process.stdout is None: + return + while await process.stdout.read(_SUBPROCESS_DRAIN_CHUNK_BYTES): + pass - if process.stdin is None: - process.kill() - raise RuntimeError("Codex subprocess has no stdin") + async def _drain_stdout_and_wait() -> None: + await _drain_remaining_stdout() + await process.wait() - process.stdin.write(args.input.encode("utf-8")) - await process.stdin.drain() - process.stdin.close() + 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, + ) + + 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()) - if process.stdout is None: - process.kill() - raise RuntimeError("Codex subprocess has no stdout") - stdout = process.stdout + 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 - 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_termination_event(args.signal)) + + 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() + + 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: break 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") - 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: - if cancel_task is not None and not cancel_task.done(): - cancel_task.cancel() - await stderr_task - if process.returncode is None: - process.kill() + 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 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_once() + 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: + 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 and process.returncode in (0, None): + raise cleanup_error + + cleanup_task = asyncio.create_task(_cleanup_resources()) + await _await_cleanup_task( + cleanup_task, + 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. @@ -228,10 +336,109 @@ 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 + + +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], + *, + 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 _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_path, "/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() 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 c1012c51b3..920b4fd179 100644 --- a/tests/extensions/experiemental/codex/test_codex_exec_thread.py +++ b/tests/extensions/experiemental/codex/test_codex_exec_thread.py @@ -1,10 +1,15 @@ from __future__ import annotations import asyncio +import contextlib import importlib import inspect import json import os +import signal +import subprocess +import sys +from collections.abc import AsyncGenerator from dataclasses import fields from pathlib import Path from typing import Any, cast @@ -31,6 +36,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"" @@ -55,6 +79,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: @@ -66,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, @@ -83,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: @@ -351,6 +391,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 @@ -442,9 +483,619 @@ 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_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.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_terminates_tree_when_descendant_holds_stdout_open( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tree_terminated = asyncio.Event() + + class PipeHoldingStdout(FakeStdout): + async def readline(self) -> bytes: + 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): + async def wait(self) -> None: + await tree_terminated.wait() + + process = PipeHoldingProcess(stdout_lines=[], returncode=None) + process.stdout = PipeHoldingStdout([]) + 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") + ) + + 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() + + +@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, +) -> 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.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, +) -> 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 + 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 stream closure") + 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.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}], " + f"stdin=subprocess.DEVNULL, stdout={helper_stdout}, stderr={helper_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() + + +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" + "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_path, "/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_path, "/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=[], 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 +1113,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 @@ -479,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() + + async def fake_create_subprocess_exec(*args: Any, **kwargs: Any) -> FakeProcess: + return process + + 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] - task = asyncio.create_task(exec_module._watch_signal(signal, process)) - signal.set() - await task + output_task = asyncio.create_task(collect_output()) + await asyncio.wait_for(stdout_started.wait(), timeout=1) + signal_event.set() - assert process.terminated is True + 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(