diff --git a/CLAUDE.md b/CLAUDE.md index 68d51ee..2596246 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -63,7 +63,7 @@ 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); minimal PTY `open_pty` / `read_pty` / `close_pty`. +- **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` / `close_pty`. - **`DEFAULT_TOOLS`**: file + process tool list for a `ToolRegistry`. ### CLI (`cli/`) diff --git a/src/plyngent/tools/__init__.py b/src/plyngent/tools/__init__.py index 7c21dd0..a1886ea 100644 --- a/src/plyngent/tools/__init__.py +++ b/src/plyngent/tools/__init__.py @@ -8,6 +8,7 @@ from .process import close_pty as close_pty from .process import open_pty as open_pty from .process import read_pty as read_pty from .process import run_command as run_command +from .process import write_pty as write_pty from .workspace import ( DEFAULT_COMMAND_DENYLIST as DEFAULT_COMMAND_DENYLIST, ) diff --git a/src/plyngent/tools/process/__init__.py b/src/plyngent/tools/process/__init__.py index 73b8bef..fe69722 100644 --- a/src/plyngent/tools/process/__init__.py +++ b/src/plyngent/tools/process/__init__.py @@ -2,5 +2,6 @@ from .close_pty import close_pty as close_pty from .open_pty import open_pty as open_pty 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 -PROCESS_TOOLS = [run_command, open_pty, read_pty, close_pty] +PROCESS_TOOLS = [run_command, open_pty, read_pty, write_pty, close_pty] diff --git a/src/plyngent/tools/process/run_command.py b/src/plyngent/tools/process/run_command.py index d0d94ba..4f791e7 100644 --- a/src/plyngent/tools/process/run_command.py +++ b/src/plyngent/tools/process/run_command.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import os import shlex +from typing import TYPE_CHECKING from plyngent.agent import tool from plyngent.tools.workspace import ( @@ -11,6 +13,9 @@ from plyngent.tools.workspace import ( resolve_path, ) +if TYPE_CHECKING: + from pathlib import Path + DEFAULT_TIMEOUT_SECONDS = 30 DEFAULT_MAX_OUTPUT_CHARS = 32_000 @@ -41,16 +46,30 @@ def _format_result( return "\n".join(parts) +def _merge_env(overrides: dict[str, str] | None) -> dict[str, str] | None: + if overrides is None: + return None + merged = dict(os.environ) + for key, value in overrides.items(): + merged[str(key)] = str(value) + return merged + + async def _run_exec( command: list[str], *, workdir: str, timeout_seconds: float, + stdin_data: bytes | None, + env: dict[str, str] | None, ) -> tuple[int | None, str, str] | str: + stdin = asyncio.subprocess.PIPE if stdin_data is not None else None try: proc = await asyncio.create_subprocess_exec( *command, cwd=workdir, + env=env, + stdin=stdin, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) @@ -60,7 +79,10 @@ async def _run_exec( return f"error: failed to start command: {exc}" try: - stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout=timeout_seconds) + stdout_b, stderr_b = await asyncio.wait_for( + proc.communicate(input=stdin_data), + timeout=timeout_seconds, + ) except TimeoutError: proc.kill() _ = await proc.communicate() @@ -71,17 +93,14 @@ async def _run_exec( return proc.returncode, stdout, stderr -@tool -async def run_command( +def _validate_run_args( command: list[str], *, - cwd: str = ".", - timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, -) -> str: - """Run a command without a shell (argv list) under the workspace. - - ``cwd`` is relative to or under the workspace root. Output is truncated. - """ + cwd: str, + timeout_seconds: float, + env: dict[str, str] | None, +) -> Path | str: + """Return resolved workdir, or an error string.""" if not command: return "error: command must not be empty" try: @@ -93,8 +112,37 @@ async def run_command( return f"error: cwd is not a directory: {cwd}" if timeout_seconds <= 0: return "error: timeout_seconds must be > 0" + _ = env # typed as dict[str, str] | None; runtime JSON already constrained by tool schema + return workdir - result = await _run_exec(command, workdir=str(workdir), timeout_seconds=timeout_seconds) + +@tool +async def run_command( + command: list[str], + *, + cwd: str = ".", + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + stdin: str | None = None, + env: dict[str, str] | None = None, +) -> str: + """Run a command without a shell (argv list) under the workspace. + + ``cwd`` is relative to or under the workspace root. Optional ``stdin`` is + written to the process stdin. Optional ``env`` overlays process environment + variables (merged with the current environment). Output is truncated. + """ + workdir = _validate_run_args(command, cwd=cwd, timeout_seconds=timeout_seconds, env=env) + if isinstance(workdir, str): + return workdir + + stdin_data = None if stdin is None else stdin.encode() + result = await _run_exec( + command, + workdir=str(workdir), + timeout_seconds=timeout_seconds, + stdin_data=stdin_data, + env=_merge_env(env), + ) if isinstance(result, str): return result returncode, stdout, stderr = result diff --git a/src/plyngent/tools/process/write_pty.py b/src/plyngent/tools/process/write_pty.py new file mode 100644 index 0000000..acbb223 --- /dev/null +++ b/src/plyngent/tools/process/write_pty.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from plyngent.agent import tool +from plyngent.tools.workspace import WorkspaceError + +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.""" + try: + PtyManager.write(session_id, data) + 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}" diff --git a/tests/test_tools/test_process.py b/tests/test_tools/test_process.py index d7d80d8..8f603f0 100644 --- a/tests/test_tools/test_process.py +++ b/tests/test_tools/test_process.py @@ -3,7 +3,7 @@ from __future__ import annotations import time from pathlib import Path -from plyngent.tools.process import close_pty, open_pty, read_pty, run_command +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 @@ -37,6 +37,24 @@ async def test_run_command_timeout(workspace: object) -> None: 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: @@ -76,3 +94,30 @@ def test_pty_echo_output(workspace: object) -> None: assert "hello-pty" in text finally: PtyManager.close_all() + + +def test_write_pty(workspace: object) -> None: + del workspace + try: + opened = call_sync(open_pty, ["cat"]) + session_id = int(opened.split("=", 1)[1]) + 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 "pty-input" in text + 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")