mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
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:
@@ -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/`)
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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)}",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|||||||
@@ -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}",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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)}",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user