Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
317 changes: 262 additions & 55 deletions src/agents/extensions/experimental/codex/exec.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
Expand Down Expand Up @@ -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] = []
Expand All @@ -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()
Comment on lines +263 to +264

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Stop waiting for stdout before reaping descendants

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 the finally cleanup that kills the process group. I reproduced this with a subprocess that prints one line, spawns a sleeping child with stdout=sys.stdout, and exits; the parent returncode was already 0 while the second anext() 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 👍 / 👎.


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.
Expand All @@ -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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Security: Contain daemonized Codex descendants during cleanup

The new cleanup only kills the subprocess' original POSIX process group. A Codex-run command can fork and call setsid()/setpgrp() before cancellation; that descendant leaves pgid pid, so stream.aclose() or an idle timeout returns while it keeps running with the configured workspace/environment. If cancellation is meant to revoke Codex execution, use a containment primitive that covers reparented descendants, or document this as best-effort and enforce hard cleanup in the sandbox/CLI.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The 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:
Expand Down
Loading