sts2-bot/run.py

#!/usr/bin/env python3
"""
run.py -- the observe -> decide -> act loop.

The loop is strictly closed. Playing a card removes it from hand and shifts
every later index, so we re-read the state after every single action. We never
precompute an action list.

usage:
    python3 run.py --dry-run --steps 5     # show decisions, touch nothing
    python3 run.py --steps 40              # actually play
    python3 run.py --steps 40 --no-jev     # heuristics only
"""

from __future__ import annotations

import argparse
import atexit
import json
import os
import pathlib
import sys
import time
import uuid

import brain
import facts as F
import sts2
from jev import JevClient, JevError, answer_record

DECK_FILE = pathlib.Path("deck.json")

# Where the game writes its own run records. This is the OUTCOME source: it
# carries win/loss, the killer, the seed and the final deck. Override with
# STS2_HISTORY_DIR when the steam id differs.
HISTORY_DIR = pathlib.Path(os.environ.get(
    "STS2_HISTORY_DIR",
    "~/Library/Application Support/SlayTheSpire2/steam/"
    "76561198141226155/modded/profile1/saves/history")).expanduser()

# One id per process, written into every trace row. A decision row and the model
# answers that produced it must be joinable, and a bare wall-clock time is not
# an identifier: JEV_TRACE timestamps only to the second, so several calls share
# a timestamp and nothing in it names the run, step or action. The random suffix
# keeps two processes started in the same second apart.
SESSION_ID = f"{time.strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4().hex[:6]}"


class RecordingClient:
    """
    Wraps a JevClient so the last call's questions and answers travel with the
    decision row they produced.

    This makes the log ATTRIBUTABLE. `JEV_TRACE` alone cannot be joined to a
    decision -- it records a time to the second and the answers, with no run,
    step or action -- so nothing in it can be traced back to a run.

    What the join supports is arm-level analysis: this decision belongs to this
    session, and the session's outcome is the run file, which is what an A/B or
    a delayed-reward model needs. It does NOT say whether an individual answer
    was correct. One run result attached to one step cannot label that step;
    per-decision accuracy needs replay, expert judgement, or a ground truth the
    code can verify on its own (lethal, legality, affordability).
    """

    def __init__(self, inner: JevClient) -> None:
        self._inner = inner
        self.last: dict | None = None

    @property
    def model(self) -> str:
        return self._inner.model

    def __repr__(self) -> str:  # keep the key redacted
        return f"RecordingClient({self._inner!r})"

    def ask(self, state, questions, model=None):
        response = self._inner.ask(state, questions, model)
        self.last = {
            "model": response.model,
            "latency_s": round(response.latency_s, 3),
            "questions": questions,
            "answers": {qid: answer_record(a)
                        for qid, a in response.answers.items()},
        }
        return response


def history_snapshot() -> set[str]:
    """Run-record filenames present right now. Diffed to attribute a session."""
    try:
        return {p.name for p in HISTORY_DIR.iterdir() if p.suffix == ".run"}
    except OSError:
        return set()


def run_outcome(name: str) -> dict:
    """The outcome fields of one game run record, or a stub if it is unreadable."""
    try:
        data = json.loads((HISTORY_DIR / name).read_text())
    except (OSError, json.JSONDecodeError):
        return {"file": name}
    players = data.get("players") or [{}]
    points = data.get("map_point_history") or []
    return {
        "file": name,
        "win": data.get("win"),
        "killed_by": str(data.get("killed_by_encounter") or "").replace(
            "ENCOUNTER.", ""),
        "seed": data.get("seed"),
        "deck_size": len(players[0].get("deck") or []),
        "map_points": sum(len(a) for a in points if isinstance(a, list)),
        "run_time_s": data.get("run_time"),
    }


def session_record(session_id: str, started_at: str, ended_at: str, steps: int,
                   sources: dict, produced) -> dict:
    """
    session id -> the run record(s) that session produced.

    Nothing else writes this mapping: the game's history file does not know our
    session id, and the A/B harness only maps a policy to a filename. Without it
    a decision cannot be attributed to the run it belongs to, so no arm-level
    outcome analysis (A/B, delayed reward) is possible.

    It is NOT a per-decision accuracy label. One run result attached to a step
    says nothing about whether that step's answer was correct; calibrating a
    gate needs labels for the individual answers.
    """
    return {
        "session": session_id,
        "started": started_at,
        "ended": ended_at,
        "steps": steps,
        "sources": sources,
        "runs": [run_outcome(name) for name in sorted(produced)],
    }


def decision_record(step: int, state_type: str, run_state, decision,
                    error: str | None, client) -> dict:
    """
    One JSONL row per decided action, carrying the model's answers with it.

    `jev` is None for a decision that did not ask the model (code paths and
    fallbacks), so an absent entry is a fact about the decision, not a gap.
    """
    return {
        "session": SESSION_ID,
        "step": step,
        "state_type": state_type,
        "run": run_state,
        "event": "decide",
        "source": decision.source,
        "action": decision.action,
        "params": decision.params,
        "reason": decision.reason,
        "confidence": decision.confidence,
        "error": error,
        "jev": getattr(client, "last", None),
    }


# Known starting decks, used only until the first combat exposes the real one.
# Source: the character_select state, which lists starting_deck per character.
STARTING_DECKS: dict[str, dict[str, int]] = {
    "The Ironclad": {"Strike": 5, "Defend": 4, "Bash": 1},
    "The Silent": {"Strike": 5, "Defend": 5, "Neutralize": 1, "Survivor": 1},
    "The Regent": {"Strike": 4, "Defend": 4, "Falling Star": 1, "Venerate": 1},
}


def load_deck() -> dict | None:
    if DECK_FILE.exists():
        try:
            return json.loads(DECK_FILE.read_text())
        except (json.JSONDecodeError, OSError):
            return None
    return None


def save_deck(deck: dict) -> None:
    DECK_FILE.write_text(json.dumps(deck, indent=2, sort_keys=True))


def observe() -> dict:
    return sts2.state()


def preflight(obs: dict) -> str | None:
    """
    Detect states the bot cannot proceed from, so a session does not silently
    burn its whole step budget doing nothing.

    Two that actually bit:
      * A pending Timeline epoch leaves the main menu with only settings/quit,
        and the mod REFUSES to automate the reveal. A whole A/B arm ran with 0
        decisions before this was noticed.
      * A parked `game_over` screen blocks every later session. Dismiss it.
    """
    if obs.get("state_type") == "game_over":
        try:
            sts2.act("menu_select", option="main_menu")
            # The dismissal is not instant. Without this wait the loop reads the
            # state again, still sees game_over, and stops the session at step 1
            # -- measured, one whole session made 0 decisions.
            time.sleep(1.5)
            print("dismissed a parked game-over screen")
        except sts2.Sts2Error as exc:
            return f"BLOCKED: parked on game_over and could not dismiss it: {exc}"
        return None

    if obs.get("state_type") != "menu" or obs.get("menu_screen") != "main":
        return None

    options = obs.get("options") or []
    names = {o if isinstance(o, str) else o.get("name") for o in options}
    if names & {"singleplayer", "continue"}:
        return None

    blocked = obs.get("blocked_options") or []
    for entry in blocked:
        if isinstance(entry, dict) and entry.get("reason") == "manual_epoch_reveal_required":
            pending = ", ".join(entry.get("pending_epoch_ids") or [])
            return (
                "BLOCKED: the Timeline has unrevealed epochs (" + pending + ").\n"
                "The mod refuses to automate this by design, and the main menu\n"
                "offers only settings/quit until it is done.\n"
                "  -> Open the Timeline IN GAME and reveal the epoch by hand,\n"
                "     then re-run."
            )

    return f"BLOCKED: main menu offers only {sorted(n for n in names if n)}; cannot start a run."


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=20)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--no-jev", action="store_true")
    ap.add_argument("--pause", type=float, default=0.6,
                    help="seconds to wait after each action")
    ap.add_argument("--stuck-seconds", type=float, default=25.0,
                    help="give up if the state does not change for this long")
    ap.add_argument("--max-duplicate-waits", type=int, default=3,
                    help="how many times to suppress an identical repeated action "
                         "on an unchanged state before trying it again")
    ap.add_argument("--max-jev-errors", type=int, default=8,
                    help="abort the session after this many consecutive model "
                         "failures, so a network outage does not silently "
                         "produce heuristic-only data")
    ap.add_argument("--stop-on-run-end", action="store_true",
                    help="end the session when the run ends instead of starting "
                         "a fresh one. REQUIRED for A/B work: without it a "
                         "session can contain several runs and the results "
                         "cannot be attributed to an arm.")
    ap.add_argument("--capture-dir", default="capture")
    ap.add_argument("--card-skip-policy", choices=("jev", "combined"), default=None,
                    help="how card-reward skips are decided (default: brain's own)")
    args = ap.parse_args()

    if not sts2.is_up():
        print(f"game not reachable at {sts2.BASE}")
        return 1

    # Fail fast on a state the bot cannot leave, instead of burning the whole
    # step budget on rejected actions.
    try:
        blocker = preflight(observe())
    except sts2.Sts2Error as exc:
        print(f"cannot read state: {exc}")
        return 1
    if blocker:
        print(blocker)
        return 2

    if args.card_skip_policy:
        brain.CARD_SKIP_POLICY = args.card_skip_policy
    print(f"card skip policy: {brain.CARD_SKIP_POLICY}")
    # Grepped by ab_card_skip.sh to stamp its attribution rows with the same id.
    print(f"session: {SESSION_ID}")

    client = None
    if not args.no_jev:
        try:
            client = RecordingClient(JevClient())
            print(f"jev ready: {client!r}")
        except JevError as exc:
            print(f"jev unavailable, using heuristics only: {exc}")

    capdir = pathlib.Path(args.capture_dir)
    capdir.mkdir(exist_ok=True)
    trace_path = capdir / "decisions.jsonl"

    def trace(record: dict) -> None:
        # One JSON line per decided action. Tail it while the bot plays:
        #   tail -f capture/decisions.jsonl | jq
        record["ts"] = time.strftime("%H:%M:%S")
        record.setdefault("session", SESSION_ID)
        with trace_path.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record, sort_keys=True, default=str) + "\n")

    stats = {"code": 0, "jev": 0, "fallback": 0}
    jev_calls = 0
    jev_tokens = 0
    started = time.monotonic()

    # The outcome half of the join: which run record(s) this session produced.
    # Registered with atexit so EVERY exit path writes it -- the normal end, the
    # step cap, a stuck screen, a model-failure abort, or an exception.
    history_before = history_snapshot()
    step = 0
    started_at = time.strftime("%Y-%m-%dT%H:%M:%S")

    def write_session_row() -> None:
        try:
            row = session_record(SESSION_ID, started_at,
                                 time.strftime("%Y-%m-%dT%H:%M:%S"), step, stats,
                                 history_snapshot() - history_before)
            with (capdir / "sessions.jsonl").open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(row, sort_keys=True, default=str) + "\n")
        except OSError:
            pass  # bookkeeping must never break the exit

    atexit.register(write_session_row)

    # The card_reward state does not expose the deck, but combat states expose
    # all four piles. Snapshot composition during combat and persist it so it
    # survives a restart.
    deck_snapshot: dict | None = load_deck()
    if deck_snapshot:
        print(f"deck snapshot loaded: {deck_snapshot}")
    waits = 0
    rejected = 0
    last_sig: str | None = None
    same_state = 0
    same_state_since = time.monotonic()
    last_ok = True
    last_action_key: tuple | None = None
    duplicate_waits = 0
    jev_errors = 0
    saw_a_run = False

    for step in range(1, args.steps + 1):
        try:
            obs = observe()
        except sts2.Sts2Error as exc:
            print(f"[{step:03d}] state read failed: {str(exc)[:160]}")
            return 1

        st = obs.get("state_type")

        # One session = one run, when asked. Otherwise the bot dies, returns to
        # the menu and starts a fresh run inside the same session, so a single
        # session yields several run records and nothing can be attributed to
        # an experimental arm.
        #
        # CRITICAL: dismiss the game-over screen BEFORE stopping. Breaking
        # first left the game parked on `game_over`, so every later session saw
        # it at step 1 and stopped instantly -- the whole A/B produced nothing.
        if args.stop_on_run_end and st == "game_over":
            print(f"[{step:03d}] run ended; dismissing game-over, then stopping")
            try:
                sts2.act("menu_select", option="main_menu")
            except sts2.Sts2Error as exc:
                print(f"[{step:03d}] could not dismiss game-over: {exc}")
            time.sleep(args.pause)
            break
        if args.stop_on_run_end and st in ("monster", "elite", "boss", "map",
                                           "rewards", "card_reward", "event",
                                           "rest_site", "shop", "treasure",
                                           "card_select", "hand_select"):
            saw_a_run = True
        if args.stop_on_run_end and saw_a_run and st == "menu" and \
                obs.get("menu_screen") == "main":
            print(f"[{step:03d}] back at the main menu; run is over, stopping session")
            break

        # Guard against re-acting while the game is still animating a transition.
        # An identical state after our own action means the action has not landed
        # yet; acting again queues duplicates (e.g. three map moves in a row).
        sig = json.dumps(obs, sort_keys=True)
        if sig == last_sig:
            same_state += 1
        else:
            same_state = 0
            last_sig = sig
            same_state_since = time.monotonic()

        # Time-based, not count-based: a boss death animation plus the rewards
        # transition can easily exceed a fixed number of reads.
        unchanged_for = time.monotonic() - same_state_since
        if unchanged_for > args.stuck_seconds:
            print(f"[{step:03d}] STUCK: state unchanged for {unchanged_for:.0f}s -- stopping")
            print(json.dumps(obs, indent=2)[:900])
            break

        # Only wait when our own action actually landed and the game is still
        # animating. If the action was rejected, fall through and pick a
        # different one instead of waiting out the stuck counter.
        #
        # NOTE: this check now happens AFTER deciding, and only suppresses a
        # REPEATED action. Waiting on "state unchanged" alone blocked
        # legitimate sequences -- character select needs select-then-embark, and
        # the screen does not change between them, so the bot stalled forever.

        if st in ("monster", "elite", "boss"):
            f = F.combat_facts(obs)
            (capdir / f"live_{step:03d}_combat.json").write_text(json.dumps(obs, indent=2))
            print(f"[{step:03d}] COMBAT  {f.describe().splitlines()[0]}")
            for line in f.describe().splitlines()[1:]:
                print(f"        {line}")
            # The snapshot carries the name->count map (card identities, which
            # are the synergy signal) plus stable all-pile aggregates. Only the
            # counts decide whether the deck actually changed.
            if f.deck_counts and f.deck_counts != (deck_snapshot or {}).get("counts"):
                deck_snapshot = {"counts": f.deck_counts, "summary": f.deck_summary}
                save_deck(deck_snapshot)

        decision = None
        decide_error = None
        if isinstance(client, RecordingClient):
            # Cleared first: a decision that asks nothing (code paths, fallbacks)
            # must not inherit the previous step's answers in the log.
            client.last = None
        try:
            decision = brain.decide(obs, client, deck_snapshot)
            jev_errors = 0
        except JevError as exc:
            jev_errors += 1
            print(f"[{step:03d}] jev error ({jev_errors}): {str(exc)[:140]}")
            decide_error = f"JevError: {str(exc)[:160]}"
            if jev_errors >= args.max_jev_errors:
                print(f"[{step:03d}] ABORT: {jev_errors} consecutive model failures. "
                      f"The run would continue on heuristics alone, which is not "
                      f"the data we want. Retry this session.")
                return 3
            # Fall back WITHOUT the model. Passing the client again just retries
            # the same failing request -- measured, a DNS blip re-raised out of
            # the "fallback" and killed the session.
            decision = brain.simple_decision(obs, None, deck_snapshot)
        except Exception as exc:  # noqa: BLE001
            # A network blip or an unexpected shape must not end the run.
            print(f"[{step:03d}] unexpected error in decide(): "
                  f"{type(exc).__name__}: {str(exc)[:160]}")
            decide_error = f"{type(exc).__name__}: {str(exc)[:160]}"
            try:
                decision = brain.simple_decision(obs, None, deck_snapshot)
            except Exception as inner:  # noqa: BLE001
                print(f"[{step:03d}] fallback also failed: {inner}")
                decision = None

        if decision is None:
            trace({"step": step, "state_type": st,
                   "event": "no_decision", "error": decide_error})
            print(f"[{step:03d}] no decision for state_type={st!r} -- stopping")
            print(json.dumps(obs, indent=2)[:800])
            break

        # Between turns there is nothing to do but look again.
        if decision.action == "__wait__":
            waits += 1
            print(f"[{step:03d}] WAIT    {decision.reason}")
            time.sleep(args.pause)
            continue

        # Suppress DUPLICATE actions on an unchanged state, but only for a
        # bounded number of reads. Proposing a DIFFERENT action is always
        # allowed, which is what makes select-then-embark and multi-purchase
        # shops work. The bound matters because some actions legitimately need
        # repeating (multi-line Ancient dialogue) and some transitions are just
        # slow -- waiting forever on those stalls the run.
        action_key = (decision.action, json.dumps(decision.params, sort_keys=True))
        if same_state == 0:
            duplicate_waits = 0

        if same_state > 0 and last_ok and action_key == last_action_key:
            duplicate_waits += 1
            if duplicate_waits <= args.max_duplicate_waits:
                waits += 1
                print(f"[{step:03d}] WAIT    same action on unchanged state ({same_state})")
                time.sleep(args.pause)
                continue
            print(f"[{step:03d}] RETRY   repeating {decision.action} after "
                  f"{duplicate_waits} waits on an unchanged state")
            duplicate_waits = 0

        stats[decision.source] = stats.get(decision.source, 0) + 1
        print(f"[{step:03d}] DECIDE  {decision}")
        trace(decision_record(step, st, obs.get("run"), decision, decide_error, client))
        last_action_key = action_key

        if args.dry_run:
            continue

        try:
            result = sts2.act(decision.action, **decision.params)
        except sts2.Sts2Error as exc:
            print(f"[{step:03d}] action failed: {str(exc)[:200]}")
            trace({"step": step, "event": "action_error",
                   "action": decision.action, "error": str(exc)[:200]})
            break

        if not result.ok:
            print(f"[{step:03d}] action rejected: {result.message}")
            trace({"step": step, "event": "action_rejected",
                   "action": decision.action, "message": result.message})
            last_ok = False
            rejected += 1
            if rejected >= 6:
                print(f"[{step:03d}] STUCK: {rejected} consecutive rejections -- stopping")
                print(json.dumps(obs, indent=2)[:900])
                break
            # Transient rejections while the game animates are normal -- a
            # rest-site `proceed` right after a heal is rejected for a moment
            # and then succeeds. Back off longer than the usual pause.
            time.sleep(args.pause * 3)
            continue
        else:
            rejected = 0
            last_ok = True

        time.sleep(args.pause)

    elapsed = time.monotonic() - started
    print()
    print(f"steps={step} waits={waits} elapsed={elapsed:.1f}s  sources={stats}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())