from __future__ import annotations import time from pathlib import Path 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 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"]) assert "exit_code=0" in out assert "hi" in out async def test_run_command_denied(workspace: object) -> None: del workspace out = await call_async(run_command, ["rm", "-rf", "."]) assert "denied" in out async def test_run_command_cwd(workspace: object) -> None: assert isinstance(workspace, Path) sub = workspace / "sub" sub.mkdir() _ = (sub / "f.txt").write_text("z", encoding="utf-8") out = await call_async(run_command, ["ls"], cwd="sub") assert "f.txt" in out async def test_run_command_timeout(workspace: object) -> None: del workspace out = await call_async(run_command, ["sleep", "5"], timeout_seconds=0.2) assert "timed out" in out async def test_run_command_stdin(workspace: object) -> None: del workspace out = await call_async(run_command, ["cat"], stdin="hello-stdin\n") assert "exit_code=0" in out assert "hello-stdin" in out async def test_run_command_env(workspace: object) -> None: del workspace out = await call_async( run_command, ["sh", "-c", "printf '%s' \"$PLYNGENT_TEST_VAR\""], env={"PLYNGENT_TEST_VAR": "from-env"}, ) assert "exit_code=0" in out assert "from-env" in out def test_pty_open_read_close(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["sleep", "30"]) assert "session_id=" in opened session_id = _session_id(opened) data = call_sync(read_pty, session_id, timeout=0.05) assert "alive=" in data assert "--- data ---" in data closed = call_sync(close_pty, session_id) assert _field(closed, "closed") == "true" assert "error" in call_sync(read_pty, session_id) finally: PtyManager.close_all() set_command_denylist(None) def test_pty_denied(workspace: object) -> None: del workspace assert "denied" in call_sync(open_pty, ["sudo", "ls"]) def test_pty_echo_output(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["/bin/echo", "hello-pty"]) 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() def test_write_pty(workspace: object) -> None: del workspace try: opened = call_sync(open_pty, ["cat"]) session_id = _session_id(opened) written = call_sync(write_pty, session_id, "pty-input\n") 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() 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.set_limit_continue_hook(None) 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) PtyManager.set_limit_continue_hook(None) def test_pty_session_limit_continue(workspace: object) -> None: del workspace previous = PtyManager.max_sessions try: PtyManager.configure(max_sessions=1) PtyManager.set_limit_continue_hook(lambda _reason: True) first = call_sync(open_pty, ["sleep", "30"]) second = call_sync(open_pty, ["sleep", "30"]) assert "session_id=" in first assert "session_id=" in second assert PtyManager.max_sessions >= 2 # noqa: PLR2004 finally: PtyManager.close_all() PtyManager.configure(max_sessions=previous) PtyManager.set_limit_continue_hook(None) 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)