-
Notifications
You must be signed in to change notification settings - Fork 4.4k
fix(codex): clean up subprocess on cancellation #3915
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
8e7af59
da45e04
c004807
378c14f
b6bd302
7552c72
8b1c969
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) | ||
|
Comment on lines
+436
to
+437
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
The new cleanup only kills the subprocess' original POSIX process group. A Codex-run command can fork and call Useful? React with 👍 / 👎.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks — this is a valid limitation. The SDK-level cleanup here is best-effort rather than a security boundary: on POSIX, a descendant that deliberately creates a new session can escape the owned process group, and covering that portably would require an OS-specific containment primitive owned by the Codex CLI or sandbox rather than recursive process scanning in the SDK. This PR will remain scoped to terminating the Codex CLI and ordinary, non-detached descendants during cancellation or stream closure. I’ll leave the implementation unchanged and defer hard containment to the CLI/sandbox layer. |
||
| finally: | ||
| if process.returncode is None: | ||
| with contextlib.suppress(ProcessLookupError): | ||
| process.kill() | ||
|
|
||
|
|
||
| def _platform_target_triple() -> str: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When the direct Codex process exits after emitting its last event but a descendant keeps stdout inherited/open, the next
stdout.readline()never returns EOF, so execution never reaches this direct-exit wait or thefinallycleanup that kills the process group. I reproduced this with a subprocess that prints one line, spawns a sleeping child withstdout=sys.stdout, and exits; the parent returncode was already0while the secondanext()hung. Race the direct-exit event with the stdout read, or trigger cleanup once the direct process exits, so clean exits with inherited stdout do not hang indefinitely.Useful? React with 👍 / 👎.