interactive-story/backend/tests/test_state_revert.py

"""Tests for undo and retry rolling back the shared `world_state`.

The state being rolled back was the scripting engine's `script_state` until
M2 removed campaign scripting. The machinery under test — `attempts.restore_state`,
`roll_back_before`, and the per-node outcome snapshot — is unchanged; only the
column it moves has. `world_state`/`world_state_after` is now the shared state
an adventure carries, so that is what these tests exercise.

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, world_state):
    user = models.User(is_guest=False)
    db.add(user)
    db.flush()
    adv = models.Adventure(user_id=user.id, title="T", world_state=world_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, type=type_, text=text,
        world_state_after=state_after,
    )
    db.add(a)
    db.flush()
    return a


def _row_count(db, adv):
    """Every action row the adventure holds, live or not, head or no head."""
    return db.query(models.Action).filter_by(adventure_id=adv.id).count()


def _all_types(db, adv):
    """The retained story in depth order, which is not the same as the story
    being told once the head has moved back behind the tip."""
    rows = (
        db.query(models.Action)
        .filter_by(adventure_id=adv.id)
        .order_by(models.Action.depth, models.Action.id)
        .all()
    )
    return [a.type for a in rows]


def _retrievable(db, adv):
    """The memories the story can currently reach, read through the same clause
    `memorybank` retrieves with — which is capped at the active head."""
    from app.context import lineage

    rows = (
        db.query(models.Memory)
        .filter(
            models.Memory.adventure_id == adv.id,
            lineage.path_of(db, adv).clause(models.Memory),
        )
        .all()
    )
    return {m.text for m in rows}


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 world_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.
    #
    # M3 rewrote what the second half of this test asserts. Undo used to delete
    # the turn, so the story was short afterwards because the rows were gone.
    # It now moves the head, so the story is short because it is being read
    # from somewhere earlier — and the rows are all still there. The state
    # assertion is unchanged, because `attempts.restore_state` is unchanged:
    # the state still comes off the node the story now ends on.
    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()
    before = _row_count(db, adv)

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

    assert adv.world_state == {"gold": 0}
    # What the story now tells.
    assert [a.type for a in page.actions] == ["start"]
    assert adv.head_depth == 0
    # What it still holds. Zero accepted turns deleted, which is the M3
    # invariant this file is the closest test to.
    assert _row_count(db, adv) == before
    assert _all_types(db, adv) == ["start", "do", "ai"]
    assert page.can_redo is True


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()

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

    assert adv.world_state == {"gold": 0}
    assert [a.type for a in page.actions] == ["start"]
    assert _all_types(db, adv) == ["start", "ai"]


def test_redo_puts_back_the_state_the_turn_left_behind(db):
    """The other half of the same mechanism: undo and redo restore the same
    snapshot from opposite directions, because it belongs to the node rather
    than to the direction of travel."""
    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)
    page = adventures.redo_turn(adv.id, db=db, adventure=adv)

    assert adv.world_state == {"gold": 10}
    assert [a.type for a in page.actions] == ["start", "do", "ai"]
    assert adv.head_depth == 2
    assert page.can_redo is False


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.world_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_stops_retrieving_a_memory_without_deleting_it(db):
    """M3 rewrote this test. Undo used to prune the memories covering the turns
    it deleted, because those turns were gone and a summary of them described
    story the adventure no longer had.

    Nothing is deleted now, and nothing needs pruning either. A memory carries
    the coordinate of the node its block ends on, so one derived from a turn
    that is now past the head falls outside the head-capped path clause and
    stops being retrievable — and becomes eligible again on Redo, without having
    been deleted and re-embedded. That is `STORY-BRANCH-SEMANTICS.md` §33
    holding as a consequence of the head rather than as its own mechanism.
    """
    user, adv = _make_adventure(db, {})
    for i in range(4):
        _add(db, adv, i, "ai" if i % 2 else "do", state_after={})
    db.commit()
    # A memory ending on the turn undo will step behind, and one ending before
    # it. The coordinate is what the clause reads; `source_*` only says which
    # stretch the summarizer covered.
    covering = models.Memory(
        adventure_id=adv.id, text="m", source_start=0, source_end=3,
        branch_id=adv.head_branch_id, depth=3,
    )
    keep = models.Memory(
        adventure_id=adv.id, text="k", source_start=0, source_end=1,
        branch_id=adv.head_branch_id, depth=1,
    )
    db.add_all([covering, keep])
    db.commit()

    adventures.undo_turn(adv.id, db=db, adventure=adv)  # head moves to depth 1

    assert _retrievable(db, adv) == {"k"}
    # Still on disk, still embedded, still attached to the adventure.
    assert {m.text for m in adv.memories} == {"k", "m"}

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

    assert _retrievable(db, adv) == {"k", "m"}


# -------------------------------------------------------- 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, type="ai", text="x")
    attempts.snapshot_outcome(adv, node)
    adv.world_state["nested"]["n"] = 99
    assert node.world_state_after == {"nested": {"n": 1}}  # unaffected by later mutation


def test_snapshot_outcome_handles_non_dict(db):
    _, adv = _make_adventure(db, {})
    adv.world_state = None
    node = models.Action(adventure_id=adv.id, type="ai", text="x")
    attempts.snapshot_outcome(adv, node)
    assert node.world_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, type="ai"))
    assert adv.world_state == {"gold": 7}
    attempts.restore_state(adv, None)
    assert adv.world_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()


    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.world_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 len(attempts.group(db, last)) == 1  # no sibling was filed
    assert last.world_state_after == {"gold": 20}  # its own outcome, untouched
    adventures.turns._active_turns.discard(adv.id)