interactive-story/backend/tools/m9_migration_proof.py

"""M9's migration claim, proved against a database M8's own code wrote.

    # from the M8 worktree, using M8's interpreter:
    python -m tools.m9_migration_proof build <db_path>

    # from the M9 tree, using M9's interpreter:
    python -m tools.m9_migration_proof open <db_path>

M9 claims to add no schema change. `git diff` proves that nothing in
`migrations.py` or `models.py` moved, which is necessary and not sufficient: a
migration can also be *missing*, and the failure then is a database that opens
and quietly answers wrongly. The M8 report set the standard here — a database
created by today's code and read by today's code proves nothing — so the
campaign below is built by a server running the signed M8 commit, from a git
worktree, and read back by M9.

`build` writes a campaign that touches every family M9 changed the handling of:
story with an alternate take, a Save Point, a manual state correction, memories
and a summary, imported knowledge including a disabled and a narrator-only
source, and per-turn context snapshots. It prints what it wrote, as JSON.

`open` opens that file with the current code, runs the migration path, and
checks every one of those against what `build` reported. It also asserts the
schema version did not move and that a second open is a no-op, which is what
"no migration" means in practice: the stamp is the same number before and after.

Neither half imports anything from the other. What crosses is the database file
and one JSON report on stdout, which is the only way the two builds can be made
to talk without one of them importing the other's code.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "tests"))


def _app(db_path: str):
    """Imports the application against `db_path`. Must run before any app import."""
    os.environ["AIDND_DB_PATH"] = db_path
    os.environ.pop("AIDND_DATABASE_URL", None)
    os.environ.pop("DATABASE_URL", None)

    from fastapi import Depends
    from fastapi.testclient import TestClient

    from app import auth, limits, memorybank, models
    from app.database import Base, SessionLocal, engine, get_db
    from app.main import app
    from app.routers import adventures
    from fakes import ScriptedProvider, state_block

    class Stub:
        async def complete(self, system, prompt, **kwargs):
            return "A memory of what had happened by then."

        async def embed(self, texts):
            return [[1.0, 0.5, 0.25] for _ in texts]

    adventures.turns.OpenAICompatibleProvider = ScriptedProvider
    memorybank.embedding_provider = lambda s: Stub()
    memorybank.summary_provider = lambda s: Stub()
    limits.check_row_cap = lambda *a, **k: None
    return {
        "Base": Base, "SessionLocal": SessionLocal, "engine": engine,
        "models": models, "app": app, "auth": auth, "get_db": get_db,
        "Depends": Depends, "TestClient": TestClient,
        "ScriptedProvider": ScriptedProvider, "state_block": state_block,
        "memorybank": memorybank,
    }


def _client(ctx, user_id: int):
    ctx["app"].dependency_overrides[ctx["auth"].get_current_user] = (
        lambda db=ctx["Depends"](ctx["get_db"]): db.get(ctx["models"].User, user_id)
    )
    return ctx["TestClient"](ctx["app"])


# ------------------------------------------------------------------- building

def build(db_path: str) -> dict:
    ctx = _app(db_path)
    ctx["Base"].metadata.create_all(bind=ctx["engine"])
    models, SessionLocal = ctx["models"], ctx["SessionLocal"]

    db = SessionLocal()
    try:
        user = models.User(is_guest=False, email="m9mig@example.com")
        db.add(user)
        db.flush()
        db.add(models.Settings(
            user_id=user.id, model="m8-model", embedding_model="stub",
            context_token_budget=4000, max_output_tokens=400,
        ))
        adventure = models.Adventure(
            user_id=user.id, title="Built by M8", auto_summarize=True,
            memory_bank_enabled=True,
            campaign_canon={"rules": ["The dead do not return."]},
        )
        db.add(adventure)
        db.flush()
        db.add(models.Action(
            adventure_id=adventure.id, type="start",
            text="Aldric sits in the Crooked Lantern with Mara.",
        ))
        db.commit()
        adv_id, user_id = adventure.id, user.id
    finally:
        db.close()

    client = _client(ctx, user_id)
    upload = client.post(
        f"/api/adventures/{adv_id}/knowledge",
        files={"file": ("canon.md", (
            "# Westhaven\n\n## The Old Abbey\n\nThe abbey crypt is sealed.\n"
        ).encode(), "text/markdown")},
        data={"classification": "canon", "always_include": "true"},
    )
    assert upload.status_code == 201, upload.text[:300]
    hidden = client.post(
        f"/api/adventures/{adv_id}/knowledge",
        files={"file": ("secret.md", (
            "# The seal\n\nIt was broken once, sixty years ago.\n"
        ).encode(), "text/markdown")},
        data={"classification": "canon", "visibility": "hidden"},
    )
    assert hidden.status_code == 201, hidden.text[:300]
    disabled = client.post(
        f"/api/adventures/{adv_id}/knowledge",
        files={"file": ("draft.md", b"# Draft\n\nAn earlier version.\n",
                        "text/markdown")},
        data={"classification": "reference"},
    )
    assert disabled.status_code == 201, disabled.text[:300]
    assert client.patch(
        f"/api/adventures/{adv_id}/knowledge/{disabled.json()['id']}",
        json={"enabled": False},
    ).status_code == 200

    state_block = ctx["state_block"]
    for turn in range(1, 8):
        ctx["ScriptedProvider"].replies = [
            f"The rain keeps on, and Mara says nothing for a while. [{turn}]\n"
            + state_block([{"type": "add_fact", "predicate": "tally",
                            "value": turn * 10, "fact_id": f"tally-{turn * 10}"}])
        ]
        response = client.post(f"/api/adventures/{adv_id}/actions",
                               json={"type": "do", "text": f"ask about turn {turn}"})
        assert response.status_code == 200, response.text[:300]
        if turn == 3:
            assert client.post(f"/api/adventures/{adv_id}/retry").status_code == 200
            point = client.post(f"/api/adventures/{adv_id}/checkpoints",
                                json={"name": "Third turn", "note": "A position."})
            assert point.status_code == 201, point.text[:300]

    correction = client.post(f"/api/adventures/{adv_id}/state/corrections", json={
        "events": [{"type": "add_fact", "predicate": "keeper", "value": "Mara",
                    "fact_id": "keeper"}],
        "note": "Established in play.",
    })
    assert correction.status_code == 201, correction.text[:300]

    import asyncio
    asyncio.run(ctx["memorybank"].run_post_turn(adv_id))

    # Two Undos, so the head is behind the retained tip when M9 opens it.
    for _ in range(2):
        assert client.post(f"/api/adventures/{adv_id}/undo").status_code == 200

    report = _describe(ctx, client, adv_id)
    ctx["app"].dependency_overrides.clear()
    return report


# -------------------------------------------------------------------- reading

def _describe(ctx, client, adv_id: int) -> dict:
    """Everything the other build has to agree with, read through the API."""
    models, SessionLocal = ctx["models"], ctx["SessionLocal"]
    page = client.get(f"/api/adventures/{adv_id}").json()
    with SessionLocal() as db:
        adventure = db.get(models.Adventure, adv_id)
        version = db.execute(_pragma()).scalar()
        counts = {
            name: db.query(model).filter(model.adventure_id == adv_id).count()
            for name, model in (
                ("actions", models.Action), ("memories", models.Memory),
                ("summaries", models.Summary), ("checkpoints", models.Checkpoint),
                ("state_events", models.StateEvent),
                ("state_proposals", models.StateProposal),
                ("knowledge_sources", models.KnowledgeSource),
                ("knowledge_chunks", models.KnowledgeChunk),
            )
        }
        head = {"branch_id": adventure.head_branch_id, "depth": adventure.head_depth}
        snapshots = db.query(models.Action).filter(
            models.Action.adventure_id == adv_id,
            models.Action.context_snapshot.isnot(None),
        ).count()
    return {
        "adventure_id": adv_id,
        "schema_version": version,
        "title": page["title"],
        "canon_rules": page["canon_rules"],
        "transcript": [a["text"] for a in page["actions"]],
        "can_undo": page["can_undo"],
        "can_redo": page["can_redo"],
        "head": head,
        "counts": counts,
        "snapshot_rows": snapshots,
        "state": client.get(f"/api/adventures/{adv_id}/state").json()["document"],
        "checkpoints": sorted(
            (c["name"], c["depth"])
            for c in client.get(f"/api/adventures/{adv_id}/checkpoints").json()
        ),
        "knowledge": sorted(
            (k["original_filename"], k["classification"], k["enabled"],
             k["visibility"], k["always_include"], k["content_hash"],
             k["index_state"])
            for k in client.get(f"/api/adventures/{adv_id}/knowledge").json()
        ),
        "events": sorted(
            (e["event_type"], e["source"], json.dumps(e["payload"], sort_keys=True))
            for e in client.get(
                f"/api/adventures/{adv_id}/state/events?limit=500").json()
        ),
    }


def _comparable(value):
    """`value` as it survives a JSON round trip, so the two builds compare like."""
    return json.loads(json.dumps(value, sort_keys=True, default=str))


def _pragma():
    from sqlalchemy import text

    return text("PRAGMA user_version")


def open_it(db_path: str, expected: dict) -> dict:
    """Opens an existing database with this build, and checks it against `expected`."""
    ctx = _app(db_path)
    from app import migrations

    # This is the migration run. `main` already called `bootstrap` at import.
    with ctx["engine"].begin() as conn:
        after_first = conn.execute(_pragma()).scalar()
    # And again, to prove idempotence: a second run must change nothing.
    migrations.bootstrap(ctx["engine"])
    with ctx["engine"].begin() as conn:
        after_second = conn.execute(_pragma()).scalar()

    adv_id = expected["adventure_id"]
    with ctx["SessionLocal"]() as db:
        user = db.query(ctx["models"].User).first()
        user_id = user.id
    client = _client(ctx, user_id)
    actual = _describe(ctx, client, adv_id)

    problems = []
    for key in ("title", "canon_rules", "transcript", "head", "counts",
                "snapshot_rows", "state", "checkpoints", "knowledge", "events",
                "can_undo", "can_redo"):
        # Compared through JSON, because that is how the other build's answer
        # arrived: a tuple written by `_describe` comes back as a list, and a
        # comparison that called that a difference would report ten differences
        # in a database nothing had changed.
        if _comparable(actual[key]) != _comparable(expected[key]):
            problems.append(f"{key}: expected {expected[key]!r}, got {actual[key]!r}")
    if expected["schema_version"] != after_first:
        problems.append(
            f"the schema version moved: {expected['schema_version']} -> {after_first}"
        )
    if after_first != after_second:
        problems.append(
            f"a second open migrated again: {after_first} -> {after_second}"
        )

    # And the campaign still works, rather than merely reading correctly.
    exported = client.get(f"/api/adventures/{adv_id}/export")
    if exported.status_code != 200:
        problems.append(f"export failed: {exported.status_code}")
    else:
        imported = client.post("/api/adventures/import", json=exported.json())
        if imported.status_code != 201:
            problems.append(f"round trip failed: {imported.text[:300]}")
        elif exported.json()["format"] != "ai-dnd-adventure-v3":
            problems.append("the M8 database did not export as v3")
    redo = client.post(f"/api/adventures/{adv_id}/redo")
    if redo.status_code != 200:
        problems.append(f"Redo failed on the migrated campaign: {redo.status_code}")

    ctx["app"].dependency_overrides.clear()
    return {
        "schema_version_before": expected["schema_version"],
        "schema_version_after": after_first,
        "schema_version_second_open": after_second,
        "problems": problems,
        "checked": {
            "families": 12, "snapshot_rows": actual["snapshot_rows"],
            "counts": actual["counts"],
        },
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("mode", choices=("build", "open"))
    parser.add_argument("db_path")
    parser.add_argument("--expected", help="the JSON `build` printed (open only)")
    args = parser.parse_args()

    if args.mode == "build":
        print(json.dumps(build(args.db_path), sort_keys=True))
        return 0

    expected = json.loads(Path(args.expected).read_text())
    result = open_it(args.db_path, expected)
    print(json.dumps(result, indent=2, sort_keys=True))
    return 1 if result["problems"] else 0


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