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.
This commit is contained in:
2026-07-14 19:47:32 +08:00
parent 73c11cee36
commit 1b53b415b6
7 changed files with 456 additions and 69 deletions
+2 -1
View File
@@ -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. - **`workspace`**: path resolve under root; path substring denylist; command basename denylist.
- **`file`**: `read_file`, `write_file`, `listdir`, `edit_replace` (first occurrence). - **`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`. - **`DEFAULT_TOOLS`**: file + process tool list for a `ToolRegistry`.
### CLI (`cli/`) ### CLI (`cli/`)
+4 -4
View File
@@ -2,11 +2,11 @@ from __future__ import annotations
from plyngent.agent import tool from plyngent.agent import tool
from .pty_session import PtyManager from .pty_session import PtyManager, format_close_result
@tool @tool
def close_pty(session_id: int) -> str: def close_pty(session_id: int) -> str:
"""Close a PTY session and terminate its process.""" """Close a PTY session (SIGTERM, then SIGKILL after a short grace period)."""
PtyManager.close(session_id) result = PtyManager.close(session_id)
return f"closed session_id={session_id}" return format_close_result(result)
+13 -2
View File
@@ -8,11 +8,22 @@ from .pty_session import PtyManager
@tool @tool
def open_pty(command: list[str], *, cwd: str = ".") -> str: 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: try:
session = PtyManager.open(command, cwd=cwd) session = PtyManager.open(command, cwd=cwd)
except WorkspaceError as exc: except WorkspaceError as exc:
return f"error: {exc}" return f"error: {exc}"
except OSError as exc: except OSError as exc:
return f"error: failed to open PTY: {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)}",
]
)
+315 -29
View File
@@ -1,3 +1,5 @@
"""Unix PTY session registry (requires ``pty`` + ``os.fork``; not supported on Windows)."""
from __future__ import annotations from __future__ import annotations
import contextlib import contextlib
@@ -13,7 +15,13 @@ from typing import ClassVar
from plyngent.tools.workspace import WorkspaceError, check_command_allowed, resolve_path from plyngent.tools.workspace import WorkspaceError, check_command_allowed, resolve_path
DEFAULT_PTY_READ_BYTES = 8192 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 _STDERR_FD = 2
_EXEC_FAIL_MARKER = b"plyngent-pty-exec-failed: "
@dataclass @dataclass
@@ -22,15 +30,58 @@ class PtySession:
master_fd: int master_fd: int
pid: int pid: int
closed: bool = False closed: bool = False
alive: bool = True
exit_code: int | None = None
created_at: float = field(default_factory=time.time) 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: class PtyManager:
"""In-process PTY session registry (minimal).""" """In-process PTY session registry (process-global; suitable for local CLI)."""
_lock: ClassVar[Lock] = Lock() _lock: ClassVar[Lock] = Lock()
_next_id: ClassVar[int] = 1 _next_id: ClassVar[int] = 1
_sessions: ClassVar[dict[int, PtySession]] = {} _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 @classmethod
def open( def open(
@@ -45,6 +96,13 @@ class PtyManager:
msg = f"cwd is not a directory: {cwd}" msg = f"cwd is not a directory: {cwd}"
raise WorkspaceError(msg) 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() master_fd, slave_fd = pty.openpty()
pid = os.fork() pid = os.fork()
if pid == 0: # child if pid == 0: # child
@@ -58,15 +116,21 @@ class PtyManager:
os.close(slave_fd) os.close(slave_fd)
_ = os.chdir(workdir) _ = os.chdir(workdir)
os.execvp(command[0], command) os.execvp(command[0], command)
except OSError: except OSError as exc:
pass with contextlib.suppress(OSError):
_ = os.write(1, _EXEC_FAIL_MARKER + str(exc).encode(errors="replace") + b"\n")
os._exit(127) os._exit(127)
os.close(slave_fd) os.close(slave_fd)
with cls._lock: with cls._lock:
session_id = cls._next_id session_id = cls._next_id
cls._next_id += 1 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 cls._sessions[session_id] = session
return session return session
@@ -76,45 +140,267 @@ class PtyManager:
return cls._sessions.get(session_id) return cls._sessions.get(session_id)
@classmethod @classmethod
def read(cls, session_id: int, *, max_bytes: int = DEFAULT_PTY_READ_BYTES, timeout: float = 0.1) -> str: def _touch(cls, session: PtySession) -> None:
session = cls.get(session_id) session.last_activity = time.time()
if session is None or session.closed:
msg = f"unknown or closed PTY session: {session_id}" @classmethod
raise WorkspaceError(msg) def _poll_exit(cls, session: PtySession) -> None:
ready, _, _ = select.select([session.master_fd], [], [], timeout) if not session.alive or session.closed:
if not ready: return
return ""
try: try:
data = os.read(session.master_fd, max_bytes) waited_pid, status = os.waitpid(session.pid, os.WNOHANG)
except OSError: except ChildProcessError:
return "" session.alive = False
return data.decode(errors="replace") 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 @classmethod
def write(cls, session_id: int, data: str) -> None: def write(cls, session_id: int, data: str) -> None:
session = cls.get(session_id) session = cls.refresh(session_id)
if session is None or session.closed: if session.closed:
msg = f"unknown or closed PTY session: {session_id}" msg = f"closed PTY session: {session_id}"
raise WorkspaceError(msg) 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 @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: with cls._lock:
session = cls._sessions.pop(session_id, None) session = cls._sessions.get(session_id)
if session is None: if session is None:
return return PtyCloseResult(
session.closed = True session_id=session_id,
with contextlib.suppress(ProcessLookupError): closed=False,
os.kill(session.pid, signal.SIGTERM) alive=False,
with contextlib.suppress(ChildProcessError): exit_code=None,
_ = os.waitpid(session.pid, 0) 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): with contextlib.suppress(OSError):
os.close(session.master_fd) 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 @classmethod
def close_all(cls) -> None: def close_all(cls) -> None:
with cls._lock: with cls._lock:
ids = list(cls._sessions.keys()) ids = list(cls._sessions.keys())
for session_id in ids: 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}",
]
)
+25 -4
View File
@@ -3,13 +3,34 @@ from __future__ import annotations
from plyngent.agent import tool from plyngent.agent import tool
from plyngent.tools.workspace import WorkspaceError from plyngent.tools.workspace import WorkspaceError
from .pty_session import PtyManager from .pty_session import DEFAULT_PTY_READ_BYTES, PtyManager, format_read_result
@tool @tool
def read_pty(session_id: int, *, max_bytes: int = 8192, timeout: float = 0.2) -> str: def read_pty(
"""Read available output from a PTY session (may be empty if nothing ready).""" 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: 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: except WorkspaceError as exc:
return f"error: {exc}" return f"error: {exc}"
return format_read_result(result)
+11 -2
View File
@@ -8,11 +8,20 @@ from .pty_session import PtyManager
@tool @tool
def write_pty(session_id: int, data: str) -> str: 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: try:
PtyManager.write(session_id, data) PtyManager.write(session_id, data)
session = PtyManager.refresh(session_id)
except WorkspaceError as exc: except WorkspaceError as exc:
return f"error: {exc}" return f"error: {exc}"
except OSError as exc: except OSError as exc:
return f"error: failed to write PTY: {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)}",
]
)
+86 -27
View File
@@ -9,6 +9,23 @@ from plyngent.tools.workspace import set_command_denylist
from tests.test_tools.helpers import call_async, call_sync 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: async def test_run_command_echo(workspace: object) -> None:
del workspace del workspace
out = await call_async(run_command, ["echo", "hi"]) out = await call_async(run_command, ["echo", "hi"])
@@ -59,12 +76,13 @@ def test_pty_open_read_close(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, ["sleep", "30"]) opened = call_sync(open_pty, ["sleep", "30"])
assert opened.startswith("session_id=") assert "session_id=" in opened
session_id = int(opened.split("=", 1)[1]) session_id = _session_id(opened)
data = call_sync(read_pty, session_id, timeout=0.05) 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) 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) assert "error" in call_sync(read_pty, session_id)
finally: finally:
PtyManager.close_all() PtyManager.close_all()
@@ -80,18 +98,12 @@ def test_pty_echo_output(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, ["/bin/echo", "hello-pty"]) opened = call_sync(open_pty, ["/bin/echo", "hello-pty"])
session_id = int(opened.split("=", 1)[1]) session_id = _session_id(opened)
chunks: list[str] = [] text = call_sync(read_pty, session_id, timeout=2.0, until="hello-pty")
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)
assert "hello-pty" in text assert "hello-pty" in text
assert _field(text, "matched") == "true"
closed = call_sync(close_pty, session_id)
assert _field(closed, "closed") == "true"
finally: finally:
PtyManager.close_all() PtyManager.close_all()
@@ -100,20 +112,12 @@ def test_write_pty(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, ["cat"]) 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") written = call_sync(write_pty, session_id, "pty-input\n")
assert "wrote" in written assert "wrote=" in written
chunks: list[str] = [] text = call_sync(read_pty, session_id, timeout=2.0, until="pty-input")
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 "pty-input" in text assert "pty-input" in text
_ = call_sync(close_pty, session_id)
finally: finally:
PtyManager.close_all() PtyManager.close_all()
@@ -121,3 +125,58 @@ def test_write_pty(workspace: object) -> None:
def test_write_pty_unknown_session(workspace: object) -> None: def test_write_pty_unknown_session(workspace: object) -> None:
del workspace del workspace
assert "error" in call_sync(write_pty, 999_999, "x") 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)