mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
tools: new_temporary_workspace with cross-platform allowlist
Create scratch dirs under system temp (tempfile), allow absolute paths under them without rebinding the project workspace or session. Cleanup owned temps on chat exit. Cap 16 concurrent temps; sanitize prefix.
This commit is contained in:
@@ -78,8 +78,8 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init (+ lightweight SQLite `
|
||||
|
||||
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`.
|
||||
- **`workspace`**: path resolve under primary root **or** temporary allowlist roots; path substring denylist; command basename denylist; clearer deny messages.
|
||||
- **`file`**: `read_file` (`with_lineno`), `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`, `new_temporary_workspace` (system temp allowlist; cleanup on chat exit; no session rebind).
|
||||
- **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` (literal) / `write_pty_keys` (escapes) / `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=)` (**CSI sanitized** so host TTY is never reset); keys via `write_pty_keys` only (`\xHH`, `ctrl+x`, `key=esc`); session limit/idle TTL/output budget; close terminate→kill.
|
||||
- CLI limit hooks: interactive confirm to raise tool-loop rounds, PTY session cap, or PTY output budget.
|
||||
|
||||
@@ -307,8 +307,10 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
finally:
|
||||
await memory.close()
|
||||
from plyngent.tools.process.pty_session import PtyManager
|
||||
from plyngent.tools.temp_workspace import cleanup_temporary_workspaces
|
||||
|
||||
PtyManager.close_all()
|
||||
_ = cleanup_temporary_workspaces()
|
||||
|
||||
|
||||
def _configure_logging(level: str) -> None:
|
||||
|
||||
@@ -22,6 +22,8 @@ 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 .process import write_pty_keys as write_pty_keys
|
||||
from .temp_workspace import cleanup_temporary_workspaces as cleanup_temporary_workspaces
|
||||
from .temp_workspace import new_temporary_workspace as new_temporary_workspace
|
||||
from .todo import TODO_TOOLS as TODO_TOOLS
|
||||
from .todo import get_todo_stack as get_todo_stack
|
||||
from .todo import set_todo_stack as set_todo_stack
|
||||
@@ -40,11 +42,16 @@ from .workspace import (
|
||||
DEFAULT_COMMAND_DENYLIST as DEFAULT_COMMAND_DENYLIST,
|
||||
)
|
||||
from .workspace import WorkspaceError as WorkspaceError
|
||||
from .workspace import add_workspace_allowlist as add_workspace_allowlist
|
||||
from .workspace import check_command_allowed as check_command_allowed
|
||||
from .workspace import clear_workspace_allowlist as clear_workspace_allowlist
|
||||
from .workspace import clear_workspace_root as clear_workspace_root
|
||||
from .workspace import get_command_denylist as get_command_denylist
|
||||
from .workspace import get_path_denylist as get_path_denylist
|
||||
from .workspace import get_workspace_root as get_workspace_root
|
||||
from .workspace import list_workspace_allowlist as list_workspace_allowlist
|
||||
from .workspace import pop_owned_temporary_workspaces as pop_owned_temporary_workspaces
|
||||
from .workspace import remove_workspace_allowlist as remove_workspace_allowlist
|
||||
from .workspace import resolve_path as resolve_path
|
||||
from .workspace import set_command_denylist as set_command_denylist
|
||||
from .workspace import set_path_denylist as set_path_denylist
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
from plyngent.tools.temp_workspace import new_temporary_workspace as new_temporary_workspace
|
||||
|
||||
from .edit_lineno import edit_lineno as edit_lineno
|
||||
from .edit_replace import edit_replace as edit_replace
|
||||
from .fs_ops import copy_path as copy_path
|
||||
@@ -22,4 +24,5 @@ FILE_TOOLS = [
|
||||
copy_path,
|
||||
move_path,
|
||||
delete_path,
|
||||
new_temporary_workspace,
|
||||
]
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from plyngent.agent import tool
|
||||
from plyngent.tools.workspace import (
|
||||
MAX_TEMPORARY_WORKSPACES,
|
||||
WorkspaceError,
|
||||
add_workspace_allowlist,
|
||||
list_workspace_allowlist,
|
||||
pop_owned_temporary_workspaces,
|
||||
remove_workspace_allowlist,
|
||||
)
|
||||
|
||||
_DEFAULT_PREFIX = "ws"
|
||||
_PREFIX_MAX_LEN = 32
|
||||
_PREFIX_SAFE = re.compile(rf"^[A-Za-z0-9_-]{{1,{_PREFIX_MAX_LEN}}}$")
|
||||
|
||||
|
||||
def _sanitize_prefix(prefix: str) -> str:
|
||||
token = (prefix or _DEFAULT_PREFIX).strip() or _DEFAULT_PREFIX
|
||||
# Replace unsafe characters so mkdtemp stays portable on Windows/POSIX.
|
||||
cleaned = re.sub(r"[^A-Za-z0-9_-]+", "-", token)
|
||||
cleaned = cleaned.strip("-_") or _DEFAULT_PREFIX
|
||||
if len(cleaned) > _PREFIX_MAX_LEN:
|
||||
cleaned = cleaned[:_PREFIX_MAX_LEN]
|
||||
if not _PREFIX_SAFE.match(cleaned):
|
||||
return _DEFAULT_PREFIX
|
||||
return cleaned
|
||||
|
||||
|
||||
def _is_under_system_temp(path: Path) -> bool:
|
||||
"""True if *path* is under the OS temporary directory (safe to rmtree)."""
|
||||
try:
|
||||
temp_root = Path(tempfile.gettempdir()).expanduser().resolve()
|
||||
resolved = path.expanduser().resolve()
|
||||
_ = resolved.relative_to(temp_root)
|
||||
except OSError, ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def cleanup_temporary_workspaces() -> int:
|
||||
"""Remove directories created by :func:`new_temporary_workspace` (chat exit).
|
||||
|
||||
Only deletes paths still under the system temp dir. Returns the number of
|
||||
directories removed.
|
||||
"""
|
||||
removed = 0
|
||||
for path in pop_owned_temporary_workspaces():
|
||||
if not _is_under_system_temp(path):
|
||||
remove_workspace_allowlist(path)
|
||||
continue
|
||||
if path.is_dir():
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
removed += 1
|
||||
remove_workspace_allowlist(path)
|
||||
return removed
|
||||
|
||||
|
||||
@tool
|
||||
def new_temporary_workspace(prefix: str = "ws") -> str:
|
||||
"""Create a scratch directory under the system temp dir and allow tool paths in it.
|
||||
|
||||
The project workspace is unchanged: relative paths still resolve there.
|
||||
Use the returned **absolute** path for file/process tools. Temporary
|
||||
workspaces are deleted when this chat process exits (not on each turn).
|
||||
|
||||
Cross-platform (uses ``tempfile``; typically ``/tmp`` on POSIX, ``%TEMP%``
|
||||
on Windows). At most 16 concurrent temporary workspaces per process.
|
||||
"""
|
||||
if len(list_workspace_allowlist()) >= MAX_TEMPORARY_WORKSPACES:
|
||||
return f"error: too many temporary workspaces (max {MAX_TEMPORARY_WORKSPACES})"
|
||||
|
||||
safe = _sanitize_prefix(prefix)
|
||||
try:
|
||||
# mkdtemp is cross-platform; prefix must end with - for readability.
|
||||
path_str = tempfile.mkdtemp(prefix=f"plyngent-{safe}-")
|
||||
path = Path(path_str).resolve()
|
||||
except OSError as exc:
|
||||
return f"error: failed to create temporary workspace: {exc}"
|
||||
|
||||
try:
|
||||
_ = add_workspace_allowlist(path, owned=True)
|
||||
except WorkspaceError as exc:
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
return f"error: {exc}"
|
||||
|
||||
return (
|
||||
f"temporary_workspace={path}\n"
|
||||
"note: project workspace unchanged; use this absolute path for tools; "
|
||||
"removed when chat exits"
|
||||
)
|
||||
@@ -26,6 +26,9 @@ DEFAULT_COMMAND_DENYLIST: frozenset[str] = frozenset(
|
||||
}
|
||||
)
|
||||
|
||||
# Max concurrent temporary workspaces registered in one process.
|
||||
MAX_TEMPORARY_WORKSPACES = 16
|
||||
|
||||
|
||||
class WorkspaceError(ValueError):
|
||||
"""Raised when a path or command violates workspace policy."""
|
||||
@@ -35,6 +38,14 @@ class _WorkspaceState:
|
||||
root: Path | None = None
|
||||
path_denylist: tuple[str, ...] = ()
|
||||
command_denylist: frozenset[str] = DEFAULT_COMMAND_DENYLIST
|
||||
# Extra roots allowed for resolve_path (e.g. temporary workspaces under system temp).
|
||||
allowlist: list[Path]
|
||||
# Paths created by new_temporary_workspace (subset of allowlist); cleaned on chat exit.
|
||||
temporary_owned: list[Path]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.allowlist = []
|
||||
self.temporary_owned = []
|
||||
|
||||
|
||||
_state = _WorkspaceState()
|
||||
@@ -59,7 +70,7 @@ def get_workspace_root() -> Path:
|
||||
|
||||
|
||||
def clear_workspace_root() -> None:
|
||||
"""Clear workspace root (mainly for tests)."""
|
||||
"""Clear workspace root (mainly for tests). Does not clear allowlist."""
|
||||
_state.root = None
|
||||
|
||||
|
||||
@@ -82,18 +93,83 @@ def get_command_denylist() -> frozenset[str]:
|
||||
return _state.command_denylist
|
||||
|
||||
|
||||
def add_workspace_allowlist(root: Path | str, *, owned: bool = False) -> Path:
|
||||
"""Allow tool paths under *root* in addition to the primary workspace.
|
||||
|
||||
When *owned* is true, the path is also registered for chat-exit cleanup
|
||||
(only paths created by :func:`new_temporary_workspace`).
|
||||
"""
|
||||
path = Path(root).expanduser().resolve()
|
||||
if not path.is_dir():
|
||||
msg = f"allowlist root is not a directory: {path}"
|
||||
raise WorkspaceError(msg)
|
||||
if path not in _state.allowlist:
|
||||
if len(_state.allowlist) >= MAX_TEMPORARY_WORKSPACES and owned:
|
||||
msg = f"too many temporary workspaces (max {MAX_TEMPORARY_WORKSPACES})"
|
||||
raise WorkspaceError(msg)
|
||||
_state.allowlist.append(path)
|
||||
if owned and path not in _state.temporary_owned:
|
||||
_state.temporary_owned.append(path)
|
||||
return path
|
||||
|
||||
|
||||
def list_workspace_allowlist() -> list[Path]:
|
||||
"""Return a copy of extra allowed roots (not including the primary workspace)."""
|
||||
return list(_state.allowlist)
|
||||
|
||||
|
||||
def clear_workspace_allowlist() -> None:
|
||||
"""Clear allowlist and owned-temp registry (tests). Does not delete directories."""
|
||||
_state.allowlist.clear()
|
||||
_state.temporary_owned.clear()
|
||||
|
||||
|
||||
def pop_owned_temporary_workspaces() -> list[Path]:
|
||||
"""Return and clear the owned temporary workspace list (for chat-exit cleanup).
|
||||
|
||||
Paths remain on the allowlist until the caller removes them via
|
||||
:func:`remove_workspace_allowlist`.
|
||||
"""
|
||||
owned = list(_state.temporary_owned)
|
||||
_state.temporary_owned.clear()
|
||||
return owned
|
||||
|
||||
|
||||
def remove_workspace_allowlist(root: Path | str) -> None:
|
||||
"""Drop *root* from the allowlist if present."""
|
||||
path = Path(root).expanduser().resolve()
|
||||
while path in _state.allowlist:
|
||||
_state.allowlist.remove(path)
|
||||
|
||||
|
||||
def _under_any_root(resolved: Path) -> bool:
|
||||
roots: list[Path] = []
|
||||
if _state.root is not None:
|
||||
roots.append(_state.root)
|
||||
roots.extend(_state.allowlist)
|
||||
for root in roots:
|
||||
try:
|
||||
_ = resolved.relative_to(root)
|
||||
except ValueError:
|
||||
continue
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def resolve_path(path: str | Path) -> Path:
|
||||
"""Resolve ``path`` under the workspace root; reject escapes and denylist hits."""
|
||||
"""Resolve ``path`` under the workspace root or an allowlisted temp root.
|
||||
|
||||
Relative paths resolve against the **primary** workspace root. Absolute
|
||||
paths may also land under a temporary workspace allowlist entry.
|
||||
"""
|
||||
root = get_workspace_root()
|
||||
candidate = Path(path)
|
||||
if not candidate.is_absolute():
|
||||
candidate = root / candidate
|
||||
resolved = candidate.expanduser().resolve()
|
||||
try:
|
||||
_ = resolved.relative_to(root)
|
||||
except ValueError as exc:
|
||||
if not _under_any_root(resolved):
|
||||
msg = f"path escapes workspace root ({root}): {path}"
|
||||
raise WorkspaceError(msg) from exc
|
||||
raise WorkspaceError(msg)
|
||||
# Normalize separators so denylist entries like ``/secrets/`` match on Windows.
|
||||
resolved_str = str(resolved).replace("\\", "/")
|
||||
for pattern in _state.path_denylist:
|
||||
|
||||
@@ -4,7 +4,11 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from plyngent.tools.workspace import clear_workspace_root, set_workspace_root
|
||||
from plyngent.tools.workspace import (
|
||||
clear_workspace_allowlist,
|
||||
clear_workspace_root,
|
||||
set_workspace_root,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
@@ -14,6 +18,8 @@ if TYPE_CHECKING:
|
||||
@pytest.fixture
|
||||
def workspace(tmp_path: Path) -> Iterator[Path]:
|
||||
clear_workspace_root()
|
||||
clear_workspace_allowlist()
|
||||
root = set_workspace_root(tmp_path)
|
||||
yield root
|
||||
clear_workspace_allowlist()
|
||||
clear_workspace_root()
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from plyngent.tools.file import read_file, write_file
|
||||
from plyngent.tools.temp_workspace import cleanup_temporary_workspaces, new_temporary_workspace
|
||||
from plyngent.tools.workspace import (
|
||||
WorkspaceError,
|
||||
clear_workspace_allowlist,
|
||||
list_workspace_allowlist,
|
||||
resolve_path,
|
||||
)
|
||||
from tests.test_tools.helpers import call_sync
|
||||
|
||||
|
||||
def _temp_path_from_tool_output(out: str) -> Path:
|
||||
line = next(part for part in out.splitlines() if part.startswith("temporary_workspace="))
|
||||
return Path(line.split("=", 1)[1].strip())
|
||||
|
||||
|
||||
def test_new_temporary_workspace_allowlist(workspace: object) -> None:
|
||||
assert isinstance(workspace, Path)
|
||||
out = call_sync(new_temporary_workspace, "unit")
|
||||
assert "temporary_workspace=" in out
|
||||
assert "project workspace unchanged" in out
|
||||
temp = _temp_path_from_tool_output(out)
|
||||
assert temp.is_dir()
|
||||
assert temp in list_workspace_allowlist()
|
||||
|
||||
# Absolute path under temp is allowed; project relative still works.
|
||||
target = temp / "scratch.txt"
|
||||
_ = target.write_text("hello-temp", encoding="utf-8")
|
||||
assert resolve_path(str(target)) == target.resolve()
|
||||
assert call_sync(read_file, str(target)) == "hello-temp"
|
||||
|
||||
_ = call_sync(write_file, "project.txt", "proj")
|
||||
assert call_sync(read_file, "project.txt") == "proj"
|
||||
|
||||
# Sibling under system temp that we did not allowlist still fails.
|
||||
outside = temp.parent / "not-ours-should-fail"
|
||||
with pytest.raises(WorkspaceError, match="escapes"):
|
||||
_ = resolve_path(str(outside))
|
||||
|
||||
|
||||
def test_cleanup_removes_owned_temps(workspace: object) -> None:
|
||||
assert isinstance(workspace, Path)
|
||||
out = call_sync(new_temporary_workspace)
|
||||
temp = _temp_path_from_tool_output(out)
|
||||
assert temp.is_dir()
|
||||
n = cleanup_temporary_workspaces()
|
||||
assert n >= 1
|
||||
assert not temp.exists()
|
||||
assert list_workspace_allowlist() == []
|
||||
|
||||
|
||||
def test_prefix_sanitized(workspace: object) -> None:
|
||||
del workspace
|
||||
out = call_sync(new_temporary_workspace, "bad/../x y!")
|
||||
assert "temporary_workspace=" in out
|
||||
assert not out.startswith("error")
|
||||
_ = cleanup_temporary_workspaces()
|
||||
clear_workspace_allowlist()
|
||||
Reference in New Issue
Block a user