From 1b53b415b6a8bf84113b0a2faa0a4174c8b27a39 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 14 Jul 2026 19:47:32 +0800 Subject: [PATCH] core/tools/process: harden PTY sessions for agents Structured alive/exit_code/data reads; until= wait; SIGTERM then SIGKILL; exec-fail marker; session limit, idle TTL, output budget. --- CLAUDE.md | 3 +- src/plyngent/tools/process/close_pty.py | 8 +- src/plyngent/tools/process/open_pty.py | 15 +- src/plyngent/tools/process/pty_session.py | 344 ++++++++++++++++++++-- src/plyngent/tools/process/read_pty.py | 29 +- src/plyngent/tools/process/write_pty.py | 13 +- tests/test_tools/test_process.py | 113 +++++-- 7 files changed, 456 insertions(+), 69 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 2596246..ddcad61 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -63,7 +63,8 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use. - **`workspace`**: path resolve under root; path substring denylist; command basename denylist. - **`file`**: `read_file`, `write_file`, `listdir`, `edit_replace` (first occurrence). -- **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` / `close_pty`. +- **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` / `close_pty` (**Unix only**: `pty`+`fork`). +- PTY: structured status (`alive`/`exit_code`/`data`); `read_pty(..., until=)`; session limit/idle TTL/output budget; close SIGTERM→SIGKILL. - **`DEFAULT_TOOLS`**: file + process tool list for a `ToolRegistry`. ### CLI (`cli/`) diff --git a/src/plyngent/tools/process/close_pty.py b/src/plyngent/tools/process/close_pty.py index 68e7648..c25dd90 100644 --- a/src/plyngent/tools/process/close_pty.py +++ b/src/plyngent/tools/process/close_pty.py @@ -2,11 +2,11 @@ from __future__ import annotations from plyngent.agent import tool -from .pty_session import PtyManager +from .pty_session import PtyManager, format_close_result @tool def close_pty(session_id: int) -> str: - """Close a PTY session and terminate its process.""" - PtyManager.close(session_id) - return f"closed session_id={session_id}" + """Close a PTY session (SIGTERM, then SIGKILL after a short grace period).""" + result = PtyManager.close(session_id) + return format_close_result(result) diff --git a/src/plyngent/tools/process/open_pty.py b/src/plyngent/tools/process/open_pty.py index afb2f1b..2e4c686 100644 --- a/src/plyngent/tools/process/open_pty.py +++ b/src/plyngent/tools/process/open_pty.py @@ -8,11 +8,22 @@ from .pty_session import PtyManager @tool def open_pty(command: list[str], *, cwd: str = ".") -> str: - """Open a PTY session running ``command`` (argv) under the workspace; returns session id.""" + """Open a Unix PTY session running ``command`` (argv) under the workspace. + + Returns structured status including session_id. Not supported on Windows. + Failed exec surfaces via later read_pty data (marker) and exit_code=127. + """ try: session = PtyManager.open(command, cwd=cwd) except WorkspaceError as exc: return f"error: {exc}" except OSError as exc: return f"error: failed to open PTY: {exc}" - return f"session_id={session.session_id}" + return "\n".join( + [ + f"session_id={session.session_id}", + "alive=true", + "exit_code=", + f"cmd={' '.join(session.command)}", + ] + ) diff --git a/src/plyngent/tools/process/pty_session.py b/src/plyngent/tools/process/pty_session.py index 9f2e89e..a018d2d 100644 --- a/src/plyngent/tools/process/pty_session.py +++ b/src/plyngent/tools/process/pty_session.py @@ -1,3 +1,5 @@ +"""Unix PTY session registry (requires ``pty`` + ``os.fork``; not supported on Windows).""" + from __future__ import annotations import contextlib @@ -13,7 +15,13 @@ from typing import ClassVar from plyngent.tools.workspace import WorkspaceError, check_command_allowed, resolve_path DEFAULT_PTY_READ_BYTES = 8192 +DEFAULT_PTY_POLL_TIMEOUT = 0.2 +DEFAULT_MAX_SESSIONS = 8 +DEFAULT_IDLE_TTL_SECONDS = 600.0 +DEFAULT_SESSION_OUTPUT_BUDGET = 256_000 +DEFAULT_CLOSE_GRACE_SECONDS = 0.5 _STDERR_FD = 2 +_EXEC_FAIL_MARKER = b"plyngent-pty-exec-failed: " @dataclass @@ -22,15 +30,58 @@ class PtySession: master_fd: int pid: int closed: bool = False + alive: bool = True + exit_code: int | None = None created_at: float = field(default_factory=time.time) + last_activity: float = field(default_factory=time.time) + bytes_read: int = 0 + command: tuple[str, ...] = () + + +@dataclass(frozen=True) +class PtyReadResult: + session_id: int + alive: bool + exit_code: int | None + data: str + truncated: bool = False + matched: bool = False + budget_exhausted: bool = False + + +@dataclass(frozen=True) +class PtyCloseResult: + session_id: int + closed: bool + alive: bool + exit_code: int | None + message: str = "" class PtyManager: - """In-process PTY session registry (minimal).""" + """In-process PTY session registry (process-global; suitable for local CLI).""" _lock: ClassVar[Lock] = Lock() _next_id: ClassVar[int] = 1 _sessions: ClassVar[dict[int, PtySession]] = {} + max_sessions: ClassVar[int] = DEFAULT_MAX_SESSIONS + idle_ttl_seconds: ClassVar[float] = DEFAULT_IDLE_TTL_SECONDS + session_output_budget: ClassVar[int] = DEFAULT_SESSION_OUTPUT_BUDGET + + @classmethod + def configure( + cls, + *, + max_sessions: int | None = None, + idle_ttl_seconds: float | None = None, + session_output_budget: int | None = None, + ) -> None: + if max_sessions is not None: + cls.max_sessions = max(1, max_sessions) + if idle_ttl_seconds is not None: + cls.idle_ttl_seconds = max(0.0, idle_ttl_seconds) + if session_output_budget is not None: + cls.session_output_budget = max(1024, session_output_budget) @classmethod def open( @@ -45,6 +96,13 @@ class PtyManager: msg = f"cwd is not a directory: {cwd}" raise WorkspaceError(msg) + _ = cls.reap_idle() + with cls._lock: + alive_count = sum(1 for s in cls._sessions.values() if not s.closed) + if alive_count >= cls.max_sessions: + msg = f"PTY session limit reached ({cls.max_sessions}); close idle sessions" + raise WorkspaceError(msg) + master_fd, slave_fd = pty.openpty() pid = os.fork() if pid == 0: # child @@ -58,15 +116,21 @@ class PtyManager: os.close(slave_fd) _ = os.chdir(workdir) os.execvp(command[0], command) - except OSError: - pass + except OSError as exc: + with contextlib.suppress(OSError): + _ = os.write(1, _EXEC_FAIL_MARKER + str(exc).encode(errors="replace") + b"\n") os._exit(127) os.close(slave_fd) with cls._lock: session_id = cls._next_id cls._next_id += 1 - session = PtySession(session_id=session_id, master_fd=master_fd, pid=pid) + session = PtySession( + session_id=session_id, + master_fd=master_fd, + pid=pid, + command=tuple(command), + ) cls._sessions[session_id] = session return session @@ -76,45 +140,267 @@ class PtyManager: return cls._sessions.get(session_id) @classmethod - def read(cls, session_id: int, *, max_bytes: int = DEFAULT_PTY_READ_BYTES, timeout: float = 0.1) -> str: - session = cls.get(session_id) - if session is None or session.closed: - msg = f"unknown or closed PTY session: {session_id}" - raise WorkspaceError(msg) - ready, _, _ = select.select([session.master_fd], [], [], timeout) - if not ready: - return "" + def _touch(cls, session: PtySession) -> None: + session.last_activity = time.time() + + @classmethod + def _poll_exit(cls, session: PtySession) -> None: + if not session.alive or session.closed: + return try: - data = os.read(session.master_fd, max_bytes) - except OSError: - return "" - return data.decode(errors="replace") + waited_pid, status = os.waitpid(session.pid, os.WNOHANG) + except ChildProcessError: + session.alive = False + return + if waited_pid == 0: + return + session.alive = False + if os.WIFEXITED(status): + session.exit_code = os.WEXITSTATUS(status) + elif os.WIFSIGNALED(status): + session.exit_code = -os.WTERMSIG(status) + else: + session.exit_code = status + + @classmethod + def refresh(cls, session_id: int) -> PtySession: + session = cls.get(session_id) + if session is None: + msg = f"unknown PTY session: {session_id}" + raise WorkspaceError(msg) + cls._poll_exit(session) + return session + + @classmethod + def _read_once(cls, session: PtySession, *, max_bytes: int, timeout: float) -> bytes: + if session.closed or (cls.session_output_budget - session.bytes_read) <= 0: + return b"" + to_read = min(max_bytes, cls.session_output_budget - session.bytes_read) + data = b"" + try: + ready, _, _ = select.select([session.master_fd], [], [], timeout) + if ready: + data = os.read(session.master_fd, to_read) + except (OSError, ValueError): + data = b"" + cls._poll_exit(session) + if not data: + return b"" + session.bytes_read += len(data) + cls._touch(session) + return data + + @classmethod + def _collect_chunks( + cls, + session: PtySession, + *, + max_bytes: int, + timeout: float, + until: str | None, + ) -> tuple[list[bytes], bool]: + chunks: list[bytes] = [] + matched = False + if until is None: + chunk = cls._read_once(session, max_bytes=max_bytes, timeout=timeout) + if chunk: + chunks.append(chunk) + return chunks, matched + + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + chunk = cls._read_once( + session, + max_bytes=max_bytes, + timeout=min(DEFAULT_PTY_POLL_TIMEOUT, remaining), + ) + if chunk: + chunks.append(chunk) + if until in b"".join(chunks).decode(errors="replace"): + matched = True + break + cls._poll_exit(session) + if not session.alive or (cls.session_output_budget - session.bytes_read) <= 0: + break + return chunks, matched + + @classmethod + def read( + cls, + session_id: int, + *, + max_bytes: int = DEFAULT_PTY_READ_BYTES, + timeout: float = DEFAULT_PTY_POLL_TIMEOUT, + until: str | None = None, + ) -> PtyReadResult: + """Read PTY output. + + Without ``until``, waits up to ``timeout`` for any data (one poll). + With ``until``, polls until the substring appears, the process dies, + ``timeout`` deadline elapses, or the session output budget is hit. + """ + if max_bytes < 1: + msg = "max_bytes must be >= 1" + raise WorkspaceError(msg) + if timeout < 0: + msg = "timeout must be >= 0" + raise WorkspaceError(msg) + + session = cls.refresh(session_id) + if session.closed: + msg = f"closed PTY session: {session_id}" + raise WorkspaceError(msg) + + if (cls.session_output_budget - session.bytes_read) <= 0: + return PtyReadResult( + session_id=session_id, + alive=session.alive, + exit_code=session.exit_code, + data="", + budget_exhausted=True, + ) + + chunks, matched = cls._collect_chunks( + session, max_bytes=max_bytes, timeout=timeout, until=until + ) + data = b"".join(chunks).decode(errors="replace") + truncated = len(data) > max_bytes + if truncated: + data = data[:max_bytes] + return PtyReadResult( + session_id=session_id, + alive=session.alive, + exit_code=session.exit_code, + data=data, + truncated=truncated, + matched=matched, + budget_exhausted=(cls.session_output_budget - session.bytes_read) <= 0, + ) @classmethod def write(cls, session_id: int, data: str) -> None: - session = cls.get(session_id) - if session is None or session.closed: - msg = f"unknown or closed PTY session: {session_id}" + session = cls.refresh(session_id) + if session.closed: + msg = f"closed PTY session: {session_id}" raise WorkspaceError(msg) - _ = os.write(session.master_fd, data.encode()) + if not session.alive: + msg = f"PTY session process is not alive: {session_id}" + raise WorkspaceError(msg) + try: + _ = os.write(session.master_fd, data.encode()) + except OSError as exc: + cls._poll_exit(session) + msg = f"failed to write PTY: {exc}" + raise WorkspaceError(msg) from exc + cls._touch(session) @classmethod - def close(cls, session_id: int) -> None: + def close(cls, session_id: int, *, grace_seconds: float = DEFAULT_CLOSE_GRACE_SECONDS) -> PtyCloseResult: with cls._lock: - session = cls._sessions.pop(session_id, None) + session = cls._sessions.get(session_id) if session is None: - return - session.closed = True - with contextlib.suppress(ProcessLookupError): - os.kill(session.pid, signal.SIGTERM) - with contextlib.suppress(ChildProcessError): - _ = os.waitpid(session.pid, 0) + return PtyCloseResult( + session_id=session_id, + closed=False, + alive=False, + exit_code=None, + message="unknown session", + ) + if session.closed: + return PtyCloseResult( + session_id=session_id, + closed=True, + alive=False, + exit_code=session.exit_code, + message="already closed", + ) + + cls._poll_exit(session) + if session.alive: + with contextlib.suppress(ProcessLookupError): + os.kill(session.pid, signal.SIGTERM) + deadline = time.monotonic() + max(0.0, grace_seconds) + while session.alive and time.monotonic() < deadline: + time.sleep(0.05) + cls._poll_exit(session) + if session.alive: + with contextlib.suppress(ProcessLookupError): + os.kill(session.pid, signal.SIGKILL) + with contextlib.suppress(ChildProcessError): + _ = os.waitpid(session.pid, 0) + session.alive = False + if session.exit_code is None: + session.exit_code = -signal.SIGKILL + else: + # Ensure reaped + with contextlib.suppress(ChildProcessError): + _ = os.waitpid(session.pid, os.WNOHANG) + with contextlib.suppress(OSError): os.close(session.master_fd) + session.closed = True + with cls._lock: + _ = cls._sessions.pop(session_id, None) + + return PtyCloseResult( + session_id=session_id, + closed=True, + alive=False, + exit_code=session.exit_code, + message="closed", + ) + + @classmethod + def reap_idle(cls) -> list[int]: + """Close sessions idle longer than ``idle_ttl_seconds`` (0 disables).""" + if cls.idle_ttl_seconds <= 0: + return [] + now = time.time() + to_close: list[int] = [] + with cls._lock: + for sid, session in list(cls._sessions.items()): + if session.closed: + continue + if now - session.last_activity >= cls.idle_ttl_seconds: + to_close.append(sid) + for sid in to_close: + _ = cls.close(sid) + return to_close @classmethod def close_all(cls) -> None: with cls._lock: ids = list(cls._sessions.keys()) for session_id in ids: - cls.close(session_id) + _ = cls.close(session_id) + + +def format_read_result(result: PtyReadResult) -> str: + exit_disp = "" if result.exit_code is None else str(result.exit_code) + lines = [ + f"session_id={result.session_id}", + f"alive={'true' if result.alive else 'false'}", + f"exit_code={exit_disp}", + f"matched={'true' if result.matched else 'false'}", + f"truncated={'true' if result.truncated else 'false'}", + f"budget_exhausted={'true' if result.budget_exhausted else 'false'}", + "--- data ---", + result.data, + ] + return "\n".join(lines) + + +def format_close_result(result: PtyCloseResult) -> str: + exit_disp = "" if result.exit_code is None else str(result.exit_code) + return "\n".join( + [ + f"session_id={result.session_id}", + f"closed={'true' if result.closed else 'false'}", + f"alive={'true' if result.alive else 'false'}", + f"exit_code={exit_disp}", + f"message={result.message}", + ] + ) diff --git a/src/plyngent/tools/process/read_pty.py b/src/plyngent/tools/process/read_pty.py index 5f21fcd..38bdaaf 100644 --- a/src/plyngent/tools/process/read_pty.py +++ b/src/plyngent/tools/process/read_pty.py @@ -3,13 +3,34 @@ from __future__ import annotations from plyngent.agent import tool from plyngent.tools.workspace import WorkspaceError -from .pty_session import PtyManager +from .pty_session import DEFAULT_PTY_READ_BYTES, PtyManager, format_read_result @tool -def read_pty(session_id: int, *, max_bytes: int = 8192, timeout: float = 0.2) -> str: - """Read available output from a PTY session (may be empty if nothing ready).""" +def read_pty( + session_id: int, + *, + max_bytes: int = DEFAULT_PTY_READ_BYTES, + timeout: float = 2.0, + until: str | None = None, +) -> str: + """Read PTY output with status. + + Returns structured text: session_id, alive, exit_code, matched, truncated, + budget_exhausted, then ``--- data ---`` and the payload. + + Without ``until``, waits up to ``timeout`` seconds for available data. + With ``until``, polls until the substring appears, the process exits, the + deadline elapses, or the session output budget is exhausted. + Empty data with alive=true means nothing was ready (not necessarily EOF). + """ try: - return PtyManager.read(session_id, max_bytes=max_bytes, timeout=timeout) + result = PtyManager.read( + session_id, + max_bytes=max_bytes, + timeout=timeout, + until=until, + ) except WorkspaceError as exc: return f"error: {exc}" + return format_read_result(result) diff --git a/src/plyngent/tools/process/write_pty.py b/src/plyngent/tools/process/write_pty.py index acbb223..466a18c 100644 --- a/src/plyngent/tools/process/write_pty.py +++ b/src/plyngent/tools/process/write_pty.py @@ -8,11 +8,20 @@ from .pty_session import PtyManager @tool def write_pty(session_id: int, data: str) -> str: - """Write text to a PTY session (e.g. interactive input). Does not append a newline.""" + """Write text to a PTY session (interactive input). Does not append a newline.""" try: PtyManager.write(session_id, data) + session = PtyManager.refresh(session_id) except WorkspaceError as exc: return f"error: {exc}" except OSError as exc: return f"error: failed to write PTY: {exc}" - return f"wrote {len(data)} characters to session_id={session_id}" + exit_disp = "" if session.exit_code is None else str(session.exit_code) + return "\n".join( + [ + f"session_id={session_id}", + f"alive={'true' if session.alive else 'false'}", + f"exit_code={exit_disp}", + f"wrote={len(data)}", + ] + ) diff --git a/tests/test_tools/test_process.py b/tests/test_tools/test_process.py index 8f603f0..60a447f 100644 --- a/tests/test_tools/test_process.py +++ b/tests/test_tools/test_process.py @@ -9,6 +9,23 @@ from plyngent.tools.workspace import set_command_denylist from tests.test_tools.helpers import call_async, call_sync +def _session_id(opened: str) -> int: + for line in opened.splitlines(): + if line.startswith("session_id="): + return int(line.split("=", 1)[1]) + msg = f"no session_id in: {opened!r}" + raise AssertionError(msg) + + +def _field(text: str, name: str) -> str: + prefix = f"{name}=" + for line in text.splitlines(): + if line.startswith(prefix): + return line[len(prefix) :] + msg = f"missing {name} in: {text!r}" + raise AssertionError(msg) + + async def test_run_command_echo(workspace: object) -> None: del workspace out = await call_async(run_command, ["echo", "hi"]) @@ -59,12 +76,13 @@ def test_pty_open_read_close(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["sleep", "30"]) - assert opened.startswith("session_id=") - session_id = int(opened.split("=", 1)[1]) + assert "session_id=" in opened + session_id = _session_id(opened) data = call_sync(read_pty, session_id, timeout=0.05) - assert isinstance(data, str) + assert "alive=" in data + assert "--- data ---" in data closed = call_sync(close_pty, session_id) - assert "closed" in closed + assert _field(closed, "closed") == "true" assert "error" in call_sync(read_pty, session_id) finally: PtyManager.close_all() @@ -80,18 +98,12 @@ def test_pty_echo_output(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["/bin/echo", "hello-pty"]) - session_id = int(opened.split("=", 1)[1]) - chunks: list[str] = [] - for _ in range(20): - chunk = call_sync(read_pty, session_id, timeout=0.1) - if chunk: - chunks.append(chunk) - if "hello-pty" in "".join(chunks): - break - time.sleep(0.05) - text = "".join(chunks) - _ = call_sync(close_pty, session_id) + session_id = _session_id(opened) + text = call_sync(read_pty, session_id, timeout=2.0, until="hello-pty") assert "hello-pty" in text + assert _field(text, "matched") == "true" + closed = call_sync(close_pty, session_id) + assert _field(closed, "closed") == "true" finally: PtyManager.close_all() @@ -100,20 +112,12 @@ def test_write_pty(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["cat"]) - session_id = int(opened.split("=", 1)[1]) + session_id = _session_id(opened) written = call_sync(write_pty, session_id, "pty-input\n") - assert "wrote" in written - chunks: list[str] = [] - for _ in range(30): - chunk = call_sync(read_pty, session_id, timeout=0.1) - if chunk: - chunks.append(chunk) - if "pty-input" in "".join(chunks): - break - time.sleep(0.05) - text = "".join(chunks) - _ = call_sync(close_pty, session_id) + assert "wrote=" in written + text = call_sync(read_pty, session_id, timeout=2.0, until="pty-input") assert "pty-input" in text + _ = call_sync(close_pty, session_id) finally: PtyManager.close_all() @@ -121,3 +125,58 @@ def test_write_pty(workspace: object) -> None: def test_write_pty_unknown_session(workspace: object) -> None: del workspace assert "error" in call_sync(write_pty, 999_999, "x") + + +def test_pty_exec_failure_surfaces(workspace: object) -> None: + del workspace + try: + opened = call_sync(open_pty, ["definitely-not-a-real-binary-xyz"]) + session_id = _session_id(opened) + text = call_sync(read_pty, session_id, timeout=2.0) + # marker and/or dead process with 127 + assert "plyngent-pty-exec-failed" in text or _field(text, "alive") == "false" + closed = call_sync(close_pty, session_id) + # exit 127 is conventional for exec failure + exit_code = _field(closed, "exit_code") + assert exit_code in {"127", "-9", ""} or exit_code.startswith("-") + finally: + PtyManager.close_all() + + +def test_pty_session_limit(workspace: object) -> None: + del workspace + previous = PtyManager.max_sessions + try: + PtyManager.configure(max_sessions=1) + first = call_sync(open_pty, ["sleep", "30"]) + assert "session_id=" in first + second = call_sync(open_pty, ["sleep", "30"]) + assert "limit" in second + finally: + PtyManager.close_all() + PtyManager.configure(max_sessions=previous) + + +def test_pty_output_budget(workspace: object) -> None: + del workspace + previous = PtyManager.session_output_budget + try: + PtyManager.configure(session_output_budget=64) + opened = call_sync(open_pty, ["sh", "-c", "yes x | head -c 1000"]) + session_id = _session_id(opened) + # Drain until budget exhausted or process ends. + budget_hit = False + last = "" + for _ in range(20): + last = call_sync(read_pty, session_id, timeout=0.5, max_bytes=32) + if _field(last, "budget_exhausted") == "true": + budget_hit = True + break + if _field(last, "alive") == "false": + break + time.sleep(0.05) + assert budget_hit or "x" in last + _ = call_sync(close_pty, session_id) + finally: + PtyManager.close_all() + PtyManager.configure(session_output_budget=previous)