diff --git a/CLAUDE.md b/CLAUDE.md index 6842d94..47bc8b4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -81,7 +81,7 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use. - **`workspace`**: path resolve under root; path substring denylist; command basename denylist; clearer deny messages. - **`file`**: `read_file`, `write_file`, `listdir`, `tree` (VCS + default noise dirs + optional `skip_dirs` / denylist walk), `glob_paths`, `grep_files` (regex, skip VCS/binary), `edit_replace`, `edit_lineno` (1-based range), `copy_path` / `move_path` / `delete_path`. - **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` / `close_pty` (POSIX: `pty`+`fork`; Windows: ConPTY via `pywinpty` env marker dep). -- PTY: backend in `pty_backend.py`; structured status (`alive`/`exit_code`/`data`); `read_pty(..., until=)`; session limit/idle TTL/output budget; close terminate→kill. +- PTY: backend in `pty_backend.py`; structured status (`alive`/`exit_code`/`data`); `read_pty(..., until=)` (CSI sanitized for host safety); `write_pty` key escapes (`\xHH`, `ctrl+x`, `key=esc`); session limit/idle TTL/output budget; close terminate→kill + **host TTY restore**. - CLI limit hooks: interactive confirm to raise tool-loop rounds, PTY session cap, or PTY output budget. - Destructive confirms: `classify_danger` + `ToolRegistry(on_confirm=…)`; CLI default deny; config `confirm_destructive` / `path_denylist`. Session YOLO: `/yolo on|off|once` and `--yes` (skip soft confirms; hard denylists unchanged; `once` expires after the next user turn). - **`vcs`**: read-only VCS tools (`vcs_kind` / `vcs_status` / `vcs_diff` / `vcs_log` / `vcs_branch`) via `VcsBackend` protocol; **git** implemented; detectors are pluggable for other systems. diff --git a/README.md b/README.md index 6503bb6..172211c 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,9 @@ Safety defaults: - Command basename denylist (e.g. dangerous shells/utilities). - Destructive tools (delete/move/overwrite) can require confirm (`confirm_destructive`; default deny in non-TTY). Override for the session with `/yolo on|off|once` or startup `--yes` (path/command denylists still apply). - PTY sessions: caps, idle TTL, output budget; master FD is non-inheritable; sessions closed on chat exit. + Prefer file tools over full-screen editors (`vim`/`nano`) for edits. `read_pty` escapes CSI so tool + results cannot reprogram the host TTY; `close_pty` / chat exit also restore the host terminal. + `write_pty` accepts escapes: `\xHH`, `ctrl+x`, `key=esc|enter|…`. ## Usage / context (CLI) diff --git a/src/plyngent/tools/process/__init__.py b/src/plyngent/tools/process/__init__.py index fe69722..7cf8366 100644 --- a/src/plyngent/tools/process/__init__.py +++ b/src/plyngent/tools/process/__init__.py @@ -1,5 +1,8 @@ from .close_pty import close_pty as close_pty from .open_pty import open_pty as open_pty +from .pty_terminal import decode_write_data as decode_write_data +from .pty_terminal import restore_host_terminal as restore_host_terminal +from .pty_terminal import sanitize_pty_output_for_tool as sanitize_pty_output_for_tool from .read_pty import read_pty as read_pty from .run_command import run_command as run_command from .write_pty import write_pty as write_pty diff --git a/src/plyngent/tools/process/pty_session.py b/src/plyngent/tools/process/pty_session.py index 6dc9635..27de518 100644 --- a/src/plyngent/tools/process/pty_session.py +++ b/src/plyngent/tools/process/pty_session.py @@ -8,6 +8,7 @@ from threading import Lock from typing import TYPE_CHECKING, ClassVar from plyngent.tools.process.pty_backend import PtyHandle, pty_available, spawn_pty +from plyngent.tools.process.pty_terminal import restore_host_terminal, sanitize_pty_output_for_tool from plyngent.tools.workspace import WorkspaceError, check_command_allowed, resolve_path if TYPE_CHECKING: @@ -286,6 +287,8 @@ class PtyManager: truncated = len(data) > max_bytes if truncated: data = data[:max_bytes] + # Escape CSI so tool results echoed to the host TTY cannot reprogram it. + data = sanitize_pty_output_for_tool(data) budget_exhausted = cls._maybe_raise_budget(session) return PtyReadResult( session_id=session_id, @@ -327,6 +330,7 @@ class PtyManager: message="unknown session", ) if session.closed: + restore_host_terminal() return PtyCloseResult( session_id=session_id, closed=True, @@ -354,6 +358,8 @@ class PtyManager: with cls._lock: _ = cls._sessions.pop(session_id, None) + # Host TTY may have been reprogrammed if tool results echoed CSI. + restore_host_terminal() return PtyCloseResult( session_id=session_id, closed=True, @@ -381,10 +387,14 @@ class PtyManager: @classmethod def close_all(cls) -> None: + """Close every session (each close restores the host TTY best-effort).""" with cls._lock: ids = list(cls._sessions.keys()) for session_id in ids: _ = cls.close(session_id) + # Extra restore if the registry was already empty. + if not ids: + restore_host_terminal() def format_read_result(result: PtyReadResult) -> str: diff --git a/src/plyngent/tools/process/pty_terminal.py b/src/plyngent/tools/process/pty_terminal.py new file mode 100644 index 0000000..1ad6576 --- /dev/null +++ b/src/plyngent/tools/process/pty_terminal.py @@ -0,0 +1,136 @@ +"""Host-terminal safety helpers for agent PTY sessions. + +Child PTY I/O is isolated on the master FD. The practical leak is that +``read_pty`` returns raw CSI/alt-screen sequences which the CLI then echoes +to the *user's* stdout as tool results — those bytes reprogram the host TTY. +""" + +from __future__ import annotations + +import contextlib +import re +import sys +from typing import Final + +# Leave alt-screen, restore cursor visibility, reset SGR, keypad, mouse, etc. +# Intentional over-reset: cheap and safe after TUI children. +_HOST_RESTORE: Final[str] = ( + "\x1b[?1049l" # leave alternate screen buffer + "\x1b[?47l" # leave alt screen (legacy) + "\x1b[?25h" # show cursor + "\x1b[0m" # reset SGR + "\x1b[?1l" # normal cursor keys + "\x1b[?1000l" # mouse off + "\x1b[?1002l" + "\x1b[?1003l" + "\x1b[?1006l" + "\x1b[?2004l" # bracketed paste off + "\x1b[?7h" # wraparound on + "\r\n" +) + +_HEX_BYTE = re.compile(r"\\x([0-9A-Fa-f]{2})") +_HEX_U = re.compile(r"\\u([0-9A-Fa-f]{4})") +_CTRL = re.compile(r"ctrl\+([a-z@\[\\\]\^_])", re.IGNORECASE) +_KEY = re.compile(r"key=(esc|escape|enter|tab|up|down|left|right)", re.IGNORECASE) +_SIMPLE = { + r"\n": "\n", + r"\r": "\r", + r"\t": "\t", + r"\e": "\x1b", + r"\E": "\x1b", + r"\0": "\x00", +} + +_KEY_MAP = { + "esc": "\x1b", + "escape": "\x1b", + "enter": "\r", + "tab": "\t", + "up": "\x1b[A", + "down": "\x1b[B", + "right": "\x1b[C", + "left": "\x1b[D", +} +_SIMPLE_ESCAPE_LEN = 2 + + +def restore_host_terminal() -> None: + """Best-effort reset of the *user* terminal (stdout), if it is a TTY.""" + try: + out = sys.stdout + if not out.isatty(): + return + buffer = getattr(out, "buffer", None) + if buffer is not None: + with contextlib.suppress(OSError, ValueError): + _ = buffer.write(_HOST_RESTORE.encode("ascii", errors="replace")) + _ = buffer.flush() + return + with contextlib.suppress(OSError, ValueError): + _ = out.write(_HOST_RESTORE) + _ = out.flush() + except AttributeError, OSError, ValueError: + return + + +def sanitize_pty_output_for_tool(data: str) -> str: + """Make PTY bytes safe to print on the host as tool/chat text. + + Escapes ESC (``\\\\x1b``) so CSI sequences are visible text instead of + control codes when the CLI echoes tool results. + """ + if not data or "\x1b" not in data: + return data + return data.replace("\x1b", "\\x1b") + + +def decode_write_data(data: str) -> str: + """Expand agent-friendly escapes into raw PTY input. + + Supported (as literal characters in ``data``): + + - ``\\n`` ``\\r`` ``\\t`` ``\\e`` / ``\\E`` (ESC) ``\\0`` + - ``\\xHH`` byte, ``\\uHHHH`` Unicode code point + - ``ctrl+c`` / ``ctrl+x`` ... (case-insensitive; A-Z or @ [ \\ ] ^ _) + - ``key=esc`` ``key=enter`` ``key=tab`` ``key=up|down|left|right`` + """ + if not data: + return data + + # Order: multi-char named tokens first, then hex, then simple backslash pairs. + out: list[str] = [] + i = 0 + n = len(data) + while i < n: + rest = data[i:] + m_key = _KEY.match(rest) + if m_key is not None: + out.append(_KEY_MAP[m_key.group(1).lower()]) + i += m_key.end() + continue + m_ctrl = _CTRL.match(rest) + if m_ctrl is not None: + ch = m_ctrl.group(1).upper() + out.append(chr((ord(ch) - ord("@")) & 0x1F)) + i += m_ctrl.end() + continue + m_hex = _HEX_BYTE.match(rest) + if m_hex is not None: + out.append(chr(int(m_hex.group(1), 16))) + i += m_hex.end() + continue + m_u = _HEX_U.match(rest) + if m_u is not None: + out.append(chr(int(m_u.group(1), 16))) + i += m_u.end() + continue + if rest.startswith("\\") and len(rest) >= _SIMPLE_ESCAPE_LEN: + pair = rest[:_SIMPLE_ESCAPE_LEN] + if pair in _SIMPLE: + out.append(_SIMPLE[pair]) + i += _SIMPLE_ESCAPE_LEN + continue + out.append(data[i]) + i += 1 + return "".join(out) diff --git a/src/plyngent/tools/process/write_pty.py b/src/plyngent/tools/process/write_pty.py index 466a18c..7d629dd 100644 --- a/src/plyngent/tools/process/write_pty.py +++ b/src/plyngent/tools/process/write_pty.py @@ -4,13 +4,20 @@ from plyngent.agent import tool from plyngent.tools.workspace import WorkspaceError from .pty_session import PtyManager +from .pty_terminal import decode_write_data @tool def write_pty(session_id: int, data: str) -> str: - """Write text to a PTY session (interactive input). Does not append a newline.""" + """Write text to a PTY session (interactive input). Does not append a newline. + + Escapes (literal backslash sequences in ``data``): ``\\n`` ``\\r`` ``\\t`` + ``\\e``/``\\E`` (ESC), ``\\xHH``, ``\\uHHHH``, ``ctrl+x`` / ``ctrl+c``, and + ``key=esc|enter|tab|up|down|left|right``. + """ try: - PtyManager.write(session_id, data) + raw = decode_write_data(data) + PtyManager.write(session_id, raw) session = PtyManager.refresh(session_id) except WorkspaceError as exc: return f"error: {exc}" @@ -22,6 +29,6 @@ def write_pty(session_id: int, data: str) -> str: f"session_id={session_id}", f"alive={'true' if session.alive else 'false'}", f"exit_code={exit_disp}", - f"wrote={len(data)}", + f"wrote={len(raw.encode())}", ] ) diff --git a/tests/test_tools/test_process.py b/tests/test_tools/test_process.py index 425880b..3d940e2 100644 --- a/tests/test_tools/test_process.py +++ b/tests/test_tools/test_process.py @@ -4,6 +4,8 @@ import asyncio import sys from pathlib import Path +import pytest + from plyngent.tools.process import close_pty, open_pty, read_pty, run_command, write_pty from plyngent.tools.process.pty_session import PtyManager from plyngent.tools.workspace import set_command_denylist @@ -289,6 +291,114 @@ def test_pty_master_not_inheritable(workspace: object) -> None: PtyManager.close_all() +def test_decode_write_data_escapes() -> None: + from plyngent.tools.process.pty_terminal import decode_write_data + + assert decode_write_data(r"\x0f") == "\x0f" + assert decode_write_data("ctrl+x") == "\x18" + assert decode_write_data("CTRL+O") == "\x0f" + assert decode_write_data(r"\e") == "\x1b" + assert decode_write_data("key=esc") == "\x1b" + assert decode_write_data("key=enter") == "\r" + assert decode_write_data(r"a\nb") == "a\nb" + assert decode_write_data("plain") == "plain" + + +def test_sanitize_pty_output_escapes_csi() -> None: + from plyngent.tools.process.pty_terminal import sanitize_pty_output_for_tool + + raw = "\x1b[?1049hhello\x1b[0m" + safe = sanitize_pty_output_for_tool(raw) + assert "\x1b" not in safe + assert "\\x1b" in safe + assert "hello" in safe + + +def test_restore_host_terminal_noop_when_not_tty(monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + from plyngent.tools.process.pty_terminal import restore_host_terminal + + class FakeOut: + def isatty(self) -> bool: + return False + + def write(self, _s: str) -> int: + raise AssertionError("should not write") + + monkeypatch.setattr(sys, "stdout", FakeOut()) + restore_host_terminal() # no-op + + +def test_restore_host_terminal_writes_when_tty(monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + from plyngent.tools.process.pty_terminal import restore_host_terminal + + written: list[bytes] = [] + + class Buf: + def write(self, data: bytes) -> int: + written.append(data) + return len(data) + + def flush(self) -> None: + return None + + class FakeOut: + buffer = Buf() + + def isatty(self) -> bool: + return True + + monkeypatch.setattr(sys, "stdout", FakeOut()) + restore_host_terminal() + assert written + blob = b"".join(written) + assert b"\x1b[?1049l" in blob + + +async def test_read_pty_sanitizes_esc(workspace: object) -> None: + del workspace + try: + # Print ESC so tool-facing read must escape it. + opened = call_sync(open_pty, _py("print('\\x1b[31mred\\x1b[0m')")) + session_id = _session_id(opened) + text = await call_async(read_pty, session_id, timeout=2.0, until="red") + assert "red" in text + # Raw ESC must not appear in tool payload. + payload = text.split("--- data ---", 1)[-1] + assert "\x1b" not in payload + assert "\\x1b" in payload or "red" in payload + _ = await call_async(close_pty, session_id) + finally: + PtyManager.close_all() + + +async def test_write_pty_ctrl_escape(workspace: object) -> None: + del workspace + try: + opened = call_sync( + open_pty, + _py( + "import sys\n" + "data = sys.stdin.buffer.read(1)\n" + "sys.stdout.buffer.write(data)\n" + "sys.stdout.buffer.flush()\n" + ), + ) + session_id = _session_id(opened) + written = call_sync(write_pty, session_id, "ctrl+c") + assert "wrote=1" in written + text = await call_async(read_pty, session_id, timeout=2.0) + # ESC sanitized; Ctrl+C is 0x03 — may appear as raw or escaped depending + # on printable; just ensure write succeeded and process got input. + assert "error" not in text.lower() or "alive=" in text + _ = await call_async(close_pty, session_id) + finally: + PtyManager.close_all() + + def test_pty_backend_available() -> None: from plyngent.tools.process.pty_backend import pty_available