nisse/nisse-dev/src/nisse_dev/tools.py

"""Executor tools (architecture L9: reduced schema, flat arguments; L13).

Every tool runs on the work item's build VM over SSH (execution.py); the
agent's own host never runs a command a model chose. Within the VM nothing is
sandboxed: the VM is the trust boundary (D7), so per-command confinement would
add failure modes without adding safety. The workspace is where work happens,
the default for relative paths and the command's directory, not a wall.

What *is* enforced here is operational: commands cannot hang the loop and
cannot flood the context window with megabytes of build log. The agent's own
credentials are never on the build VM to inherit.
"""

import base64
import io
import shlex
import tarfile
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from nisse_dev.context import head_and_tail
from nisse_dev.execution import HostUnavailable, Result, SSHHost

# Tools that change or verify something. Reading a file or listing a
# directory is a model looking around, not doing the work, so neither counts
# as evidence that work was carried out.
MUTATING_TOOLS = frozenset({"write_file", "run_command"})

DEFAULT_TIMEOUT_S = 120
MAX_OUTPUT_CHARS = 4000

EXECUTION_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read a UTF-8 text file.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}},
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Write a UTF-8 text file, creating parent directories.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "content": {"type": "string"},
                },
                "required": ["path", "content"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "list_dir",
            "description": "List the entries of a directory.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}},
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_command",
            "description": (
                "Run a shell command in the workspace. Returns its exit code "
                "and output; a non-zero exit is a result, not an error."
            ),
            "parameters": {
                "type": "object",
                "properties": {"command": {"type": "string"}},
                "required": ["command"],
            },
        },
    },
]


@dataclass(frozen=True)
class ToolOutcome:
    result: str
    success: bool


class Workspace:
    """The work item's workspace: a directory on its tool VM (L13, D3).

    Every operation runs there over SSH; nothing touches the agent's own
    host. ``root`` is the path on that host.
    """

    def __init__(
        self,
        host: SSHHost,
        root: Path | str,
        *,
        timeout_s: int = DEFAULT_TIMEOUT_S,
        max_output_chars: int = MAX_OUTPUT_CHARS,
        secrets: Mapping[str, str] | None = None,
    ) -> None:
        self.host = host
        self.root = Path(root)
        self._timeout_s = timeout_s
        self._max_output_chars = max_output_chars
        # The project secrets the model's commands may see (D7); redacted from
        # whatever comes back.
        self._secrets = dict(secrets or {})
        self._run(f"mkdir -p -- {shlex.quote(str(self.root))}")

    def with_timeout(self, timeout_s: int) -> "Workspace":
        """The same workspace with a longer limit per command (installs)."""
        return Workspace(
            self.host,
            self.root,
            timeout_s=timeout_s,
            max_output_chars=self._max_output_chars,
            secrets=self._secrets,
        )

    def _run(
        self, command: str, *, stdin: str | None = None, cwd: str | None = None
    ) -> Result:
        return self.host.execute(
            command, cwd=cwd, stdin=stdin, timeout_s=self._timeout_s
        )

    def _truncate(self, text: str) -> str:
        return head_and_tail(text, self._max_output_chars)

    def _resolve(self, path: str) -> str:
        candidate = Path(path)
        return str(candidate if candidate.is_absolute() else self.root / candidate)

    def _inside(self, target: str) -> bool:
        result = self._run(
            f"realpath -m -- {shlex.quote(target)} {shlex.quote(str(self.root))}"
        )
        lines = result.stdout.splitlines()
        if not result.ok or len(lines) != 2:
            return False
        resolved, root = (Path(line) for line in lines)
        return resolved == root or root in resolved.parents

    def dispatch(self, tool: str, arguments: dict[str, Any]) -> ToolOutcome:
        try:
            match tool:
                case "read_file":
                    return self.read_file(str(arguments["path"]))
                case "write_file":
                    return self.write_file(
                        str(arguments["path"]), str(arguments["content"])
                    )
                case "list_dir":
                    return self.list_dir(str(arguments.get("path", ".")))
                case "run_command":
                    return self.run_command(str(arguments["command"]))
                case _:
                    return ToolOutcome(f"no such tool: {tool}", success=False)
        except KeyError as exc:
            return ToolOutcome(f"missing argument {exc} for {tool}", success=False)
        except HostUnavailable as exc:
            # Not the work failing: the build VM is out of reach.
            return ToolOutcome(f"the build VM is unreachable: {exc}", success=False)

    def read_file(self, path: str) -> ToolOutcome:
        target = self._resolve(path)
        # base64 on the wire, so the bytes arrive intact and decoding is
        # strict here, as it was for a local file.
        result = self._run(f"base64 -w0 -- {shlex.quote(target)}")
        if not result.ok:
            return ToolOutcome(
                f"could not read {path}: {result.stderr.strip()}", success=False
            )
        try:
            text = base64.b64decode(result.stdout).decode()
        except (ValueError, UnicodeDecodeError) as exc:
            return ToolOutcome(f"could not read {path}: {exc}", success=False)
        return ToolOutcome(self._truncate(text), success=True)

    def write_file(self, path: str, content: str) -> ToolOutcome:
        target = self._resolve(path)
        # Not a sandbox — the VM is the trust boundary (D7). This is about the
        # work existing at all: a file written outside the repo is invisible
        # to git, so the milestone would publish an empty PR while the model
        # truthfully reports success. Models do guess plausible absolute paths
        # like /workspace, so say where the workspace actually is.
        if not self._inside(target):
            return ToolOutcome(
                f"refusing to write outside the workspace: {path}. "
                f"The workspace root is {self.root}; write there, using "
                "relative paths.",
                success=False,
            )
        quoted = shlex.quote(target)
        result = self._run(
            f'mkdir -p -- "$(dirname -- {quoted})" && base64 -d > {quoted}',
            stdin=base64.b64encode(content.encode()).decode(),
        )
        if not result.ok:
            return ToolOutcome(
                f"could not write {path}: {result.stderr.strip()}", success=False
            )
        return ToolOutcome(f"wrote {len(content)} characters to {path}", success=True)

    def list_dir(self, path: str) -> ToolOutcome:
        target = self._resolve(path)
        result = self._run(f"ls -1Ap -- {shlex.quote(target)}")
        if not result.ok:
            return ToolOutcome(
                f"could not list {path}: {result.stderr.strip()}", success=False
            )
        entries = sorted(line for line in result.stdout.splitlines() if line)
        return ToolOutcome(
            self._truncate("\n".join(entries) or "(empty)"), success=True
        )

    def describe(self) -> str:
        """Told to the model, so it does not have to guess where it is."""
        return (
            f"Your workspace is {self.root} on the build VM. It is a git "
            "repository and the only place your work is recorded; commands run "
            "there by default. Use relative paths."
        )

    def run_command(self, command: str) -> ToolOutcome:
        try:
            result = self.host.execute(
                command,
                cwd=str(self.root),
                secrets=self._secrets,
                timeout_s=self._timeout_s,
            )
        except HostUnavailable as exc:
            return ToolOutcome(f"the build VM is unreachable: {exc}", success=False)
        if result.timed_out:
            return ToolOutcome(
                f"command timed out after {self._timeout_s}s: {command}",
                success=False,
            )
        body = self._truncate(
            f"exit code: {result.exit_code}\nstdout:\n{result.stdout}"
            f"\nstderr:\n{result.stderr}"
        )
        return ToolOutcome(body, success=result.exit_code == 0)

    def upload_tree(self, local: Path, dest: str = ".") -> None:
        """Copy a local directory's contents into the workspace (the bench
        seeding a task); raises RuntimeError if it fails."""
        buffer = io.BytesIO()
        with tarfile.open(fileobj=buffer, mode="w:gz") as archive:
            for entry in sorted(Path(local).iterdir()):
                archive.add(entry, arcname=entry.name)
        target = shlex.quote(self._resolve(dest))
        result = self._run(
            f"mkdir -p -- {target} && base64 -d | tar -xzf - -C {target}",
            stdin=base64.b64encode(buffer.getvalue()).decode(),
        )
        if not result.ok:
            raise RuntimeError(f"upload to {dest} failed: {result.stderr.strip()}")

    def reset(self) -> None:
        """Discard workspace contents; used by tests, not by the agent."""
        root = shlex.quote(str(self.root))
        self._run(f"rm -rf -- {root} && mkdir -p -- {root}")