chemenu/tools/chemenu/commands/run_budget.py

"""Iteration/cost budget gate: a hard, code-enforced cap on how many wikitool
commands a single agent session may run before requiring explicit human
confirmation, plus a loop-breaker that trips immediately if the last few
calls are near-identical (same command + same arguments).

This closes the gap documented in AGENTS.md's "Gates" section: unlike a
prompt instruction ("stop after N steps"), this check runs
in-process on every `wikitool` invocation and cannot be skipped by the
calling agent "politely trying again". It mirrors the Mass-Update Gate
pattern (see git_publish.py / kb/concepts/workflows/Mass-Update Gate.md), but that
gate is scoped to the *size* of a single publish, while this one is scoped to
*iteration volume* across a whole session (e.g. a wiki-ingest or wiki-lint run
that could otherwise loop unbounded over many entity/concept pages).

Session scoping: a "session" is approximated by the parent process of this
CLI invocation (the agent's shell), via the WIKITOOL_SESSION_ID env var if the
caller sets one, otherwise os.getppid(). A new terminal/session therefore
starts with a fresh budget.
"""
from __future__ import annotations

import json
import os
import time
from contextlib import contextmanager
from pathlib import Path

import typer

from chemenu import cli_contract, config
from chemenu.commands._util import fail, success
from chemenu.session import session_id as _shared_session_id
from chemenu.session import session_id_source as _shared_session_id_source
from chemenu.telemetry import emit

app = typer.Typer(help="Session iteration/cost budget gate (see the tooling contract's 'Iteration and Cost Limits').")

STATE_DIR = config.ROOT / "tools" / ".wikitool_session"
STATE_FILE = STATE_DIR / "budget.json"
LOCK_FILE = STATE_DIR / "budget.lock"

# Calibration, measured in this instance rather than inherited: ~5-15 calls for
# a simple task, ~20-35 for a complex multi-tool workflow such as an ingest.
#
# The upper band used to read 15-25, taken from an industry rule of thumb (see
# kb/concepts/Iteration and Cost Limits.md, which still cites it as such). Four
# consecutive real ingests measured 24, 26, 29 and 30 calls - every one of them
# at or above the old band's ceiling while doing nothing unusual. A guideline
# that the normal case exceeds is not a guideline; it teaches an agent that the
# numbers are decorative.
#
# The limit sits well above the band on purpose. It is not a target but the
# point past which a session is presumed stuck. At 30 the ingest of 2026-08-30
# hit it on overhead alone - reading a report back, a corrected retry, checking
# the tree before publishing - which is the gate firing on the tool rather than
# on the task.
DEFAULT_CALL_LIMIT = 60

# Loop-breaker: abort if the last N calls all share the same command + args,
# even if the overall call limit hasn't been reached yet.
DEFAULT_LOOP_WINDOW = 3

# Sessions untouched for this long are dropped on the next write. Without this
# the state file grows one entry per session forever - and the getppid()
# fallback makes new keys cheap (a new shell is a new session).
SESSION_TTL_SECONDS = 7 * 24 * 3600

# Never gate the gate's *read* side, or reporting the situation to the user
# would become impossible exactly when the limit trips. `budget reset` is
# deliberately NOT exempt: it clears the counter, so exempting it would make
# the whole gate a formality an agent could step around by resetting first.
# It is gated on `--yes` instead, the same way `publish` is.
#
# Which commands are exempt is no longer a list here at all (Gitea #121 D5,
# B8) - it is each command's own `cli_contract` record, the same `budget:`
# property `wikitool -h`'s index prints. Two commands used to need their own
# branch below because a hand-maintained pair-set cannot express them:
# `search`/`doctor`/`review` are flat commands whose first argument is a query
# or an option, not a subcommand (`_resolve_record` tries the bare command
# path before ever treating `args[0]` as one), and bare `wikitool version`
# aliases `version show`. `version regrade`'s `budget: exempt_without_args`
# is the third shape a plain exempt/counted split cannot express: the bare
# listing reads, but any index writes `CHANGES.md` and must stay counted like
# `version bump`.


def _resolve_record(command: str, args: list[str]) -> tuple["cli_contract.CommandRecord | None", bool]:
    """The `cli_contract` record this invocation maps to, and whether it was
    matched *without* consuming `args[0]` as a subcommand (a flat command, or
    the bare `version` alias) - which is what `is_exempt` needs to answer
    `version regrade`'s own arguments-dependent question correctly."""
    direct = cli_contract.get(command)
    if direct is not None:
        return direct, True
    subcommand = args[0] if args and not args[0].startswith("-") else ""
    if command == "version" and not subcommand:
        # Bare `wikitool version` (and `version --json`) is an alias for
        # `show` - see version_cmd.py's own `@app.callback()`.
        return cli_contract.get("version show"), True
    return cli_contract.get(f"{command} {subcommand}".strip()), False


def is_exempt(command: str, args: list[str]) -> bool:
    """Whether this invocation is outside the budget entirely, read from the
    matched command's own `cli_contract` record."""
    record, bare = _resolve_record(command, args)
    if record is None:
        return False
    budget = record.properties.budget
    if budget == cli_contract.Budget.EXEMPT:
        return True
    if budget == cli_contract.Budget.EXEMPT_WITHOUT_ARGS:
        # `version regrade`'s own shape: exempt only for the bare listing.
        # `bare` is only True here for a flat command with no group at all,
        # which `version regrade` is not, so the subcommand token itself is
        # the one argument the bare listing consumes.
        consumed = 0 if bare else 1
        return len(args) == consumed
    return False


def _session_id() -> str:
    return _shared_session_id()


def _session_id_source() -> str:
    return _shared_session_id_source()


def _entry_for(state: dict, session_id: str) -> dict:
    """The state entry for this session id, starting a fresh counter if the
    same id string now carries a different origin than the one that wrote it.

    Two different id spaces (a `getppid()` integer, a harness UUID, an
    explicit `WIKITOOL_SESSION_ID`) are vanishingly unlikely to collide as
    strings - but "unlikely" is not "impossible", and inheriting a stranger's
    count on collision is exactly the silent mis-key #110 exists to close.
    An entry written before this field existed carries no `source` at all and
    is treated as compatible: it keeps its count rather than being reset the
    first time this ships, which would throw away real, in-flight state.
    """
    source = _session_id_source()
    entry = state.get(session_id)
    if entry is None or (entry.get("source") is not None and entry["source"] != source):
        entry = {"count": 0, "recent": []}
        state[session_id] = entry
    entry.setdefault("source", source)
    return entry


def _load_state() -> dict:
    if not STATE_FILE.exists():
        return {}
    try:
        return json.loads(STATE_FILE.read_text())
    except (json.JSONDecodeError, OSError):
        return {}


def prune_state(state: dict, now: float, ttl: float = SESSION_TTL_SECONDS) -> dict:
    """Drop sessions whose last recorded call is older than the TTL. Entries
    written before `last_seen` existed are kept (they get a timestamp on their
    next recorded call)."""
    return {
        session_id: entry
        for session_id, entry in state.items()
        if "last_seen" not in entry or now - entry["last_seen"] <= ttl
    }


def _save_state(state: dict) -> None:
    """Write the state atomically: build the payload, write it to a sibling
    temp file, then rename it over the real file. A crash or concurrent read
    mid-write can never observe a truncated/partial JSON file this way -
    os.replace() is atomic on POSIX."""
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    payload = json.dumps(prune_state(state, time.time()), indent=2)
    tmp_file = STATE_FILE.with_suffix(STATE_FILE.suffix + ".tmp")
    tmp_file.write_text(payload)
    os.replace(tmp_file, STATE_FILE)


@contextmanager
def _state_lock():
    """Exclusive cross-process lock guarding the budget state's load-modify-
    save cycle. Without this, two `wikitool` calls racing in the same session
    (e.g. two parallel subagents) can both load count=N, both compute N+1, and
    both save - losing an increment and letting the session run past the gate
    it exists to enforce. POSIX-only (fcntl); best-effort no-op if unavailable,
    since the loop-breaker's identical-call check still degrades gracefully."""
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    try:
        import fcntl
    except ImportError:  # pragma: no cover - non-POSIX platform
        yield
        return
    with open(LOCK_FILE, "w") as lock_fh:
        fcntl.flock(lock_fh, fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_fh, fcntl.LOCK_UN)


def loop_breaker_message(call_signature: str, loop_window: int) -> str:
    return (
        f"Loop-Breaker: the last {loop_window} wikitool calls in this session were "
        f"identical ('{call_signature}'). This usually means the agent is stuck retrying "
        "the same failing operation instead of changing approach - per the tooling contract's "
        "Tool Error Contracts, that is exactly the case for stopping and escalating rather than "
        "retrying again. Stop, explain the situation to the user, and get explicit direction "
        "before continuing. Only re-run with --override-budget once the user has confirmed "
        "the repeat is intentional - never add it on the agent's own initiative."
    )


def call_limit_message(count: int, call_limit: int) -> str:
    return (
        f"Iteration Budget Gate: this session has made {count} wikitool calls, exceeding the "
        f"limit of {call_limit}. Per the tooling contract's 'Iteration and Cost Limits' section, a "
        "single task should typically need roughly 5-15 calls (simple) or 20-35 (complex multi-tool "
        "workflow like wiki-ingest/wiki-lint). This far past that band is a documented sign of poor "
        "task decomposition or a stuck loop. Stop, summarize progress and the blocker to the "
        "user, and get explicit direction before continuing. Only re-run with --override-budget "
        "once the user has approved continuing this session - never add it unprompted."
    )


def record_and_check(
    command: str,
    args: list[str],
    override: bool,
    call_limit: int = DEFAULT_CALL_LIMIT,
    loop_window: int = DEFAULT_LOOP_WINDOW,
) -> bool:
    """Record this invocation against the session budget and enforce the gate.
    Called once per process from main() before Typer dispatches to a
    subcommand, so every wikitool command is covered uniformly.

    A refused call is *not* recorded: it never ran, so counting it would keep
    inflating the number quoted back to the user on every subsequent attempt.
    The loop-breaker still trips on the next identical call, because the
    history that made it identical is already stored.

    Returns whether a slot was actually charged, so the caller knows whether
    there is anything to hand back via `refund()`.
    """
    if not command or is_exempt(command, args):
        return False

    with _state_lock():
        session_id = _session_id()
        state = _load_state()
        entry = _entry_for(state, session_id)
        recent = entry["recent"]

        call_signature = f"{command} {' '.join(args)}".strip()

        # Checked against history *before* this call is appended, so it answers
        # "were the last `loop_window` calls already identical to this one?".
        is_repeat_of_recent = (
            len(recent) >= loop_window
            and all(c == call_signature for c in recent[-loop_window:])
        )

        if not override:
            if is_repeat_of_recent:
                emit(
                    "wikitool",
                    "gate.refused",
                    {
                        "gate": "loop-breaker",
                        "command": command,
                        "args": args,
                        "call_signature": call_signature,
                        "loop_window": loop_window,
                        "count": entry["count"],
                    },
                )
                fail(loop_breaker_message(call_signature, loop_window))
            if entry["count"] + 1 > call_limit:
                emit(
                    "wikitool",
                    "gate.refused",
                    {
                        "gate": "iteration-budget",
                        "command": command,
                        "args": args,
                        "count": entry["count"] + 1,
                        "limit": call_limit,
                    },
                )
                fail(call_limit_message(entry["count"] + 1, call_limit))

        entry["count"] += 1
        recent.append(call_signature)
        entry["recent"] = recent[-max(loop_window, 10):]
        entry["last_seen"] = time.time()
        _save_state(state)
        return True


def refund() -> None:
    """Give the current session its last charged slot back.

    Called when the command declined instead of acting: a rejected argument,
    or a read-only check reporting findings (`_util.fail`, exit 1). The
    tooling contract answers a rejected argument with "fix it and retry once",
    so charging for the rejection makes the prescribed response cost two slots
    for one operation - and the budget exists to bound iteration on the wiki,
    which a call that changed nothing did not do.

    The call stays in `recent`. Repeating the same broken invocation is a real
    failure, and the loop-breaker is the instrument for it: it needs the
    history, not the counter.
    """
    with _state_lock():
        state = _load_state()
        entry = state.get(_session_id())
        if not entry or entry.get("count", 0) <= 0:
            return
        entry["count"] -= 1
        _save_state(state)


@cli_contract.record(cli_contract.CommandRecord(
    path="budget status",
    summary="Show the current session's `wikitool` call count and recent command history.",
    synopsis=(cli_contract.Variant(usage="budget status"),),
    properties=cli_contract.Properties(
        effect=cli_contract.Effect.READ,
        idempotent=cli_contract.Idempotent.YES,
        atomic="Read-only",
        budget=cli_contract.Budget.EXEMPT,
    ),
    notes="Recent command history is never counted against the budget. Never fails. Safe to "
    "retry freely.",
    failures=(),
))
def status_command():
    """Show the current session's call count and recent command history."""
    state = _load_state()
    entry = state.get(_session_id())
    typer.echo(f"Session: {_session_id()} (from {_session_id_source()})")
    if not entry:
        success("No recorded calls yet for this session.")
        return
    typer.echo(f"Calls so far: {entry['count']} (limit {DEFAULT_CALL_LIMIT})")
    typer.echo("Recent calls:")
    for c in entry["recent"]:
        typer.echo(f"  - {c}")


def reset_message() -> str:
    return (
        "`budget reset` clears the Iteration Budget Gate - the check that exists to stop a "
        "session looping or sprawling unnoticed. Resetting it on the agent's own initiative "
        "would make the gate advisory, which is exactly what it was built not to be. Stop, "
        "summarize what the session has done so far and why it needs more calls, and only "
        "re-run with --yes once the user has approved continuing."
    )


@cli_contract.record(cli_contract.CommandRecord(
    path="budget reset",
    summary="Clear the current session's (or every session's) iteration budget state.",
    synopsis=(cli_contract.Variant(usage="budget reset --yes [--all]"),),
    properties=cli_contract.Properties(
        effect=cli_contract.Effect.WRITE,
        idempotent=cli_contract.Idempotent.NO,
        atomic="Read/rewrite of one JSON file (or its deletion, with `--all`)",
        budget=cli_contract.Budget.COUNTED,
    ),
    notes="Requires `--yes`: clearing the counter is itself a way around the gate, so it needs "
    "the same explicit human approval",
    failures=(cli_contract.Failure(
        label="",
        cause="`--yes` not passed",
        reaction="Get the user's approval, then re-run with `--yes`",
    ),),
))
def reset_command(
    all_sessions: bool = typer.Option(
        False, "--all", help="Clear every session's budget, not just the current one."
    ),
    yes: bool = typer.Option(
        False,
        "--yes",
        "-y",
        help="Confirm clearing the budget. Only pass this after a human has approved "
        "continuing the session - never set it automatically to work around the gate.",
    ),
):
    """Clear the current session's (or all sessions') recorded budget."""
    if not yes:
        fail(reset_message())
    with _state_lock():
        if all_sessions:
            if STATE_FILE.exists():
                STATE_FILE.unlink()
            success("Cleared budget state for all sessions.")
            return
        state = _load_state()
        if state.pop(_session_id(), None) is not None:
            _save_state(state)
    success(f"Cleared budget state for session {_session_id()}.")


app.command("status")(status_command)
app.command("reset")(reset_command)