interactive-story/backend/tests/test_state_revert.py

"""Tests for undo and retry rolling back the shared `script_state`
(plan/11-state-revert-and-retry-fix.md).

Phase 14 SP4 reversed the snapshots. An action used to carry the state as
it stood before it ran, and rolling back read the snapshot off the action
being removed. Now it carries the state it left behind, and rolling back
reads that state off the node in front of it. This is the same value
reached from the other direction, and it is the only version a retry can
use, because attempts at one turn share a starting position and differ
only in their outcome.

Run from the backend dir:  python -m pytest tests/test_state_revert.py -v
"""

# Point the app at a throwaway SQLite file before importing anything that
# binds the engine at import time. `app.database` reads `AIDND_DB_PATH`
# on import.

import pytest
from fastapi import HTTPException

from app import attempts, memorybank, models
from app.database import Base, SessionLocal, engine
from app.routers import adventures


@pytest.fixture()
def db():
    Base.metadata.create_all(bind=engine)
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()
        Base.metadata.drop_all(bind=engine)
        adventures.turns._active_turns.clear()


def _make_adventure(db, script_state):
    user = models.User(is_guest=False)
    db.add(user)
    db.flush()
    adv = models.Adventure(user_id=user.id, title="T", script_state=script_state)
    db.add(adv)
    db.flush()
    return user, adv


def _add(db, adv, index, type_, text="x", state_after=None):
    a = models.Action(
        adventure_id=adv.id, index=index, type=type_, text=text,
        state_after=state_after,
    )
    db.add(a)
    db.flush()
    return a


def _forget_snapshots(db, adv):
    """Blank every outcome, the way a row written before SP4 looks.

    Straight SQL, because `tree.stamp_outcome` runs on every flush precisely so
    that a node written through the ORM cannot end up without one.
    """
    db.query(models.Action).filter_by(adventure_id=adv.id).update(
        {"state_after": None, "world_state_after": None}, synchronize_session=False
    )
    db.commit()
    db.expire_all()


# ---------------------------------------------------------------- undo

def test_undo_reverts_state_to_before_the_turn(db):
    # A turn moved script_state from {gold:0} to {gold:10}. The node in
    # front of the turn records where it started. The current state is
    # the mutated one.
    user, adv = _make_adventure(db, {"gold": 10})
    _add(db, adv, 0, "start", state_after={"gold": 0})
    _add(db, adv, 1, "do", state_after={"gold": 0})
    _add(db, adv, 2, "ai", state_after={"gold": 10})
    db.commit()

    adventures.undo_turn(adv.id, db=db, adventure=adv)

    assert adv.script_state == {"gold": 0}
    assert [a.type for a in adv.actions] == ["start"]


def test_undo_of_bare_continue_uses_the_node_in_front(db):
    # A "continue" turn has no player action, so the opening is what the story
    # falls back to.
    user, adv = _make_adventure(db, {"gold": 5})
    _add(db, adv, 0, "start", state_after={"gold": 0})
    _add(db, adv, 1, "ai", state_after={"gold": 5})
    db.commit()

    adventures.undo_turn(adv.id, db=db, adventure=adv)

    assert adv.script_state == {"gold": 0}
    assert [a.type for a in adv.actions] == ["start"]


def test_undo_leaves_state_untouched_when_snapshot_missing(db):
    # A row the SP4 migration could not derive an outcome for: leave the live
    # state alone rather than resetting it to nothing.
    user, adv = _make_adventure(db, {"gold": 10})
    _add(db, adv, 0, "start")
    _add(db, adv, 1, "do")
    _add(db, adv, 2, "ai")
    _forget_snapshots(db, adv)

    adventures.undo_turn(adv.id, db=db, adventure=adv)

    assert adv.script_state == {"gold": 10}


def test_undo_raises_when_nothing_to_undo(db):
    user, adv = _make_adventure(db, {})
    _add(db, adv, 0, "start")
    db.commit()
    with pytest.raises(HTTPException) as exc:
        adventures.undo_turn(adv.id, db=db, adventure=adv)
    assert exc.value.status_code == 400


def test_undo_blocked_by_active_turn_lock(db):
    user, adv = _make_adventure(db, {})
    _add(db, adv, 0, "start")
    _add(db, adv, 1, "ai", state_after={})
    db.commit()

    adventures.turns.acquire_turn_lock(adv.id)  # a turn is "generating"
    try:
        with pytest.raises(HTTPException) as exc:
            adventures.undo_turn(adv.id, db=db, adventure=adv)
        assert exc.value.status_code == 409
        # The failed undo must not have released someone else's lock.
        assert adv.id in adventures.turns._active_turns
    finally:
        adventures.turns._active_turns.discard(adv.id)


def test_undo_prunes_memory_covering_removed_actions(db):
    user, adv = _make_adventure(db, {})
    for i in range(4):
        _add(db, adv, i, "ai" if i % 2 else "do", state_after={})
    # A memory summarizing actions up to index 3, which undo will delete.
    covering = models.Memory(adventure_id=adv.id, text="m", source_start=0, source_end=3)
    keep = models.Memory(adventure_id=adv.id, text="k", source_start=0, source_end=1)
    db.add_all([covering, keep])
    db.commit()

    adventures.undo_turn(adv.id, db=db, adventure=adv)  # removes indexes 2 & 3

    texts = {m.text for m in adv.memories}
    assert texts == {"k"}


# -------------------------------------------------------- withdrawing a node

def test_forget_node_withdraws_only_what_that_node_produced(db):
    """Phase 14 SP3: a memory attaches to the node where its block ends, so
    removing a node is a lookup rather than a scan for memories that
    reference actions the story no longer has."""
    user, adv = _make_adventure(db, {})
    _add(db, adv, 0, "do")
    second = _add(db, adv, 1, "ai")
    db.add_all([
        models.Memory(adventure_id=adv.id, text="hangs off node 1",
                      source_start=0, source_end=1),
        models.Memory(adventure_id=adv.id, text="hangs off node 0",
                      source_start=0, source_end=0),
    ])
    db.commit()

    removed = memorybank.forget_node(db, adv, second)
    db.commit()
    db.refresh(adv)  # expire_on_commit=False: reload the memories collection

    assert removed == 1
    assert {m.text for m in adv.memories} == {"hangs off node 0"}


# ---------------------------------------------------------------- snapshot

def test_snapshot_outcome_is_an_independent_deep_copy(db):
    _, adv = _make_adventure(db, {"nested": {"n": 1}})
    node = models.Action(adventure_id=adv.id, index=0, type="ai", text="x")
    attempts.snapshot_outcome(adv, node)
    adv.script_state["nested"]["n"] = 99
    assert node.state_after == {"nested": {"n": 1}}  # unaffected by later mutation


def test_snapshot_outcome_handles_non_dict(db):
    _, adv = _make_adventure(db, {})
    adv.script_state = None
    node = models.Action(adventure_id=adv.id, index=0, type="ai", text="x")
    attempts.snapshot_outcome(adv, node)
    assert node.state_after == {}


def test_restore_state_ignores_a_node_with_no_outcome(db):
    _, adv = _make_adventure(db, {"gold": 7})
    attempts.restore_state(adv, models.Action(adventure_id=adv.id, index=0, type="ai"))
    assert adv.script_state == {"gold": 7}
    attempts.restore_state(adv, None)
    assert adv.script_state == {"gold": 7}


# ---------------------------------------------------------------- retry

def test_retry_restores_the_state_the_turn_started_from(db, monkeypatch):
    # Retry must roll script_state back to what the node in front of the AI
    # action left behind, so regeneration does not stack output mutations
    # on top of the attempt being replaced.
    user, adv = _make_adventure(db, {"gold": 20})  # 20 = double-applied bug value
    _add(db, adv, 0, "start", state_after={"gold": 0})
    _add(db, adv, 1, "do", state_after={"gold": 10})
    _add(db, adv, 2, "ai", state_after={"gold": 20})
    db.commit()

    monkeypatch.setattr(adventures.limits, "rate_limit", lambda *a, **k: None)
    monkeypatch.setattr(adventures.turns, "check_demo_cap", lambda *a, **k: None)

    async def _noop(*a, **k):
        if False:
            yield  # make it an async generator
    monkeypatch.setattr(adventures.turns, "generate_turn", _noop)

    adventures.retry_action(adv.id, request=None, db=db, user=user, adventure=adv)

    assert adv.script_state == {"gold": 10}
    # Nothing is written until a replacement actually arrives: the attempt on
    # screen is left exactly as it was, and stays the live one.
    assert [a.type for a in adv.actions] == ["start", "do", "ai"]
    last = adv.actions[-1]
    assert last.live is True
    assert last.variant_count == 0
    assert last.state_after == {"gold": 20}  # its own outcome, untouched
    adventures.turns._active_turns.discard(adv.id)