nisse/nisse-dev/tests/test_bench.py

"""The driver benchmark."""

import json
import shutil
from pathlib import Path
from typing import Any

import pytest

from nisse_core.canonical import AssistantText, ToolCall
from nisse_core.catalog import ModelFacts
from nisse_core.execution import SSHHost
from nisse_dev.bench.results import (
    append,
    load,
    render,
    results_file,
    summarise,
    wilson,
)
from nisse_dev.bench.runner import MeteringClient, run_task
from nisse_dev.bench.tasks import CHECK_DIR, load_tasks
from nisse_dev.config import DriverTuning, Tuning

# The build VM for every test here: the session's throwaway sshd.
HOST: SSHHost


@pytest.fixture(autouse=True)
def _build_host(ssh_host: SSHHost) -> None:
    global HOST
    HOST = ssh_host


TASKS = load_tasks()
FACTS = ModelFacts(
    alias="local-fast",
    backing="openai/qwen3-8b-awq",
    provider="openai",
    supports_tools=True,
    max_input_tokens=32768,
    max_output_tokens=8192,
    supports_tool_choice=True,
)


# --- the tasks themselves ---------------------------------------------------


def test_the_bundled_tasks_are_the_eight_planned() -> None:
    assert [t.name for t in TASKS] == [
        "fix-bug",
        "long-milestone",
        "long-output",
        "multi-file",
        "refactor-rename",
        "spec-parser",
        "trace-bug",
        "write-function",
    ]


@pytest.mark.parametrize("task", TASKS, ids=lambda t: t.name)
def test_the_seed_fails_its_check_and_the_solution_passes(task, tmp_path: Path) -> None:
    # A check that passes on the untouched seed measures nothing; a task its
    # own reference solution can't pass is unfair to every model.
    seed = tmp_path / "seed"
    task.seed_into(seed)
    assert not task.run_check(seed)[0], f"{task.name}: the seed already passes"

    solved = tmp_path / "solved"
    task.seed_into(solved)
    task.apply_solution(solved)
    passed, output = task.run_check(solved)
    assert passed, f"{task.name}: the reference solution fails:\n{output}"


@pytest.mark.parametrize("task", TASKS, ids=lambda t: t.name)
def test_the_hidden_check_is_not_in_the_seed(task, tmp_path: Path) -> None:
    task.seed_into(tmp_path)
    assert not (tmp_path / CHECK_DIR).exists()
    assert "Call finish_milestone" in task.plan  # as the driver must end


def test_long_output_really_is_long(tmp_path: Path) -> None:
    """The point of the task is output far past tools.max_output_chars."""
    import subprocess
    import sys

    [task] = load_tasks(["long-output"])
    task.seed_into(tmp_path)
    run = subprocess.run(
        [sys.executable, "-m", "pytest", "-v"],
        cwd=tmp_path,
        capture_output=True,
        text=True,
        check=False,
    )
    assert len(run.stdout) > 4 * Tuning().tools.max_output_chars
    assert "R237" in run.stdout[-3000:]  # the failure is named at the end


def test_trace_bug_refuses_a_fix_at_the_call_site(tmp_path: Path) -> None:
    """Fixing only the currency path passes the visible test, but the same
    shared helper is wrong for discounts and tax: the check must notice."""
    [task] = load_tasks(["trace-bug"])
    task.seed_into(tmp_path)
    currency = tmp_path / "invoicing" / "currency.py"
    currency.write_text(
        currency.read_text()
        .replace("apply_rate, money, to_cents", "money, to_cents")
        .replace(
            "return apply_rate(amount, rate(from_currency, to_currency))",
            "return to_cents(amount * rate(from_currency, to_currency))",
        )
    )
    passed, output = task.run_check(tmp_path)
    assert not passed
    assert "test_fractional_discount" in output


def test_refactor_rename_refuses_a_blind_replace(tmp_path: Path) -> None:
    """A plain text replace also renames load_config_defaults and
    reload_config, which must keep their names."""
    [task] = load_tasks(["refactor-rename"])
    task.seed_into(tmp_path)
    for path in [*tmp_path.rglob("*.py"), *tmp_path.rglob("*.md")]:
        path.write_text(path.read_text().replace("load_config", "read_settings"))
    passed, output = task.run_check(tmp_path)
    assert not passed
    assert "test_the_similar_names_were_left_alone" in output


def test_refactor_rename_refuses_a_partial_rename(tmp_path: Path) -> None:
    [task] = load_tasks(["refactor-rename"])
    task.seed_into(tmp_path)
    task.apply_solution(tmp_path)
    readme = tmp_path / "README.md"
    readme.write_text(readme.read_text().replace("read_settings", "load_config", 1))
    passed, output = task.run_check(tmp_path)
    assert not passed
    assert "README.md" in output


def test_a_changed_task_gets_a_new_hash(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    shutil.copytree(task.path, tmp_path / "write-function")
    [copy] = load_tasks(root=tmp_path)
    assert copy.content_hash == task.content_hash
    (copy.path / "seed" / "README.md").write_text("changed\n")
    assert copy.content_hash != task.content_hash


def test_an_unknown_task_is_refused() -> None:
    with pytest.raises(ValueError, match="unknown task"):
        load_tasks(["nope"])


# --- running a task ----------------------------------------------------------


class Scripted:
    """A model that plays back fixed turns, reporting itself as `answered`."""

    def __init__(self, turns: list[list[Any]], answered: str = "local-fast") -> None:
        self.turns = turns
        self.last_usage = {"input": 1000, "output": 50}
        self.last_cost = 0.0
        self.last_model = answered
        self.trimmed_requests = 0

    def complete(self, events, tools, model, system=None):
        if not self.turns:
            return [AssistantText(text="nothing left to do", model=model)]
        return self.turns.pop(0)


def call(tool: str, cid: str, **arguments: Any) -> ToolCall:
    return ToolCall(call_id=cid, tool=tool, arguments=arguments)


SOLVE_SLUGIFY = [
    [
        call(
            "write_file",
            "w",
            path="slugify.py",
            content=(
                "import re\n\n\ndef slugify(text):\n"
                "    return re.sub(r'[^a-z0-9]+', '-', text.lower()).strip('-')\n"
            ),
        )
    ],
    [call("run_command", "r", command="python -c 'import slugify'")],
    [call("finish_milestone", "f", summary="added slugify", final="yes")],
]


def test_a_solved_task_is_a_success(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    record = run_task(
        task, FACTS, Scripted(list(SOLVE_SLUGIFY)), Tuning(), host=HOST, note="n"
    )

    assert record.success
    assert record.outcome == "completed"
    assert record.turns == 3
    assert record.tokens_in == 3000
    assert record.task_hash == task.content_hash
    assert record.model_info["supports_tool_choice"] is True
    assert record.tuning["driver"]["max_turns"] == 40


def test_saying_finished_is_not_success(tmp_path: Path) -> None:
    # Success is the hidden check passing, never the model's own claim.
    [task] = load_tasks(["write-function"])
    wrong = [
        [
            call(
                "write_file",
                "w",
                path="slugify.py",
                content="def slugify(t):\n    return t\n",
            )
        ],
        [call("finish_milestone", "f", summary="done", final="yes")],
    ]
    record = run_task(task, FACTS, Scripted(wrong), Tuning(), host=HOST)
    assert record.outcome == "completed"
    assert not record.success


def test_misbehaviour_is_counted(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    turns = [
        [call("chat", "x1")],
        [call("run_command", "r1", command="exit 1")],
        [AssistantText(text="let me explain", model="m")],
    ]
    # A higher failure limit, so the invented tool and the failed command
    # don't escalate before the prose turn is reached.
    tuning = Tuning(driver=DriverTuning(failure_limit=5))
    record = run_task(task, FACTS, Scripted(turns), tuning, host=HOST)
    assert record.invented_tools == 1
    assert record.failed_commands == 1
    assert record.prose_turns >= 1
    assert not record.success


def test_asking_the_operator_is_an_escalation(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    asks = [[call("ask_human", "q", question="which file?")]]
    record = run_task(task, FACTS, Scripted(asks), Tuning(), host=HOST)
    assert (record.outcome, record.success) == ("escalated", False)


def test_the_driver_gets_the_loops_own_prompt_and_tools(tmp_path: Path) -> None:
    from nisse_dev.loop.agent import SYSTEM

    seen: list[tuple[str, set[str]]] = []

    class Spy(Scripted):
        def complete(self, events, tools, model, system=None):  # type: ignore[no-untyped-def]
            seen.append((system or "", {t["function"]["name"] for t in tools}))
            return super().complete(events, tools, model, system)

    [task] = load_tasks(["write-function"])
    run_task(task, FACTS, Spy(list(SOLVE_SLUGIFY)), Tuning(), host=HOST)
    system, tools = seen[0]
    assert system.startswith(SYSTEM.split("{")[0])
    assert {"start_milestone", "finish_milestone", "run_command"} <= tools
    assert not {"delegate", "ask_expert", "restate"} & tools


def test_a_fallback_answer_can_never_be_a_success(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    client = Scripted(list(SOLVE_SLUGIFY), answered="claude-sonnet-5")
    record = run_task(task, FACTS, client, Tuning(), host=HOST)
    assert record.fallbacks == ["claude-sonnet-5"] * 3
    assert not record.success  # the check passed, but not for this model


# --- storing and reporting ---------------------------------------------------


def test_wilson_intervals_behave() -> None:
    low, high = wilson(5, 5)
    assert high == 1.0 and 0.5 < low < 0.6
    # 3 of 5 and 5 of 5 overlap: five runs do not separate them.
    assert wilson(3, 5)[1] > wilson(5, 5)[0]
    assert wilson(0, 0) == (0.0, 1.0)


def record_json(**overrides: Any) -> dict[str, Any]:
    [task] = load_tasks(["write-function"])
    base = run_task(
        task, FACTS, Scripted(list(SOLVE_SLUGIFY)), Tuning(), host=HOST
    ).to_json()
    base.update(overrides)
    return base


def test_results_round_trip_and_report(tmp_path: Path) -> None:
    from nisse_dev.bench.runner import RunRecord

    rec = RunRecord(**record_json())
    path = results_file(tmp_path, "local-fast")
    append(path, rec)
    append(path, RunRecord(**record_json(success=False)))
    stored = load(tmp_path)
    assert len(stored) == 2
    assert json.loads(path.read_text().splitlines()[0])["model"] == "local-fast"

    current = {t.name: t.content_hash for t in TASKS}
    [row] = summarise(stored, current)
    assert (row.successes, row.runs) == (1, 2)
    assert not row.stale
    table = render([row])
    assert "openai/qwen3-8b-awq" in table
    assert "1/2" in table


def test_runs_of_an_older_task_version_are_flagged() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    rows = summarise([record_json(task_hash="0123456789abcdef")], current)
    assert rows[0].stale
    assert "(old)" in render(rows)


def test_models_served_differently_are_kept_apart() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    rows = summarise(
        [record_json(note="vLLM, AWQ"), record_json(note="llama.cpp, Q4_K_M")], current
    )
    assert len(rows) == 2


def test_metering_counts_prose_and_cost() -> None:
    class Inner:
        def __init__(self) -> None:
            self.last_usage = {"input": 10, "output": 2}
            self.last_cost = 0.5
            self.last_model = "local-fast"

        def complete(self, events, tools, model, system=None):
            return [AssistantText(text="hi", model=model)]

    meter = MeteringClient(Inner(), FACTS)
    meter.complete([], None, "local-fast")
    meter.complete([], None, "local-fast")
    assert (meter.calls, meter.prose_turns, meter.cost) == (2, 2, 1.0)


def test_metering_counts_rate_limit_waits_even_on_failure() -> None:
    class Inner:
        def __init__(self) -> None:
            self.last_usage = {"input": 10, "output": 2}
            self.last_cost = 0.0
            self.last_model = "local-fast"
            self.last_rate_limit_wait_s = 0.0
            self.calls = 0

        def complete(self, events, tools, model, system=None):
            self.calls += 1
            self.last_rate_limit_wait_s = 20.0
            if self.calls == 2:
                raise RuntimeError("upstream returned 429")
            return [AssistantText(text="hi", model=model)]

    meter = MeteringClient(Inner(), FACTS)
    meter.complete([], None, "local-fast")
    with pytest.raises(RuntimeError):
        meter.complete([], None, "local-fast")
    assert meter.rate_limit_wait_s == 40.0


def test_a_run_records_its_rate_limit_waits(tmp_path: Path) -> None:
    [task] = load_tasks(["write-function"])
    client = Scripted(list(SOLVE_SLUGIFY))
    client.last_rate_limit_wait_s = 20.0  # every call waited once
    record = run_task(task, FACTS, client, Tuning(), host=HOST)
    assert record.rate_limit_wait_s == 60.0


def test_the_report_leaves_rate_limit_waits_out_of_time() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    rows = summarise(
        [
            record_json(wall_s=50.0, rate_limit_wait_s=40.0),
            record_json(wall_s=10.0, rate_limit_wait_s=0.0),
        ],
        current,
    )
    [row] = rows
    assert row.mean_wall_s == 10.0
    assert row.rate_limited == 1
    assert "429 = runs that waited" in render(rows)


def test_results_stored_before_waits_were_recorded_still_load() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    old = record_json(wall_s=12.0)
    del old["rate_limit_wait_s"]
    [row] = summarise([old], current)
    assert (row.mean_wall_s, row.rate_limited) == (12.0, 0)


def test_runs_with_different_tuning_are_kept_apart() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    capped = record_json()
    capped["tuning"]["driver"]["max_turns"] = 7
    rows = summarise([record_json(), capped], current)
    assert sorted(row.tuning for row in rows) == ["", "driver.max_turns=7"]
    table = render(rows)
    assert "driver.max_turns=7" in table
    assert "tuning = settings that differ" in table


def test_a_report_at_default_tuning_has_no_tuning_column() -> None:
    current = {t.name: t.content_hash for t in TASKS}
    table = render(summarise([record_json()], current))
    assert "tuning" not in table


def test_runs_from_before_the_driver_read_as_its_defaults() -> None:
    # Stored before M2.8, the driver's section was called executor.
    current = {t.name: t.content_hash for t in TASKS}
    old = record_json()
    old["tuning"]["executor"] = old["tuning"].pop("driver")
    [row] = summarise([old], current)
    assert row.tuning == ""