from __future__ import annotations import asyncio import sys from pathlib import Path import pytest from plyngent.tools.process import ( ask_into_pty, close_pty, open_pty, read_pty, run_argv, run_argv_batch, write_pty, write_pty_keys, ) 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) def _py(code: str) -> list[str]: """Cross-platform argv that runs a short Python snippet.""" return [sys.executable, "-c", code] async def test_run_argv_echo(workspace: object) -> None: del workspace out = await call_async(run_argv, _py("print('hi')")) assert "exit_code=0" in out assert "hi" in out async def test_run_argv_string_command_error(workspace: object) -> None: del workspace out = await call_async(run_argv, "ls -la") assert "error:" in out assert "only accepts an array" in out assert "executable not found" not in out async def test_run_argv_stringified_json_array_error(workspace: object) -> None: del workspace out = await call_async(run_argv, '["ls", "-la"]') assert "error:" in out assert "only accepts an array" in out assert "executable not found" not in out async def test_run_argv_non_string_args_error(workspace: object) -> None: del workspace out = await call_async(run_argv, ["ls", 1]) assert "error:" in out assert "array of string args" in out async def test_open_pty_string_command_error(workspace: object) -> None: del workspace out = await call_async(open_pty, "echo hi") assert "error:" in out assert "only accepts an array" in out async def test_run_argv_batch_serial(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ {"argv": _py("print('a')")}, {"argv": _py("print('b')")}, ], ) assert "steps=2" in out assert "ran=2" in out assert "stopped_early=false" in out assert "--- step 0 ---" in out assert "--- step 1 ---" in out assert "a" in out and "b" in out async def test_run_argv_batch_stop_on_error(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ {"argv": _py("import sys; sys.exit(2)")}, {"argv": _py("print('should-not-run')")}, ], stop_on_error=True, ) assert "ran=1" in out assert "stopped_early=true" in out assert "should-not-run" not in out async def test_run_argv_batch_pipe(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ { "argv": _py("print('piped-payload', end='')"), "pipe": {"argv": _py("import sys; print(sys.stdin.read())")}, }, ], ) assert "ran=2" in out assert "piped-payload" in out assert "pipe=true" in out async def test_run_argv_batch_pipe_out_is_gone(workspace: object) -> None: """The old provider-side flag is refused, pointing at the pipe node.""" del workspace out = await call_async( run_argv_batch, [{"argv": _py("print(1)"), "pipe_out": True}], ) assert "error:" in out assert '{"pipe": [...]}' in out async def test_run_argv_batch_mix_stderr(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ { "argv": _py("import sys; sys.stdout.write('out'); sys.stderr.write('err')"), "mix_stderr": True, "pipe": {"argv": _py("import sys; print(repr(sys.stdin.read()))")}, }, ], ) assert "ran=2" in out # Merged capture should include both (order depends on OS scheduling). assert "out" in out and "err" in out async def test_run_argv_batch_then_runs_on_success(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [{"argv": _py("print('first')"), "then": {"argv": _py("print('then-body')")}}], ) assert "then-body" in out assert "ran=2" in out assert "depth=1" in out assert "via=then of step 0" in out async def test_run_argv_batch_else_handles_the_failure(workspace: object) -> None: """A failure with an ``else`` is answered, so the batch keeps going.""" del workspace out = await call_async( run_argv_batch, [ { "argv": _py("import sys; sys.exit(2)"), "else": {"argv": _py("print('else-body')")}, }, {"argv": _py("print('sibling-ran')")}, ], ) assert "else-body" in out assert "sibling-ran" in out assert "via=else of step 0" in out assert "stopped_early=false" in out async def test_run_argv_batch_runs_a_then_list_in_order(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ { "argv": _py("print('first')"), "then": [{"argv": _py("print('second')")}, {"argv": _py("print('third')")}], } ], ) assert "ran=3" in out assert out.index("first") < out.index("second") < out.index("third") async def test_run_argv_batch_a_failing_branch_stops_the_batch(workspace: object) -> None: """Nobody answered a failure inside a continuation, so the batch stops.""" del workspace out = await call_async( run_argv_batch, [ {"argv": _py("print('first')"), "then": {"argv": _py("import sys; sys.exit(4)")}}, {"argv": _py("print('sibling-ran')")}, ], ) assert "stopped_early=true" in out assert "sibling-ran" not in out async def test_run_argv_batch_stop_on_error_false_skips_only_its_continuations(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [ { "argv": _py("import sys; sys.exit(2)"), "stop_on_error": False, "then": {"argv": _py("print('then-ran')")}, }, {"argv": _py("print('sibling-ran')")}, ], ) assert "then-ran" not in out assert "sibling-ran" in out assert "stopped_early=false" in out async def test_run_argv_batch_branches_on_the_pipe_result(workspace: object) -> None: """A piped node branches on the last command of its chain.""" del workspace out = await call_async( run_argv_batch, [ { "argv": _py("print('payload', end='')"), "pipe": {"argv": _py("import sys; sys.stdin.read(); sys.exit(3)")}, "then": {"argv": _py("print('pipe-ok')")}, "else": {"argv": _py("print('pipe-failed')")}, } ], ) assert "pipe-failed" in out assert "pipe-ok" not in out assert "ran=3" in out async def test_run_argv_batch_pipe_step_takes_no_branch(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [{"argv": _py("print(1)"), "pipe": {"argv": _py("print(2)"), "then": {"argv": _py("print(3)")}}}], ) assert "error:" in out assert "pipe step takes no then/else" in out async def test_run_argv_batch_pipe_step_takes_no_stdin(workspace: object) -> None: del workspace out = await call_async( run_argv_batch, [{"argv": _py("print(1)"), "pipe": {"argv": _py("print(2)"), "stdin": "own"}}], ) assert "error:" in out assert "stdin comes from the pipe" in out async def test_run_argv_batch_rejects_a_deep_tree(workspace: object) -> None: del workspace node: dict[str, object] = {"argv": _py("print(1)")} for _ in range(6): node = {"argv": _py("print(1)"), "then": node} out = await call_async(run_argv_batch, [node]) assert "error:" in out assert "nest at most 4 levels deep" in out async def test_run_argv_batch_counts_nested_nodes(workspace: object) -> None: """Every node of the tree costs one, so a nested tree hits the same cap.""" del workspace out = await call_async(run_argv_batch, [{"argv": _py("print(1)")} for _ in range(21)]) assert "error:" in out assert "more than 20 nodes" in out async def test_run_argv_batch_names_the_node_grammar(workspace: object) -> None: del workspace out = await call_async(run_argv_batch, [{"argv": _py("print(1)"), "when": "always"}]) assert "error:" in out assert "unknown key(s) when" in out assert "then, else, pipe" in out async def test_run_argv_denied(workspace: object) -> None: del workspace out = await call_async(run_argv, ["rm", "-rf", "."]) assert "denied" in out async def test_run_argv_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_argv, _py("import os; print('\\n'.join(os.listdir()))"), cwd="sub") assert "f.txt" in out async def test_run_argv_timeout(workspace: object) -> None: del workspace out = await call_async(run_argv, _py("import time; time.sleep(5)"), timeout_seconds=0.2) assert "timed_out=true" in out async def test_run_argv_timeout_keeps_partial_output(workspace: object) -> None: del workspace # Print then sleep past the timeout so communicate has partial stdout after kill. out = await call_async( run_argv, _py("import sys, time; sys.stdout.write('partial-out'); sys.stdout.flush(); time.sleep(5)"), timeout_seconds=0.3, ) assert "timed_out=true" in out assert "partial-out" in out async def test_run_argv_returns_the_full_stdout(workspace: object) -> None: del workspace # The result is stored whole; the request-prep pass caps it on the wire and # names its document, so nothing is cut at the tool. out = await call_async(run_argv, _py("print('x' * 40000)")) body = out.split("--- stdout ---", 1)[1].split("--- stderr ---", 1)[0].strip() assert body == "x" * 40000 assert "truncate_token=" not in out async def test_run_argv_small_stdout_no_marker(workspace: object) -> None: del workspace out = await call_async(run_argv, _py("print('tiny')")) assert "Truncated" not in out assert "truncate_token=" not in out async def test_run_argv_stdin(workspace: object) -> None: del workspace out = await call_async( run_argv, _py("import sys; print(sys.stdin.read(), end='')"), stdin="hello-stdin\n", ) assert "exit_code=0" in out assert "hello-stdin" in out async def test_run_argv_env(workspace: object) -> None: del workspace out = await call_async( run_argv, _py("import os; print(os.environ.get('PLYNGENT_TEST_VAR', ''), end='')"), env={"PLYNGENT_TEST_VAR": "from-env"}, ) assert "exit_code=0" in out assert "from-env" in out async def test_pty_open_read_close(workspace: object) -> None: del workspace try: opened = await call_async(open_pty, _py("import time; time.sleep(30)")) assert "session_id=" in opened session_id = _session_id(opened) data = await call_async(read_pty, session_id, timeout=0.05) assert "alive=" in data assert "--- data ---" in data closed = await call_async(close_pty, session_id) assert _field(closed, "closed") == "true" assert "error" in await call_async(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"]) async def test_pty_echo_output(workspace: object) -> None: del workspace try: opened = await call_async(open_pty, _py("print('hello-pty')")) session_id = _session_id(opened) text = await call_async(read_pty, session_id, timeout=2.0, until="hello-pty") assert "hello-pty" in text assert _field(text, "matched") == "true" closed = await call_async(close_pty, session_id) assert _field(closed, "closed") == "true" finally: PtyManager.close_all() async def test_write_pty(workspace: object) -> None: del workspace try: # Line-oriented echo (portable stand-in for ``cat``). opened = await call_async( open_pty, _py( "import sys\n" "while True:\n" " line = sys.stdin.readline()\n" " if not line:\n" " break\n" " sys.stdout.write(line)\n" " sys.stdout.flush()\n" ), ) session_id = _session_id(opened) written = await call_async(write_pty, session_id, "pty-input\n") assert "wrote=" in written text = await call_async(read_pty, session_id, timeout=2.0, until="pty-input") assert "pty-input" in text _ = await call_async(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") async def test_ask_into_pty_writes_without_leaking_secret(workspace: object) -> None: """Human answer goes to PTY only; tool result must not contain the secret.""" del workspace from plyngent.prompting import temporary_backend from tests.test_prompting import ScriptedBackend secret = "super-secret-password-xyz" backend = ScriptedBackend([], secrets=[secret]) try: opened = await call_async( open_pty, _py( "import sys\n" "while True:\n" " line = sys.stdin.readline()\n" " if not line:\n" " break\n" " sys.stdout.write(line)\n" " sys.stdout.flush()\n" ), ) session_id = _session_id(opened) with temporary_backend(backend): result = await call_async( ask_into_pty, session_id, "Enter password", secret=True, ) assert "error:" not in result assert "wrote=true" in result assert "secret=true" in result assert "source=human" in result assert secret not in result # Length of secret must not appear as wrote=N either. assert f"wrote={len(secret)}" not in result assert f"wrote={len(secret) + 1}" not in result echoed = await call_async(read_pty, session_id, timeout=2.0, until=secret) assert secret in echoed closed = await call_async(close_pty, session_id) assert "session_id=" in closed finally: PtyManager.close_all() async def test_ask_into_pty_empty_cancels(workspace: object) -> None: del workspace from plyngent.prompting import temporary_backend from tests.test_prompting import ScriptedBackend backend = ScriptedBackend([], secrets=[""]) try: opened = await call_async(open_pty, _py("import time; time.sleep(30)")) session_id = _session_id(opened) with temporary_backend(backend): result = await call_async(ask_into_pty, session_id, "Pass?", secret=True) assert "error:" in result assert "empty" in result.lower() or "cancel" in result.lower() finally: PtyManager.close_all() async def test_ask_into_pty_nonsecret_line(workspace: object) -> None: del workspace from plyngent.prompting import temporary_backend from tests.test_prompting import ScriptedBackend backend = ScriptedBackend(["yes"]) try: opened = await call_async( open_pty, _py("import sys\nline = sys.stdin.readline()\nsys.stdout.write(line)\nsys.stdout.flush()\n"), ) session_id = _session_id(opened) with temporary_backend(backend): result = await call_async(ask_into_pty, session_id, "Continue?", secret=False) assert "wrote=true" in result assert "secret=false" in result assert "yes" not in result # answer not in tool result even for non-secret text = await call_async(read_pty, session_id, timeout=2.0, until="yes") assert "yes" in text finally: PtyManager.close_all() async def test_pty_exec_failure_surfaces(workspace: object) -> None: del workspace try: opened = await call_async(open_pty, ["definitely-not-a-real-binary-xyz"]) # Windows ConPTY may fail at open; POSIX may open then fail on exec. if "error:" in opened and "session_id=" not in opened: assert "not found" in opened.lower() or "failed" in opened.lower() return session_id = _session_id(opened) text = await call_async(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 = await call_async(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, _py("import time; time.sleep(30)")) assert "session_id=" in first second = call_sync(open_pty, _py("import time; time.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, _py("import time; time.sleep(30)")) second = call_sync(open_pty, _py("import time; time.sleep(30)")) assert "session_id=" in first assert "session_id=" in second assert PtyManager.max_sessions >= 2 finally: PtyManager.close_all() PtyManager.configure(max_sessions=previous) PtyManager.set_limit_continue_hook(None) async def test_pty_output_budget(workspace: object) -> None: del workspace previous = PtyManager.session_output_budget try: PtyManager.set_limit_continue_hook(None) PtyManager.configure(session_output_budget=64) opened = await call_async(open_pty, _py("print('x' * 1000)")) session_id = _session_id(opened) # Drain until budget exhausted or process ends. budget_hit = False last = "" for _ in range(20): last = await call_async(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 await asyncio.sleep(0.05) assert budget_hit or "x" in last _ = await call_async(close_pty, session_id) finally: PtyManager.close_all() PtyManager.configure(session_output_budget=previous) PtyManager.set_limit_continue_hook(None) async def test_pty_output_budget_is_per_session(workspace: object) -> None: del workspace previous = PtyManager.session_output_budget try: # configure clamps budget to >= 1024 PtyManager.configure(session_output_budget=1024) class_budget = PtyManager.session_output_budget PtyManager.set_limit_continue_hook(lambda _reason: True) opened = await call_async(open_pty, _py("print('x' * 200)")) session_id = _session_id(opened) session = PtyManager.get(session_id) assert session is not None before = session.output_budget # Force budget exhaustion path by setting bytes_read high. session.bytes_read = session.output_budget _ = await call_async(read_pty, session_id, timeout=0.1) session2 = PtyManager.get(session_id) assert session2 is not None assert session2.output_budget > before # Raising is per-session; class default for new sessions stays put. assert PtyManager.session_output_budget == class_budget _ = await call_async(close_pty, session_id) finally: PtyManager.close_all() PtyManager.configure(session_output_budget=previous) PtyManager.set_limit_continue_hook(None) def test_pty_master_not_inheritable(workspace: object) -> None: del workspace import os if sys.platform == "win32": import pytest pytest.skip("master FD inheritance is POSIX-only") try: opened = call_sync(open_pty, _py("import time; time.sleep(5)")) session_id = _session_id(opened) session = PtyManager.get(session_id) assert session is not None assert session.master_fd is not None assert os.get_inheritable(session.master_fd) is False finally: 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(r"\e") == "\x1b" assert decode_write_data(r"a\nb") == "a\nb" assert decode_write_data("plain") == "plain" # Markers: names, modifiers, repeats, and a literal `<`. assert decode_write_data("") == "\x1b" assert decode_write_data("") == "\x1b" assert decode_write_data("") == "\r" assert decode_write_data("") == "\x1b[Z" assert decode_write_data("") == "\x18" assert decode_write_data("") == "\x0f" assert decode_write_data("") == "\x1b" assert decode_write_data("") == "\x1bb" assert decode_write_data("") == "\x1b\t" assert decode_write_data("") == "\x1b[A" * 3 assert decode_write_data("<") == "" assert decode_write_data("") == "\x1b[H" assert decode_write_data("") == "\x1b[2~" assert decode_write_data("") == "\x7f" assert decode_write_data("") == "\x1b[5~" assert decode_write_data("") == "\x1b[15~" # A modifier on a key that has one: xterm's 1 + shift/alt/ctrl parameter. assert decode_write_data("") == "\x1b[1;2A" assert decode_write_data("") == "\x1b[1;3A" assert decode_write_data("") == "\x1b[1;5A" assert decode_write_data("") == "\x1b[1;3D" assert decode_write_data("") == "\x1b[6;2~" assert decode_write_data("") == "\x1b[3;5~" assert decode_write_data("") == "\x1b[15;6~" assert decode_write_data("") == "\x1b[1;7H" assert decode_write_data("") == "\x00" # A `<` that does not open a marker stays literal, spaces included. assert decode_write_data("sort < file") == "sort < file" assert decode_write_data("a < b > c") == "a < b > c" def test_decode_write_data_refuses_an_unknown_key() -> None: """A typo must be an error, not text typed into the terminal.""" from plyngent.tools.process.pty_terminal import KeySyntaxError, decode_write_data for data in ("", "", "", "", "", "", "", ""): with pytest.raises(KeySyntaxError): _ = decode_write_data(data) async def test_write_pty_keys_sends_modified_keys(workspace: object) -> None: """`` leaves the session as the bytes a terminal would send.""" del workspace try: opened = await call_async( open_pty, _py( "import sys, tty\n" "tty.setraw(0)\n" "sys.stdout.write('READY\\n')\n" "sys.stdout.flush()\n" "data = sys.stdin.buffer.read(9)\n" "sys.stdout.write('GOT ' + repr(data))\n" "sys.stdout.flush()\n" ), ) session_id = _session_id(opened) # Raw mode first: a cooked PTY would hold the keys until a newline. _ = await call_async(read_pty, session_id, timeout=5.0, until="READY") written = await call_async(write_pty_keys, session_id, "") assert "wrote=9" in written text = await call_async(read_pty, session_id, timeout=2.0, until="GOT ") assert "GOT " + repr(b"\x1b[1;5A\x1b[Z") in text _ = await call_async(close_pty, session_id) finally: PtyManager.close_all() def test_sanitize_pty_output_escapes_csi() -> None: from plyngent.tools.process.pty_terminal import sanitize_pty_output_for_tool raw = chr(0x1B) + "[?1049hhello" + chr(0x1B) + "[0m" + chr(7) safe = sanitize_pty_output_for_tool(raw) assert chr(0x1B) not in safe assert "\\x1b" in safe assert "hello" in safe assert "\\x07" in safe def test_close_all_empty_is_noop() -> None: """Ctrl-D / chat exit with no sessions must not touch the host TTY.""" PtyManager.close_all() async def test_read_pty_sanitizes_esc(workspace: object) -> None: del workspace try: opened = await call_async(open_pty, _py("print(chr(0x1B)+'[31mred'+chr(0x1B)+'[0m')")) session_id = _session_id(opened) text = await call_async(read_pty, session_id, timeout=2.0, until="red") assert "red" in text payload = text.split("--- data ---", 1)[-1] assert chr(0x1B) 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_keys_ctrl_escape(workspace: object) -> None: del workspace try: opened = await call_async( 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 = await call_async(write_pty_keys, session_id, "") assert "wrote=1" in written text = await call_async(read_pty, session_id, timeout=2.0) assert "error" not in text.lower() or "alive=" in text _ = await call_async(close_pty, session_id) finally: PtyManager.close_all() async def test_write_pty_keys_refuses_an_unknown_marker(workspace: object) -> None: """The typo comes back as an error and never reaches the session.""" del workspace try: opened = await call_async(open_pty, _py("print('up')")) session_id = _session_id(opened) result = await call_async(write_pty_keys, session_id, "") assert result.startswith("error: unknown key ") _ = await call_async(close_pty, session_id) finally: PtyManager.close_all() def test_write_pty_is_literal() -> None: """Plain write_pty must not call the keys decoder.""" import inspect from plyngent.tools.process.pty_terminal import decode_write_data assert decode_write_data("press to cancel") == "press " + chr(3) + " to cancel" src = inspect.getsource(write_pty.handler) assert "decode_write_data" not in src doc = write_pty.description or write_pty.handler.__doc__ or "" assert "literal" in doc.lower() def test_pty_backend_available() -> None: from plyngent.tools.process.pty_backend import pty_available assert pty_available() is True