chemenu/tools/chemenu/tests/test_run_budget.py

import json
import os
import subprocess
import sys

import pytest
import typer

from chemenu import cli, config  # noqa: F401 - importing the full CLI registers every
# command's cli_contract record, which is_exempt (Gitea #121 B8) now reads instead of a
# hand-maintained list; run_budget.py cannot do this import itself (cli.py imports run_budget).
from chemenu.commands import run_budget


@pytest.fixture(autouse=True)
def isolated_state(tmp_path, monkeypatch):
    """Point the budget state file at a scratch location and pin the session
    id, so tests never touch the real .wikitool_session/ dir or bleed into
    the actual calling shell's session."""
    state_file = tmp_path / "budget.json"
    monkeypatch.setattr(run_budget, "STATE_DIR", tmp_path)
    monkeypatch.setattr(run_budget, "STATE_FILE", state_file)
    monkeypatch.setattr(run_budget, "LOCK_FILE", tmp_path / "budget.lock")
    monkeypatch.setenv("WIKITOOL_SESSION_ID", "test-session")
    return state_file


def test_default_thresholds_match_the_contract():
    assert run_budget.DEFAULT_CALL_LIMIT == 60
    assert run_budget.DEFAULT_LOOP_WINDOW == 3


def test_calls_below_limit_pass_silently():
    for i in range(5):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    state = run_budget._load_state()
    assert state["test-session"]["count"] == 5


def test_call_limit_trips_past_threshold():
    limit = run_budget.DEFAULT_CALL_LIMIT
    for i in range(limit):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("new", ["entity", "--name", "OneTooMany"], override=False)


def test_call_limit_message_mentions_contract_and_override():
    message = run_budget.call_limit_message(61, 60)
    assert "Iteration Budget Gate and loop-breaker" in message
    assert "--override-budget" in message
    assert "61" in message and "60" in message


def test_message_sections_exist_in_the_documents_they_cite():
    """Both refusal messages point a reader at a section by name (Gitea
    #147: they used to point at "the tooling contract's 'Iteration and Cost
    Limits'" and "...Tool Error Contracts", neither of which exists). Pin
    that the cited headings are real, so the two drift apart again only if
    this test is touched too."""
    agents_md = (config._PACKAGE_ROOT / "AGENTS.md").read_text(encoding="utf-8")
    gates_md = (config._PACKAGE_ROOT / "instructions" / "gates.md").read_text(encoding="utf-8")
    assert "## Tool error contract" in agents_md
    assert "## Iteration Budget Gate and loop-breaker" in gates_md


def test_record_and_check_reports_whether_it_charged():
    assert run_budget.record_and_check("new", ["entity"], override=False) is True
    assert run_budget.record_and_check("search", ["anything"], override=False) is False


def test_refund_gives_back_the_slot_but_keeps_the_history():
    """A declined call did not iterate on the wiki, so it costs nothing - but
    the loop-breaker still has to see that it happened."""
    for i in range(3):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)
    run_budget.refund()
    entry = run_budget._load_state()["test-session"]
    assert entry["count"] == 3
    assert entry["recent"][-1] == "new source --set raw_files=nope"


def test_refund_never_drives_the_counter_negative():
    run_budget.refund()
    run_budget.record_and_check("new", ["entity"], override=False)
    run_budget.refund()
    run_budget.refund()
    assert run_budget._load_state()["test-session"]["count"] == 0


def test_three_identical_declined_calls_still_trip_the_loop_breaker():
    """The counter is refunded, the history is not - which is what makes the
    loop-breaker the right instrument for a repeated broken invocation."""
    for _ in range(3):
        run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)
        run_budget.refund()
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("new", ["source", "--set", "raw_files=nope"], override=False)


def test_refused_call_is_not_counted():
    """A refused call never ran, so it must not inflate the number quoted
    back to the user on the next attempt."""
    limit = run_budget.DEFAULT_CALL_LIMIT
    for i in range(limit):
        run_budget.record_and_check("new", ["entity", "--name", f"Page{i}"], override=False)
    for _ in range(3):
        with pytest.raises(typer.Exit):
            run_budget.record_and_check("new", ["entity", "--name", "Blocked"], override=False)
    assert run_budget._load_state()["test-session"]["count"] == limit


def test_stale_sessions_are_pruned_on_write():
    now = 1_000_000.0
    state = {
        "fresh": {"count": 1, "recent": [], "last_seen": now - 60},
        "stale": {"count": 1, "recent": [], "last_seen": now - run_budget.SESSION_TTL_SECONDS - 1},
        "legacy-no-timestamp": {"count": 1, "recent": []},
    }
    pruned = run_budget.prune_state(state, now)
    assert set(pruned) == {"fresh", "legacy-no-timestamp"}


def test_recorded_call_stamps_last_seen():
    run_budget.record_and_check("lint", [], override=False)
    assert "last_seen" in run_budget._load_state()["test-session"]


def test_loop_breaker_trips_on_identical_repeats():
    args = ["add", "--a", "X", "--b", "Y"]
    for _ in range(3):
        run_budget.record_and_check("xref", args, override=False)
    with pytest.raises(typer.Exit):
        run_budget.record_and_check("xref", args, override=False)


def test_loop_breaker_does_not_trip_on_varying_calls():
    for i in range(5):
        run_budget.record_and_check("xref", ["add", "--a", f"X{i}", "--b", "Y"], override=False)
    # No exception raised - varying args are not a loop.


def test_loop_breaker_message_mentions_signature():
    message = run_budget.loop_breaker_message("xref add --a X --b Y", 3)
    assert "xref add --a X --b Y" in message
    assert "Loop-Breaker" in message


def test_override_bypasses_both_gates():
    args = ["add", "--a", "X", "--b", "Y"]
    for _ in range(40):
        run_budget.record_and_check("xref", args, override=True)
    state = run_budget._load_state()
    assert state["test-session"]["count"] == 40


def test_a_bucket_without_a_recorded_origin_keeps_its_count():
    """Grandfathering (#110 decision 1): an entry written before this field
    existed must not be reset the moment this ships - that would throw away
    real, in-flight state on every existing instance's first call after
    upgrading."""
    run_budget._save_state({"test-session": {"count": 5, "recent": ["lint"]}})
    run_budget.record_and_check("lint", [], override=False)
    entry = run_budget._load_state()["test-session"]
    assert entry["count"] == 6
    assert entry["source"] == "WIKITOOL_SESSION_ID"


def test_a_bucket_with_a_different_recorded_origin_starts_over():
    """The same id string, stamped by a different origin than the one
    recorded, is treated as a stranger's bucket rather than inherited - the
    mechanism behind #110's 'no bucket is silently reinterpreted' criterion."""
    run_budget._save_state(
        {"test-session": {"count": 40, "recent": ["lint"], "source": "getppid() fallback"}}
    )
    run_budget.record_and_check("lint", [], override=False)
    entry = run_budget._load_state()["test-session"]
    assert entry["count"] == 1
    assert entry["source"] == "WIKITOOL_SESSION_ID"


def test_a_bucket_with_the_same_recorded_origin_keeps_counting():
    run_budget._save_state(
        {"test-session": {"count": 7, "recent": ["lint"], "source": "WIKITOOL_SESSION_ID"}}
    )
    run_budget.record_and_check("lint", [], override=False)
    assert run_budget._load_state()["test-session"]["count"] == 8


def test_save_state_writes_atomically_and_leaves_no_tmp_file(isolated_state):
    run_budget._save_state({"test-session": {"count": 1, "recent": []}})
    assert isolated_state.exists()
    tmp_file = isolated_state.with_suffix(isolated_state.suffix + ".tmp")
    assert not tmp_file.exists()
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_state_lock_can_be_acquired_and_released_sequentially():
    """Smoke test for the cross-process lock: acquiring and releasing it must
    not raise, and a later locked call still lands normally (flock locks the
    *open file description*, so nesting two separate acquisitions in one
    process would deadlock - this deliberately tests sequential use only)."""
    with run_budget._state_lock():
        pass
    run_budget.record_and_check("lint", [], override=False)
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_budget_status_is_never_gated():
    """Reading the gate must stay possible after the gate trips - that report is
    what the agent owes the user."""
    for _ in range(50):
        run_budget.record_and_check("budget", ["status"], override=False)
    state = run_budget._load_state()
    assert "test-session" not in state


def test_budget_reset_is_counted_like_any_other_call():
    """`reset` clears the counter, so exempting it would let a session step
    around the gate by resetting first."""
    run_budget.record_and_check("budget", ["reset"], override=False)
    assert run_budget._load_state()["test-session"]["count"] == 1


def test_search_is_never_gated():
    """Retrieval is reading, not iterating. Charging for a search would tax the
    one habit that lowers token cost - looking before reading."""
    for _ in range(50):
        run_budget.record_and_check("search", ["Longhorn"], override=False)
    assert "test-session" not in run_budget._load_state()


def test_search_exemption_survives_a_query_that_looks_like_a_subcommand():
    """`is_exempt` reads args[0] as a subcommand for grouped commands; for
    `search` that slot holds the user's query, so the exemption has to be
    command-level or it depends on what was searched for."""
    assert run_budget.is_exempt("search", ["status"])
    assert run_budget.is_exempt("search", ["anything at all"])
    assert not run_budget.is_exempt("publish", ["--message", "x"])


def test_version_regrade_bare_listing_is_exempt():
    """`version regrade` only reads when called with no further arguments at
    all - the listing form - which is exactly `budget: exempt_without_args`
    in its `cli_contract` record."""
    assert run_budget.is_exempt("version", ["regrade"])


def test_version_regrade_with_indices_is_counted():
    assert not run_budget.is_exempt("version", ["regrade", "1", "--impact", "high"])
    assert not run_budget.is_exempt("version", ["regrade", "3"])


def test_exempt_budget_matches_the_pre_121_skip_sets():
    """Parity test (Gitea #121 acceptance criteria): the set of commands whose
    `cli_contract` record carries `budget: exempt` is exactly the union of the
    old `SKIP_COMMANDS`/`SKIP_COMMAND_PATHS` sets this replaced, plus bare
    `version` (`version show`'s own path). `version regrade` is deliberately
    not in this set - its `exempt_without_args` budget is a third shape, and
    is checked separately below."""
    from chemenu import cli_contract

    expected = {
        "search", "doctor", "review",  # old SKIP_COMMANDS
        "budget status", "eval score", "eval sessions", "cite id", "links show",
        "version show", "version check", "version notes",
        "migrate list", "migrate status", "migrate verify",
    }
    actual = {
        path
        for path, record in cli_contract.all_records().items()
        if record.properties.budget == cli_contract.Budget.EXEMPT
    }
    assert actual == expected


def test_only_version_regrade_uses_exempt_without_args():
    from chemenu import cli_contract

    without_args = {
        path
        for path, record in cli_contract.all_records().items()
        if record.properties.budget == cli_contract.Budget.EXEMPT_WITHOUT_ARGS
    }
    assert without_args == {"version regrade"}


def test_reset_command_requires_yes():
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    with pytest.raises(typer.Exit):
        run_budget.reset_command(all_sessions=False, yes=False)
    assert "test-session" in run_budget._load_state()


def test_reset_command_clears_current_session():
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    run_budget.reset_command(all_sessions=False, yes=True)
    state = run_budget._load_state()
    assert "test-session" not in state


def test_reset_all_clears_state_file(isolated_state):
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    assert isolated_state.exists()
    run_budget.reset_command(all_sessions=True, yes=True)
    assert not isolated_state.exists()


def test_status_command_reports_count(capsys):
    run_budget.record_and_check("new", ["entity", "--name", "X"], override=False)
    run_budget.status_command()
    out = capsys.readouterr().out
    assert "Calls so far: 1" in out


# --- gate reproduced across separate processes (Gitea #110) ---
#
# `isolated_state`'s in-process monkeypatching cannot exercise the actual bug:
# `os.getppid()` is constant within one test process. These spawn a fresh
# Python subprocess per call - the same shape as Claude Code's Bash tool,
# which runs every `wikitool` invocation in a freshly initialised shell - so
# the parent pid really does differ call to call, and only a harness variable
# (standing in for `CLAUDE_CODE_SESSION_ID`) can hold the run together.
# Before the fallback chain existed, both tests below would be green *and*
# blind: every call landed in its own one-or-two-call bucket, and neither
# gate could ever see enough of one session to trip.

def _spawn_call(tmp_path, monkeypatch, *, command="lint", args=(), override=False):
    monkeypatch.setenv("CHEMENU_ROOT", str(tmp_path))
    monkeypatch.setenv("CLAUDE_CODE_SESSION_ID", "harness-run")
    monkeypatch.delenv("WIKITOOL_SESSION_ID", raising=False)
    code = (
        "from chemenu.commands import run_budget\n"
        f"run_budget.record_and_check({command!r}, {list(args)!r}, override={override!r})\n"
    )
    return subprocess.run(
        [sys.executable, "-c", code],
        cwd=config._PACKAGE_ROOT / "tools",
        capture_output=True, text=True,
    )


def test_the_iteration_budget_gate_trips_across_separate_shells(tmp_path, monkeypatch):
    for i in range(run_budget.DEFAULT_CALL_LIMIT):
        result = _spawn_call(tmp_path, monkeypatch, args=[f"--pass-{i}"])
        assert result.returncode == 0, result.stdout + result.stderr

    result = _spawn_call(tmp_path, monkeypatch, args=["--one-too-many"])
    assert result.returncode != 0
    assert "Iteration Budget Gate" in result.stdout


def test_the_loop_breaker_trips_across_separate_shells(tmp_path, monkeypatch):
    args = ["add", "--a", "X", "--b", "Y"]
    for _ in range(run_budget.DEFAULT_LOOP_WINDOW):
        result = _spawn_call(tmp_path, monkeypatch, command="xref", args=args)
        assert result.returncode == 0, result.stdout + result.stderr

    result = _spawn_call(tmp_path, monkeypatch, command="xref", args=args)
    assert result.returncode != 0
    assert "Loop-Breaker" in result.stdout


# --- a refusal exits without a traceback (Gitea #147) ---
#
# The two tests above call `run_budget.record_and_check` directly, so they
# never exercise `cli.main()` - the place `_util.fail()`'s `typer.Exit` used
# to go uncaught, since it is raised *before* `app()` starts and nothing
# outside `app()` catches it. These drive the real entry point
# (`python -m chemenu.cli`, not a `-c` snippet) so the whole `main()` codepath
# runs, and check the property the fix promises: exit 1, no traceback on
# either stream, `gate.refused` recorded, no `wikitool.call` for the refused
# invocation, and the budget counter left exactly where it was.

def _spawn_cli_call(tmp_path, trace_dir, monkeypatch, *, session_id, args):
    monkeypatch.setenv("CHEMENU_ROOT", str(tmp_path))
    monkeypatch.setenv("WIKI_TRACE_DIR", str(trace_dir))
    monkeypatch.setenv("WIKITOOL_SESSION_ID", session_id)
    monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False)
    return subprocess.run(
        [sys.executable, "-m", "chemenu.cli", *args],
        cwd=config._PACKAGE_ROOT / "tools",
        capture_output=True, text=True,
    )


def _trace_events(trace_dir, session_id):
    path = trace_dir / session_id / "trace.jsonl"
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def test_loop_breaker_refusal_has_no_traceback(tmp_path, monkeypatch):
    session_id = "test-147-loop"
    trace_dir = tmp_path / "trace"
    args = ["xref", "add", "--a", "X", "--b", "Y"]
    for _ in range(run_budget.DEFAULT_LOOP_WINDOW):
        _spawn_cli_call(tmp_path, trace_dir, monkeypatch, session_id=session_id, args=args)

    result = _spawn_cli_call(tmp_path, trace_dir, monkeypatch, session_id=session_id, args=args)
    assert result.returncode == 1
    assert result.stdout.startswith("ERROR Loop-Breaker")
    assert "Traceback" not in result.stdout
    assert "Traceback" not in result.stderr

    events = _trace_events(trace_dir, session_id)
    assert [e["event"] for e in events[-1:]] == ["gate.refused"]
    assert events[-1]["attrs"]["gate"] == "loop-breaker"
    assert sum(e["event"] == "wikitool.call" for e in events) == run_budget.DEFAULT_LOOP_WINDOW

    state = json.loads((tmp_path / "tools" / ".wikitool_session" / "budget.json").read_text())
    assert state[session_id]["count"] == run_budget.DEFAULT_LOOP_WINDOW


def test_iteration_budget_refusal_has_no_traceback(tmp_path, monkeypatch):
    session_id = "test-147-limit"
    trace_dir = tmp_path / "trace"
    for i in range(run_budget.DEFAULT_CALL_LIMIT):
        _spawn_cli_call(tmp_path, trace_dir, monkeypatch, session_id=session_id, args=["lint", f"--pass-{i}"])

    result = _spawn_cli_call(
        tmp_path, trace_dir, monkeypatch, session_id=session_id, args=["lint", "--one-too-many"]
    )
    assert result.returncode == 1
    assert result.stdout.startswith("ERROR Iteration Budget Gate")
    assert "Traceback" not in result.stdout
    assert "Traceback" not in result.stderr

    events = _trace_events(trace_dir, session_id)
    assert [e["event"] for e in events[-1:]] == ["gate.refused"]
    assert events[-1]["attrs"]["gate"] == "iteration-budget"
    assert sum(e["event"] == "wikitool.call" for e in events) == run_budget.DEFAULT_CALL_LIMIT

    state = json.loads((tmp_path / "tools" / ".wikitool_session" / "budget.json").read_text())
    assert state[session_id]["count"] == run_budget.DEFAULT_CALL_LIMIT