mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 22:44:59 +08:00
core/tools/process: run_command stdin/env and write_pty
Optional stdin and env overlay for run_command; write_pty for interactive PTY input.
This commit is contained in:
@@ -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.
|
- **`workspace`**: path resolve under root; path substring denylist; command basename denylist.
|
||||||
- **`file`**: `read_file`, `write_file`, `listdir`, `edit_replace` (first occurrence).
|
- **`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`.
|
- **`DEFAULT_TOOLS`**: file + process tool list for a `ToolRegistry`.
|
||||||
|
|
||||||
### CLI (`cli/`)
|
### CLI (`cli/`)
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ from .process import close_pty as close_pty
|
|||||||
from .process import open_pty as open_pty
|
from .process import open_pty as open_pty
|
||||||
from .process import read_pty as read_pty
|
from .process import read_pty as read_pty
|
||||||
from .process import run_command as run_command
|
from .process import run_command as run_command
|
||||||
|
from .process import write_pty as write_pty
|
||||||
from .workspace import (
|
from .workspace import (
|
||||||
DEFAULT_COMMAND_DENYLIST as DEFAULT_COMMAND_DENYLIST,
|
DEFAULT_COMMAND_DENYLIST as DEFAULT_COMMAND_DENYLIST,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -2,5 +2,6 @@ from .close_pty import close_pty as close_pty
|
|||||||
from .open_pty import open_pty as open_pty
|
from .open_pty import open_pty as open_pty
|
||||||
from .read_pty import read_pty as read_pty
|
from .read_pty import read_pty as read_pty
|
||||||
from .run_command import run_command as run_command
|
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]
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import os
|
||||||
import shlex
|
import shlex
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from plyngent.agent import tool
|
from plyngent.agent import tool
|
||||||
from plyngent.tools.workspace import (
|
from plyngent.tools.workspace import (
|
||||||
@@ -11,6 +13,9 @@ from plyngent.tools.workspace import (
|
|||||||
resolve_path,
|
resolve_path,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
DEFAULT_TIMEOUT_SECONDS = 30
|
DEFAULT_TIMEOUT_SECONDS = 30
|
||||||
DEFAULT_MAX_OUTPUT_CHARS = 32_000
|
DEFAULT_MAX_OUTPUT_CHARS = 32_000
|
||||||
|
|
||||||
@@ -41,16 +46,30 @@ def _format_result(
|
|||||||
return "\n".join(parts)
|
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(
|
async def _run_exec(
|
||||||
command: list[str],
|
command: list[str],
|
||||||
*,
|
*,
|
||||||
workdir: str,
|
workdir: str,
|
||||||
timeout_seconds: float,
|
timeout_seconds: float,
|
||||||
|
stdin_data: bytes | None,
|
||||||
|
env: dict[str, str] | None,
|
||||||
) -> tuple[int | None, str, str] | str:
|
) -> tuple[int | None, str, str] | str:
|
||||||
|
stdin = asyncio.subprocess.PIPE if stdin_data is not None else None
|
||||||
try:
|
try:
|
||||||
proc = await asyncio.create_subprocess_exec(
|
proc = await asyncio.create_subprocess_exec(
|
||||||
*command,
|
*command,
|
||||||
cwd=workdir,
|
cwd=workdir,
|
||||||
|
env=env,
|
||||||
|
stdin=stdin,
|
||||||
stdout=asyncio.subprocess.PIPE,
|
stdout=asyncio.subprocess.PIPE,
|
||||||
stderr=asyncio.subprocess.PIPE,
|
stderr=asyncio.subprocess.PIPE,
|
||||||
)
|
)
|
||||||
@@ -60,7 +79,10 @@ async def _run_exec(
|
|||||||
return f"error: failed to start command: {exc}"
|
return f"error: failed to start command: {exc}"
|
||||||
|
|
||||||
try:
|
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:
|
except TimeoutError:
|
||||||
proc.kill()
|
proc.kill()
|
||||||
_ = await proc.communicate()
|
_ = await proc.communicate()
|
||||||
@@ -71,17 +93,14 @@ async def _run_exec(
|
|||||||
return proc.returncode, stdout, stderr
|
return proc.returncode, stdout, stderr
|
||||||
|
|
||||||
|
|
||||||
@tool
|
def _validate_run_args(
|
||||||
async def run_command(
|
|
||||||
command: list[str],
|
command: list[str],
|
||||||
*,
|
*,
|
||||||
cwd: str = ".",
|
cwd: str,
|
||||||
timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
|
timeout_seconds: float,
|
||||||
) -> str:
|
env: dict[str, str] | None,
|
||||||
"""Run a command without a shell (argv list) under the workspace.
|
) -> Path | str:
|
||||||
|
"""Return resolved workdir, or an error string."""
|
||||||
``cwd`` is relative to or under the workspace root. Output is truncated.
|
|
||||||
"""
|
|
||||||
if not command:
|
if not command:
|
||||||
return "error: command must not be empty"
|
return "error: command must not be empty"
|
||||||
try:
|
try:
|
||||||
@@ -93,8 +112,37 @@ async def run_command(
|
|||||||
return f"error: cwd is not a directory: {cwd}"
|
return f"error: cwd is not a directory: {cwd}"
|
||||||
if timeout_seconds <= 0:
|
if timeout_seconds <= 0:
|
||||||
return "error: timeout_seconds must be > 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):
|
if isinstance(result, str):
|
||||||
return result
|
return result
|
||||||
returncode, stdout, stderr = result
|
returncode, stdout, stderr = result
|
||||||
|
|||||||
@@ -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}"
|
||||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
|||||||
import time
|
import time
|
||||||
from pathlib import Path
|
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.process.pty_session import PtyManager
|
||||||
from plyngent.tools.workspace import set_command_denylist
|
from plyngent.tools.workspace import set_command_denylist
|
||||||
from tests.test_tools.helpers import call_async, call_sync
|
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
|
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:
|
def test_pty_open_read_close(workspace: object) -> None:
|
||||||
del workspace
|
del workspace
|
||||||
try:
|
try:
|
||||||
@@ -76,3 +94,30 @@ def test_pty_echo_output(workspace: object) -> None:
|
|||||||
assert "hello-pty" in text
|
assert "hello-pty" in text
|
||||||
finally:
|
finally:
|
||||||
PtyManager.close_all()
|
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")
|
||||||
|
|||||||
Reference in New Issue
Block a user