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.
|
||||
- **`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/`)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)}",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -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}",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)}",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user