dscli-agent 0.1.0__py3-none-any.whl

dscli/tools.py

@@ -0,0 +1,674 @@
+from __future__ import annotations
+
+import difflib
+import locale
+import os
+import shutil
+import subprocess
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable
+
+from .protocol import CANONICAL_TOOLS, ToolCall, canonical_tool_name
+
+
+class ToolError(ValueError):
+    pass
+
+
+class PathDenied(ToolError):
+    pass
+
+
+class ApprovalRequired(RuntimeError):
+    pass
+
+
+class ToolDenied(RuntimeError):
+    pass
+
+
+ConfirmFn = Callable[[str, str | None], bool]
+
+SKIP_DIRS = {
+    ".git",
+    ".hg",
+    ".svn",
+    "node_modules",
+    "__pycache__",
+    ".venv",
+    "venv",
+    "env",
+    ".mypy_cache",
+    ".pytest_cache",
+    ".ruff_cache",
+    ".next",
+    ".nuxt",
+    ".turbo",
+    "dist",
+    "build",
+    "target",
+    "coverage",
+    ".idea",
+    ".tox",
+}
+TEXT_SKIP_EXTS = {
+    ".pyc",
+    ".pyo",
+    ".so",
+    ".dll",
+    ".exe",
+    ".bin",
+    ".zip",
+    ".gz",
+    ".7z",
+    ".tar",
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".gif",
+    ".webp",
+    ".ico",
+    ".pdf",
+    ".mp4",
+    ".mov",
+    ".mp3",
+    ".wav",
+    ".woff",
+    ".woff2",
+    ".ttf",
+}
+
+
+@dataclass
+class ToolResult:
+    name: str
+    arguments: dict[str, Any]
+    ok: bool
+    output: str
+    changed: list[str] = field(default_factory=list)
+
+
+@dataclass
+class _TextFile:
+    text: str
+    encoding: str
+    newline: str
+    had_bom: bool
+    size: int
+
+
+def _human_size(size: int) -> str:
+    value = float(size)
+    for unit in ("B", "KB", "MB"):
+        if value < 1024 or unit == "MB":
+            return f"{value:.0f} {unit}" if unit == "B" else f"{value:.1f} {unit}"
+        value /= 1024
+    return f"{value:.1f} MB"
+
+
+def _truncate(text: str, limit: int, note: str = "output truncated") -> str:
+    if len(text) <= limit:
+        return text
+    keep = max(0, limit - 80)
+    return text[:keep] + f"\n... [{note}: {len(text) - keep} more characters] ..."
+
+
+def _truncate_middle(text: str, limit: int) -> str:
+    if len(text) <= limit:
+        return text
+    half = max(0, (limit - 80) // 2)
+    return (
+        text[:half]
+        + f"\n... [{len(text) - 2 * half} characters omitted] ...\n"
+        + text[-half:]
+    )
+
+
+def _decode(raw: bytes) -> str:
+    for enc in ("utf-8-sig", "utf-8", locale.getpreferredencoding(False), "cp1252"):
+        if not enc:
+            continue
+        try:
+            return raw.decode(enc)
+        except (UnicodeDecodeError, LookupError):
+            continue
+    return raw.decode("utf-8", errors="replace")
+
+
+def _rel(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return str(path)
+
+
+def workspace_preview(root: Path, max_entries: int = 40) -> str:
+    try:
+        entries = sorted(root.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower()))
+    except OSError as exc:
+        return f"(cannot list workspace: {exc})"
+    lines: list[str] = []
+    for entry in entries:
+        if entry.name in SKIP_DIRS:
+            continue
+        if entry.is_dir():
+            lines.append(f"{entry.name}/")
+        else:
+            try:
+                size = _human_size(entry.stat().st_size)
+            except OSError:
+                size = "?"
+            lines.append(f"{entry.name} ({size})")
+        if len(lines) >= max_entries:
+            lines.append("... (more entries; use list_files)")
+            break
+    return "\n".join(lines) or "(empty workspace)"
+
+
+class Toolbox:
+    def __init__(
+        self,
+        root: Path,
+        *,
+        allow_outside: bool = False,
+        allow_shell: bool = False,
+        confirm: ConfirmFn | None = None,
+        max_read_chars: int = 40000,
+        max_write_chars: int = 1_000_000,
+        max_output_chars: int = 20000,
+        shell_timeout_sec: int = 120,
+    ) -> None:
+        self.root = root.resolve()
+        self.allow_outside = allow_outside
+        self.allow_shell = allow_shell
+        self.confirm = confirm
+        self.max_read_chars = max_read_chars
+        self.max_write_chars = max_write_chars
+        self.max_output_chars = max_output_chars
+        self.shell_timeout_sec = shell_timeout_sec
+        self._read_any: set[str] = set()
+        self._changed: list[str] = []
+
+    # ------------------------------------------------------------------ paths
+    def _resolve(self, value: Any, *, must_exist: bool = False, for_write: bool = False) -> Path:
+        if value is None or str(value).strip() == "":
+            value = "."
+        if not isinstance(value, (str, os.PathLike)):
+            raise ToolError(f"path must be a string, got {type(value).__name__}")
+        path = Path(str(value)).expanduser()
+        if not path.is_absolute():
+            path = self.root / path
+        resolved = path.resolve()
+        try:
+            inside = resolved.is_relative_to(self.root)
+        except ValueError:
+            inside = False
+        if not inside and not self.allow_outside:
+            raise PathDenied(
+                f"Path '{value}' is outside the workspace ({self.root}). "
+                "Pass --allow-outside to permit this."
+            )
+        if for_write and inside:
+            rel_parts = resolved.relative_to(self.root).parts
+            if ".git" in rel_parts:
+                raise PathDenied("Refusing to modify anything inside .git.")
+        if must_exist and not resolved.exists():
+            raise ToolError(f"Path does not exist: {value}")
+        return resolved
+
+    def _display(self, path: Path) -> str:
+        return _rel(path, self.root)
+
+    def _read_text(self, path: Path) -> _TextFile:
+        if not path.exists():
+            raise ToolError(f"File does not exist: {self._display(path)}")
+        if path.is_dir():
+            raise ToolError(f"{self._display(path)} is a directory; use list_files.")
+        raw = path.read_bytes()
+        if b"\x00" in raw[:8192]:
+            raise ToolError(
+                f"{self._display(path)} looks like a binary file and cannot be edited as text."
+            )
+        had_bom = raw.startswith(b"\xef\xbb\xbf")
+        text: str | None = None
+        encoding = "utf-8-sig" if had_bom else "utf-8"
+        try:
+            text = raw.decode(encoding)
+        except UnicodeDecodeError:
+            for fallback in ("cp1252", locale.getpreferredencoding(False) or "cp1252"):
+                try:
+                    text = raw.decode(fallback)
+                    encoding = fallback
+                    break
+                except (UnicodeDecodeError, LookupError):
+                    continue
+        if text is None:
+            raise ToolError(f"Cannot decode {self._display(path)} as text.")
+        crlf = text.count("\r\n")
+        lf = text.count("\n") - crlf
+        newline = "\r\n" if crlf > lf else "\n"
+        normalized = text.replace("\r\n", "\n").replace("\r", "\n")
+        return _TextFile(
+            text=normalized,
+            encoding=encoding,
+            newline=newline,
+            had_bom=had_bom,
+            size=len(raw),
+        )
+
+    def _write_text(self, path: Path, text: str, *, encoding: str = "utf-8", newline: str = "\n") -> None:
+        data = text
+        if newline != "\n":
+            data = data.replace("\n", newline)
+        try:
+            encoded = data.encode(encoding)
+        except UnicodeEncodeError as exc:
+            raise ToolError(
+                f"Cannot encode the new content as {encoding}; the file may use a legacy codepage."
+            ) from exc
+        path.parent.mkdir(parents=True, exist_ok=True)
+        mode = None
+        if path.exists():
+            try:
+                mode = path.stat().st_mode
+            except OSError:
+                pass
+        fd, tmp_name = tempfile.mkstemp(prefix=".dscli-", suffix=".tmp", dir=str(path.parent))
+        try:
+            with os.fdopen(fd, "wb") as fh:
+                fh.write(encoded)
+            if mode is not None:
+                try:
+                    os.chmod(tmp_name, mode)
+                except OSError:
+                    pass
+            os.replace(tmp_name, path)
+        finally:
+            try:
+                if os.path.exists(tmp_name):
+                    os.unlink(tmp_name)
+            except OSError:
+                pass
+
+    def _approve(self, title: str, detail: str | None = None) -> None:
+        if self.confirm is None:
+            raise ApprovalRequired(f"{title} needs confirmation, but no confirmer is configured.")
+        if not self.confirm(title, detail):
+            raise ToolDenied("The user denied this action.")
+
+    def _diff(self, old: str, new: str, path: Path) -> str:
+        rel = self._display(path)
+        diff = difflib.unified_diff(
+            old.splitlines(keepends=True),
+            new.splitlines(keepends=True),
+            fromfile=f"a/{rel}",
+            tofile=f"b/{rel}",
+            n=3,
+        )
+        return "".join(diff) or "(no changes)"
+
+    def _mark_written(self, path: Path) -> None:
+        key = str(path)
+        self._read_any.add(key)
+        rel = self._display(path)
+        if rel not in self._changed:
+            self._changed.append(rel)
+
+    # ------------------------------------------------------------------ tools
+    def execute(self, call: ToolCall) -> ToolResult:
+        name = canonical_tool_name(call.name)
+        args = self._normalize_args(name, call.arguments or {})
+        try:
+            if name == "read_file":
+                result = self._read_file(args)
+            elif name == "write_file":
+                result = self._write_file(args)
+            elif name == "apply_patch":
+                result = self._apply_patch(args)
+            elif name == "list_files":
+                result = self._list_files(args)
+            elif name == "delete_file":
+                result = self._delete_file(args)
+            elif name == "run_command":
+                result = self._run_command(args)
+            elif name == "finish":
+                raise ToolError("finish is handled by the harness; use the final field instead.")
+            else:
+                raise ToolError(
+                    "Unknown tool. Valid tools: " + ", ".join(sorted(CANONICAL_TOOLS)) + "."
+                )
+            result.name = name
+            result.arguments = args
+            if result.changed:
+                for rel in result.changed:
+                    if rel not in self._changed:
+                        self._changed.append(rel)
+            return result
+        except ToolDenied as exc:
+            return ToolResult(name=name, arguments=args, ok=False, output=str(exc))
+        except ApprovalRequired:
+            raise
+        except (ToolError, OSError) as exc:
+            return ToolResult(name=name, arguments=args, ok=False, output=str(exc))
+        except Exception as exc:  # keep the agent alive on unexpected tool errors
+            return ToolResult(
+                name=name,
+                arguments=args,
+                ok=False,
+                output=f"Internal tool error: {type(exc).__name__}: {exc}",
+            )
+
+    def _normalize_args(self, name: str, args: dict[str, Any]) -> dict[str, Any]:
+        args = dict(args)
+
+        def pick(*keys: str) -> Any:
+            for key in keys:
+                if key in args and args[key] is not None:
+                    return args[key]
+            return None
+
+        path = pick("path", "file_path", "filename", "file", "dir", "directory")
+        if path is None and name in {"read_file", "write_file", "apply_patch", "delete_file"} and "input" in args:
+            path = args["input"]
+        if path is not None:
+            args["path"] = path
+
+        if name == "write_file":
+            content = pick("content", "contents", "text", "body", "new_content", "data")
+            if content is not None:
+                args["content"] = content
+        if name == "apply_patch":
+            search = pick("search", "old_string", "old", "find", "target", "pattern")
+            replace = pick("replace", "new_string", "new", "replacement", "replace_with")
+            if search is not None:
+                args["search"] = search
+            if replace is not None:
+                args["replace"] = replace
+            all_flag = pick("replace_all", "all", "global")
+            if all_flag is not None:
+                if isinstance(all_flag, str):
+                    all_flag = all_flag.strip().lower() in {"1", "true", "yes", "all"}
+                args["replace_all"] = bool(all_flag)
+        if name == "run_command":
+            command = pick("command", "cmd", "script", "input")
+            if command is not None:
+                args["command"] = command
+            timeout = pick("timeout_sec", "timeout", "timeout_seconds")
+            if timeout is not None:
+                args["timeout_sec"] = timeout
+        if name == "list_files" and args.get("path") is None:
+            args["path"] = "."
+        return args
+
+    def _read_file(self, args: dict[str, Any]) -> ToolResult:
+        path = self._resolve(args.get("path"), must_exist=True)
+        tf = self._read_text(path)
+        lines = tf.text.splitlines()
+        total_lines = len(lines)
+        try:
+            start = int(args.get("start_line") or 1)
+            end = int(args.get("end_line") or 0)
+        except (TypeError, ValueError):
+            raise ToolError("start_line and end_line must be integers.")
+        if start < 1:
+            start = 1
+        if end and end < start:
+            raise ToolError("end_line must be greater than or equal to start_line.")
+        if start > total_lines and total_lines > 0:
+            raise ToolError(f"start_line {start} is past the end ({total_lines} lines).")
+        end = end or total_lines
+        end = min(end, total_lines)
+        selected = "\n".join(lines[start - 1 : end])
+        truncated = False
+        if len(selected) > self.max_read_chars:
+            selected = selected[: self.max_read_chars]
+            truncated = True
+        if start == 1 and end >= total_lines and not truncated:
+            pass
+        self._read_any.add(str(path))
+        header = (
+            f"{self._display(path)} | lines {start}-{end} of {total_lines} | "
+            f"{tf.size} bytes | encoding {tf.encoding}"
+        )
+        if truncated:
+            header += f" | TRUNCATED at {self.max_read_chars} chars; use start_line/end_line"
+        output = header + "\n" + selected
+        if not selected and total_lines == 0:
+            output = header + "\n(empty file)"
+        return ToolResult(name="read_file", arguments=args, ok=True, output=output)
+
+    def _write_file(self, args: dict[str, Any]) -> ToolResult:
+        path = self._resolve(args.get("path"), for_write=True)
+        content = args.get("content")
+        if not isinstance(content, str):
+            raise ToolError("write_file requires a string 'content'.")
+        if len(content) > self.max_write_chars:
+            raise ToolError(
+                f"Content is too large ({len(content)} chars; limit {self.max_write_chars})."
+            )
+        key = str(path)
+        existed = path.exists()
+        if existed and path.is_dir():
+            raise ToolError(f"{self._display(path)} is a directory.")
+        if existed and path.stat().st_size > 0:
+            raise ToolError(
+                f"{self._display(path)} already exists. write_file only creates new files; "
+                "use apply_patch to modify an existing file."
+            )
+        new_norm = content.replace("\r\n", "\n").replace("\r", "\n")
+        if existed:
+            # Existing empty file: replacing it is safe and unambiguous.
+            diff = self._diff("", new_norm, path)
+            self._approve(f"Fill empty file {self._display(path)}?", diff)
+        else:
+            diff = self._diff("", new_norm, path)
+            self._approve(f"Create {self._display(path)}?", diff)
+        self._write_text(path, new_norm, encoding="utf-8", newline="\n")
+        self._mark_written(path)
+        return ToolResult(
+            name="write_file",
+            arguments=args,
+            ok=True,
+            output=f"Wrote {self._display(path)} ({len(new_norm)} chars, {new_norm.count(chr(10)) + 1} lines).",
+            changed=[self._display(path)],
+        )
+
+    def _apply_patch(self, args: dict[str, Any]) -> ToolResult:
+        path = self._resolve(args.get("path"), must_exist=True, for_write=True)
+        if not path.is_file():
+            raise ToolError(f"{self._display(path)} is not a file.")
+        search = args.get("search")
+        replace = args.get("replace")
+        if not isinstance(search, str) or not search:
+            raise ToolError("apply_patch requires a non-empty string 'search'.")
+        if not isinstance(replace, str):
+            raise ToolError("apply_patch requires a string 'replace'.")
+        key = str(path)
+        if key not in self._read_any:
+            raise ToolError(
+                f"Read {self._display(path)} with read_file before patching it."
+            )
+        tf = self._read_text(path)
+        search_norm = search.replace("\r\n", "\n").replace("\r", "\n")
+        replace_norm = replace.replace("\r\n", "\n").replace("\r", "\n")
+        count = tf.text.count(search_norm)
+        if count == 0:
+            raise ToolError(
+                "The search text was not found. Re-read the file and copy the exact "
+                "text, including indentation."
+            )
+        replace_all = bool(args.get("replace_all"))
+        if count > 1 and not replace_all:
+            raise ToolError(
+                f"The search text occurs {count} times. Add more surrounding context "
+                "or pass replace_all=true."
+            )
+        new_text = tf.text.replace(search_norm, replace_norm, -1 if replace_all else 1)
+        if new_text == tf.text:
+            return ToolResult(
+                name="apply_patch", arguments=args, ok=True,
+                output=f"No changes produced for {self._display(path)}.",
+            )
+        if len(new_text) > self.max_write_chars:
+            raise ToolError(
+                f"Patched file would be too large ({len(new_text)} chars; limit {self.max_write_chars})."
+            )
+        diff = self._diff(tf.text, new_text, path)
+        self._approve(f"Apply patch to {self._display(path)}?", diff)
+        self._write_text(path, new_text, encoding=tf.encoding, newline=tf.newline)
+        self._mark_written(path)
+        return ToolResult(
+            name="apply_patch",
+            arguments=args,
+            ok=True,
+            output=f"Patched {self._display(path)} ({count} match(es); +{len(new_text) - len(tf.text)} chars).",
+            changed=[self._display(path)],
+        )
+
+    def _list_files(self, args: dict[str, Any]) -> ToolResult:
+        path = self._resolve(args.get("path") or ".")
+        if not path.exists():
+            raise ToolError(f"Path does not exist: {self._display(path)}")
+        if not path.is_dir():
+            raise ToolError(f"{self._display(path)} is not a directory.")
+        try:
+            depth = int(args.get("depth") or 2)
+        except (TypeError, ValueError):
+            depth = 2
+        depth = max(1, min(depth, 6))
+        max_entries = 300
+        lines: list[str] = []
+        truncated = False
+        for current, dirs, files in os.walk(path):
+            current_path = Path(current)
+            rel = current_path.relative_to(path)
+            level = 0 if str(rel) == "." else len(rel.parts)
+            if level >= depth:
+                dirs[:] = []
+                continue
+            dirs[:] = sorted(d for d in dirs if d not in SKIP_DIRS)
+            for d in dirs:
+                if len(lines) >= max_entries:
+                    truncated = True
+                    break
+                child = current_path / d
+                if child.is_symlink():
+                    continue
+                lines.append(f"{self._display(child)}/")
+            for f in sorted(files):
+                if len(lines) >= max_entries:
+                    truncated = True
+                    break
+                child = current_path / f
+                try:
+                    size = child.stat().st_size
+                except OSError:
+                    size = -1
+                if child.suffix.lower() in TEXT_SKIP_EXTS and size > 2_000_000:
+                    continue
+                lines.append(f"{self._display(child)} ({_human_size(size) if size >= 0 else '?'})")
+            if truncated:
+                break
+        output = "\n".join(lines) if lines else "(directory is empty)"
+        if truncated:
+            output += f"\n... [listing truncated at {max_entries} entries; narrow the path or depth]"
+        return ToolResult(name="list_files", arguments=args, ok=True, output=output)
+
+    def _delete_file(self, args: dict[str, Any]) -> ToolResult:
+        path = self._resolve(args.get("path"), must_exist=True, for_write=True)
+        if path.is_dir():
+            raise ToolError("delete_file only deletes single files, not directories.")
+        if not path.is_file():
+            raise ToolError(f"Not a regular file: {self._display(path)}")
+        try:
+            size = path.stat().st_size
+        except OSError:
+            size = 0
+        self._approve(f"Delete {self._display(path)} ({_human_size(size)})?", None)
+        path.unlink()
+        key = str(path)
+        self._read_any.discard(key)
+        self._changed.append(self._display(path))
+        return ToolResult(
+            name="delete_file",
+            arguments=args,
+            ok=True,
+            output=f"Deleted {self._display(path)}.",
+            changed=[self._display(path)],
+        )
+
+    def _run_command(self, args: dict[str, Any]) -> ToolResult:
+        if not self.allow_shell:
+            raise ToolError(
+                "Shell access is disabled for this run. Ask the user to rerun with --allow-shell "
+                "and approve the command."
+            )
+        command = args.get("command")
+        if not isinstance(command, str) or not command.strip():
+            raise ToolError("run_command requires a non-empty string 'command'.")
+        if len(command) > 8000:
+            raise ToolError("Command is too long.")
+        try:
+            timeout = int(args.get("timeout_sec") or self.shell_timeout_sec)
+        except (TypeError, ValueError):
+            timeout = self.shell_timeout_sec
+        timeout = max(1, min(timeout, 600))
+        self._approve(f"Run command in {self.root}?", command)
+        if os.name == "nt":
+            pwsh = shutil.which("pwsh") or shutil.which("powershell")
+            if pwsh:
+                cmd = [
+                    pwsh,
+                    "-NoLogo",
+                    "-NoProfile",
+                    "-NonInteractive",
+                    "-Command",
+                    "$ProgressPreference='SilentlyContinue'; "
+                    "[Console]::OutputEncoding=[System.Text.Encoding]::UTF8; " + command,
+                ]
+            else:
+                cmd = [os.environ.get("COMSPEC", "cmd.exe"), "/d", "/s", "/c", command]
+        else:
+            cmd = ["/bin/sh", "-c", command]
+        try:
+            proc = subprocess.run(
+                cmd,
+                cwd=str(self.root),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                timeout=timeout,
+            )
+        except subprocess.TimeoutExpired as exc:
+            stdout = _decode(exc.stdout or b"")
+            stderr = _decode(exc.stderr or b"")
+            body = (stdout + stderr).strip()
+            return ToolResult(
+                name="run_command",
+                arguments=args,
+                ok=False,
+                output=f"Command timed out after {timeout}s." + (f"\n{body}" if body else ""),
+            )
+        stdout = _decode(proc.stdout or b"")
+        stderr = _decode(proc.stderr or b"")
+        combined = stdout
+        if stderr:
+            combined += ("\n" if combined and not combined.endswith("\n") else "") + stderr
+        combined = combined.rstrip()
+        if len(combined) > self.max_output_chars:
+            combined = _truncate_middle(combined, self.max_output_chars)
+        if not combined:
+            combined = "(no output)"
+        ok = proc.returncode == 0
+        header = f"exit code {proc.returncode} (cwd {self._display(self.root) or '.'})"
+        return ToolResult(
+            name="run_command",
+            arguments=args,
+            ok=ok,
+            output=header + "\n" + combined,
+        )
+
+    @property
+    def changed_files(self) -> list[str]:
+        return list(self._changed)