diff --git a/CLAUDE.md b/CLAUDE.md index 89bce3e..5350ae7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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. diff --git a/src/plyngent/cli/app.py b/src/plyngent/cli/app.py index 74ea553..18a7e71 100644 --- a/src/plyngent/cli/app.py +++ b/src/plyngent/cli/app.py @@ -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: diff --git a/src/plyngent/tools/__init__.py b/src/plyngent/tools/__init__.py index 7c740a8..86e746b 100644 --- a/src/plyngent/tools/__init__.py +++ b/src/plyngent/tools/__init__.py @@ -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 diff --git a/src/plyngent/tools/file/__init__.py b/src/plyngent/tools/file/__init__.py index 427a5c8..cc46280 100644 --- a/src/plyngent/tools/file/__init__.py +++ b/src/plyngent/tools/file/__init__.py @@ -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, ] diff --git a/src/plyngent/tools/temp_workspace.py b/src/plyngent/tools/temp_workspace.py new file mode 100644 index 0000000..5a78323 --- /dev/null +++ b/src/plyngent/tools/temp_workspace.py @@ -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" + ) diff --git a/src/plyngent/tools/workspace.py b/src/plyngent/tools/workspace.py index b952faf..8172c71 100644 --- a/src/plyngent/tools/workspace.py +++ b/src/plyngent/tools/workspace.py @@ -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: diff --git a/tests/test_tools/conftest.py b/tests/test_tools/conftest.py index b3ba328..6c92e27 100644 --- a/tests/test_tools/conftest.py +++ b/tests/test_tools/conftest.py @@ -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() diff --git a/tests/test_tools/test_temp_workspace.py b/tests/test_tools/test_temp_workspace.py new file mode 100644 index 0000000..a9dc310 --- /dev/null +++ b/tests/test_tools/test_temp_workspace.py @@ -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()