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

"""Running one task once with one executor model, and recording it."""

import io
import shlex
import subprocess
import tempfile
import time
import uuid
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from nisse_dev.bench.tasks import Task
from nisse_dev.canonical import (
    CanonicalEvent,
    PlanSubmitted,
    StoredEvent,
    ToolCall,
    ToolResult,
    UserMessage,
)
from nisse_dev.catalog import ModelFacts
from nisse_dev.config import Tuning
from nisse_dev.execution import SSHHost
from nisse_dev.executor import LLMExecutor
from nisse_dev.llm import RoleMap
from nisse_dev.oplog import OpLog
from nisse_dev.state import StateStore
from nisse_dev.tools import Workspace


class MeteringClient:
    """Counts what a run spends, and notices a fallback: an answer from a
    model other than the one asked for is not a measurement of it."""

    def __init__(self, inner: Any, facts: ModelFacts) -> None:
        self._inner = inner
        self._facts = facts
        self.calls = 0
        self.prose_turns = 0
        self.tokens_in = 0
        self.tokens_out = 0
        self.cost = 0.0
        self.fallbacks: list[str] = []
        self.rate_limit_wait_s = 0.0

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        self.calls += 1
        try:
            produced = self._inner.complete(events, tools, model, system=system)
        finally:
            # Counted even when the call fails: a rate limit that never
            # cleared still cost the run its wait.
            self.rate_limit_wait_s += float(
                getattr(self._inner, "last_rate_limit_wait_s", 0.0) or 0.0
            )
        usage = getattr(self._inner, "last_usage", {}) or {}
        self.tokens_in += int(usage.get("input", 0))
        self.tokens_out += int(usage.get("output", 0))
        self.cost += float(getattr(self._inner, "last_cost", 0.0) or 0.0)
        answered = getattr(self._inner, "last_model", None)
        if not self._facts.answered(answered):
            self.fallbacks.append(str(answered))
        if not any(isinstance(e, ToolCall) for e in produced):
            self.prose_turns += 1
        return produced


@dataclass
class RunRecord:
    """One run, with everything needed to compare it with runs made later
    against other models, other servers, or other tuning."""

    ts: str
    model: str
    backing: str
    provider: str
    model_info: dict[str, Any]
    note: str
    task: str
    task_hash: str
    commit: str
    tuning: dict[str, Any]
    success: bool
    outcome: str
    outcome_note: str
    turns: int
    prose_turns: int
    failed_commands: int
    invented_tools: int
    trimmed_requests: int
    tokens_in: int
    tokens_out: int
    cost: float
    wall_s: float
    fallbacks: list[str] = field(default_factory=list)
    # Part of wall_s spent waiting out the proxy's rate limit, not working.
    rate_limit_wait_s: float = 0.0
    check_output: str = ""

    def to_json(self) -> dict[str, Any]:
        return asdict(self)


def agent_commit(repo: Path | None = None) -> str:
    """The agent's commit, marked when the tree had uncommitted changes."""
    where = repo or Path(__file__).resolve().parent
    try:
        sha = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=where,
            capture_output=True,
            text=True,
            check=True,
        ).stdout.strip()
        dirty = subprocess.run(
            ["git", "status", "--porcelain"],
            cwd=where,
            capture_output=True,
            text=True,
            check=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return "unknown"
    return f"{sha}+dirty" if dirty else sha


def _count(events: list[StoredEvent]) -> tuple[int, int]:
    tools = {
        e.event.call_id: e.event.tool for e in events if isinstance(e.event, ToolCall)
    }
    failed = invented = 0
    for stored in events:
        event = stored.event
        if isinstance(event, ToolResult) and not event.success:
            if event.result.startswith("no such tool"):
                invented += 1
            elif tools.get(event.call_id) == "run_command":
                failed += 1
    return failed, invented


def run_task(
    task: Task,
    facts: ModelFacts,
    client: Any,
    tuning: Tuning,
    *,
    host: SSHHost,
    note: str = "",
    commit: str = "",
    workdir: Path | None = None,
) -> RunRecord:
    """Run one task once: seed a workspace, hand the executor the approved
    plan, then judge the result with the hidden check.

    Delegation and ask_expert are off: an escalation is a failure of the
    model under test, not a frontier model's success."""
    with tempfile.TemporaryDirectory(dir=workdir) as tmp:
        root = Path(tmp)
        # The workspace is on the build VM, like the agent's own (L13).
        workspace = Workspace(
            host,
            f"/tmp/nisse-bench/{uuid.uuid4().hex}",
            timeout_s=tuning.tools.command_timeout_s,
            max_output_chars=tuning.tools.max_output_chars,
        )
        task.seed_workspace(workspace)
        store = StateStore(root / "state.db", oplog=OpLog("nisse-bench", io.StringIO()))
        # One work item per run, as the agent's own tasks are (L1).
        store.open_work_item(
            "bench", title=task.name, position="EXECUTING", cause="benchmark run"
        )
        store.transcript("bench", "planner").append(
            PlanSubmitted(plan=task.plan, model="benchmark")
        )
        store.transcript("bench", "operator").append(
            UserMessage(text="start milestone 1", source="benchmark")
        )

        meter = MeteringClient(client, facts)
        trims_before = getattr(client, "trimmed_requests", 0)
        model = facts.alias
        roles = RoleMap(executor=model, planner=model, helper=model, expert=model)
        executor = LLMExecutor(
            meter,
            roles,
            workspace,
            max_turns=tuning.executor.max_turns,
            failure_limit=tuning.executor.failure_limit,
        )
        started = time.monotonic()
        try:
            outcome = executor.run(store.transcript("bench", "executor"), 1)
            status, detail = outcome.status, outcome.note
        except Exception as exc:  # noqa: BLE001 -- recorded, the benchmark goes on
            status, detail = "error", repr(exc)
        wall = time.monotonic() - started

        success, check_output = task.check_workspace(workspace)
        host.execute(f"rm -rf -- {shlex.quote(str(workspace.root))}", timeout_s=60)
        failed, invented = _count(store.transcript("bench", "executor").read())
        store.close()

    return RunRecord(
        ts=datetime.now(UTC).isoformat(timespec="seconds"),
        model=model,
        backing=facts.backing,
        provider=facts.provider,
        model_info={
            "max_input_tokens": facts.max_input_tokens,
            "max_output_tokens": facts.max_output_tokens,
            "supports_tool_choice": facts.supports_tool_choice,
            "supports_function_calling": facts.supports_tools,
        },
        note=note,
        task=task.name,
        task_hash=task.content_hash,
        commit=commit,
        tuning=asdict(tuning),
        success=success and not meter.fallbacks,
        outcome=status,
        outcome_note=detail[:500],
        turns=meter.calls,
        prose_turns=meter.prose_turns,
        failed_commands=failed,
        invented_tools=invented,
        trimmed_requests=getattr(client, "trimmed_requests", 0) - trims_before,
        tokens_in=meter.tokens_in,
        tokens_out=meter.tokens_out,
        cost=round(meter.cost, 6),
        wall_s=round(wall, 1),
        fallbacks=meter.fallbacks,
        rate_limit_wait_s=meter.rate_limit_wait_s,
        check_output=check_output,
    )