nisse/nisse-dev/src/nisse_dev/bench/tasks.py

"""Benchmark tasks: a seed project, an approved plan, and a hidden check.

Each task is a directory under ``tasks/``:

- ``task.yaml``: ``title``, ``plan`` (the approved milestone plan, written the
  way the driver restates them), and ``check`` (a shell command run
  in the workspace afterwards).
- ``seed/``: the project the driver starts from.
- ``hidden/``: files the check needs, copied in only after the run, into
  ``.bench_check/``. The driver never sees them.
- ``solution/``: a reference solution, overlaid on the seed by the tests to
  prove each task can be solved and that its check tells the difference.
"""

import hashlib
import shlex
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import yaml

if TYPE_CHECKING:
    from nisse_dev.tools import Workspace

TASKS_DIR = Path(__file__).parent / "tasks"
CHECK_DIR = ".bench_check"


@dataclass(frozen=True)
class Task:
    name: str
    title: str
    plan: str
    check: str
    path: Path

    @property
    def content_hash(self) -> str:
        """What the task is, so a changed task is never compared with runs of
        its earlier version. Covers everything the model or the check sees;
        not the reference solution."""
        digest = hashlib.sha256()
        for part in ("task.yaml", "seed", "hidden"):
            root = self.path / part
            files = [root] if root.is_file() else sorted(root.rglob("*"))
            for file in files:
                if file.is_file() and "__pycache__" not in file.parts:
                    digest.update(str(file.relative_to(self.path)).encode())
                    digest.update(file.read_bytes())
        return digest.hexdigest()[:16]

    def seed_into(self, workspace: Path) -> None:
        shutil.copytree(self.path / "seed", workspace, dirs_exist_ok=True)

    def apply_solution(self, workspace: Path) -> None:
        shutil.copytree(self.path / "solution", workspace, dirs_exist_ok=True)

    def run_check(self, workspace: Path, timeout_s: int = 300) -> tuple[bool, str]:
        """Copy the hidden files in and run the check. Returns whether it
        passed, and its output for the record."""
        target = workspace / CHECK_DIR
        if target.exists():
            shutil.rmtree(target)  # never trust a directory the model made
        shutil.copytree(self.path / "hidden", target)
        try:
            result = subprocess.run(
                self.check,
                shell=True,
                cwd=workspace,
                capture_output=True,
                text=True,
                timeout=timeout_s,
                check=False,
            )
        except subprocess.TimeoutExpired:
            return False, f"check timed out after {timeout_s}s"
        output = (result.stdout + result.stderr)[-2000:]
        return result.returncode == 0, output

    # -- on the build VM, where the agent's runs happen (L13) ---------------

    def seed_workspace(self, workspace: "Workspace") -> None:
        workspace.upload_tree(self.path / "seed")

    def check_workspace(
        self, workspace: "Workspace", timeout_s: int = 300
    ) -> tuple[bool, str]:
        """run_check, on the build VM: copy the hidden files in fresh (never
        trusting a directory the model made) and run the check there."""
        root = shlex.quote(str(workspace.root / CHECK_DIR))
        workspace.host.execute(f"rm -rf -- {root}", timeout_s=60)
        workspace.upload_tree(self.path / "hidden", CHECK_DIR)
        result = workspace.host.execute(
            self.check, cwd=str(workspace.root), timeout_s=timeout_s
        )
        if result.timed_out:
            return False, f"check timed out after {timeout_s}s"
        return result.exit_code == 0, (result.stdout + result.stderr)[-2000:]


def load_task(path: Path) -> Task:
    spec = yaml.safe_load((path / "task.yaml").read_text())
    return Task(
        name=path.name,
        title=str(spec["title"]),
        plan=str(spec["plan"]).strip(),
        check=str(spec["check"]),
        path=path,
    )


def load_tasks(names: list[str] | None = None, root: Path = TASKS_DIR) -> list[Task]:
    available = {
        p.name: p for p in sorted(root.iterdir()) if (p / "task.yaml").exists()
    }
    if names:
        unknown = [n for n in names if n not in available]
        if unknown:
            raise ValueError(
                f"unknown task(s): {', '.join(unknown)}; "
                f"available: {', '.join(available)}"
            )
        return [load_task(available[n]) for n in names]
    return [load_task(p) for p in available.values()]