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 / wiki/concepts/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 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.
#
# `eval score` and `eval sessions` read a trace and re-run lint's checks in
# process. Reading back what a session already did is not iteration on the wiki,
# and charging for it would discourage checking one's own work.
# `version show`/`check`/`notes` only read - `VERSION`, the release stamp, the
# changelog, or a remote release feed. `version bump` writes two files and
# stays counted like every other mutation.
SKIP_COMMAND_PATHS = {
    ("budget", "status"),
    ("eval", "score"),
    ("eval", "sessions"),
    ("cite", "id"),
    # Retrieval, like `search`: an agent that has to ration looking up what
    # points at a page starts guessing instead - and under authored directional
    # edges this is the *only* way to ask that question.
    ("links", "show"),
    ("version", "show"),
    ("version", "check"),
    ("version", "notes"),
    # Bare `wikitool version` (and `version --json`) is an alias for `show`;
    # the subcommand slot is empty, so it needs its own entry to be exempt
    # alongside the command it delegates to.
    ("version", ""),
    # `migrate list/status/verify` only read - the migration documents, the KB
    # state file, and git history. `verify` especially: a migration runs it
    # once per unit by design, and charging for the check would push an agent
    # toward skipping the one step that catches a dropped reference.
    # `migrate done`/`baseline` write the state file and stay counted.
    ("migrate", "list"),
    ("migrate", "status"),
    ("migrate", "verify"),
    # `upstream verify` only reads two git revisions and reports what changed -
    # the same argument as `migrate verify`: a check that costs budget is one
    # an agent starts skipping. `upstream merge` stays counted: it mutates the
    # branch and can leave an open merge behind on refusal, so it belongs on
    # the non-idempotent list (AGENTS.md's tool error contract) rather than
    # the exempt one.
    ("upstream", "verify"),
}

# Commands exempt regardless of their first argument, because that argument is
# a query rather than a subcommand. `search` is here because retrieval is
# reading, not iterating: the budget exists to stop an agent looping over the
# wiki's *state*, and charging for a search would penalise the one habit that
# lowers cost - looking before reading. `doctor` is here for the same reason:
# it only reads and reports, never mutates anything. Every command that
# mutates anything stays counted.
SKIP_COMMANDS = {"search", "doctor"}


def is_exempt(command: str, args: list[str]) -> bool:
    """Whether this invocation is outside the budget entirely."""
    if command in SKIP_COMMANDS:
        return True
    subcommand = args[0] if args and not args[0].startswith("-") else ""
    return (command, subcommand) in SKIP_COMMAND_PATHS


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


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


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 = state.setdefault(session_id, {"count": 0, "recent": []})
        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)


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."
    )


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)