Rewind a session and migrate its database

"""Replay the rewind demo offline, and run the session migration CLI on a fresh legacy database.

Run: python -m unittest -v test_pins.py
The model is replaced by its recorded responses; the Runner, DatabaseSessionService, rewind and
`adk migrate session` are real. Nothing here calls a model or the network.
"""
import asyncio
import json
import logging
import subprocess
import sys
import tempfile
import unittest
import warnings
from pathlib import Path
from unittest import mock

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.genai import types  # noqa: E402

import check_migrated  # noqa: E402
import legacy  # noqa: E402
import pins  # noqa: E402

HERE = Path(__file__).resolve().parent
RECORD = json.loads((HERE / "recorded-runs.json").read_text(encoding="utf-8"))
ADK = str(Path(sys.executable).with_name("adk"))


class ReplayLlm(BaseLlm):
    turns: list

    async def generate_content_async(self, llm_request, stream=False):
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)["content"]))


class Rewind(unittest.TestCase):
    def setUp(self):
        warnings.simplefilter("ignore")
        self.tmp = tempfile.TemporaryDirectory()
        model = ReplayLlm(model="replay", turns=list(RECORD["rewind"]["model_turns"]))
        self.agent = pins.root_agent.clone(update={"model": model})

    def tearDown(self):
        self.tmp.cleanup()

    def run_demo(self, rewind_to_unknown: bool = False):
        with mock.patch.object(pins, "root_agent", self.agent):
            runner = pins.build_runner(f"sqlite+aiosqlite:///{Path(self.tmp.name) / 'sessions.db'}")

        async def go():
            session = await runner.session_service.create_session(app_name=pins.APP, user_id=pins.USER)
            try:
                await pins.ask(runner, session.id, RECORD["rewind"]["turns"][0])
                second = await pins.ask(runner, session.id, RECORD["rewind"]["turns"][1])
                target = "e-0000" if rewind_to_unknown else second[0].invocation_id
                await runner.rewind_async(user_id=pins.USER, session_id=session.id, rewind_before_invocation_id=target)
                return await pins.snapshot(runner, session.id)
            finally:
                await runner.close()

        return asyncio.run(go())

    def test_rewind_restores_session_state_and_the_artifact(self):
        snap = self.run_demo()
        expected = RECORD["rewind"]["snapshots"]["after rewind"]
        self.assertEqual((snap["pinned"], snap["artifact"]), (expected["pinned"], expected["artifact"]))
        self.assertEqual(snap["artifact_versions"], expected["artifact_versions"])  # restored as a new version

    def test_user_state_is_not_restored(self):
        snap = self.run_demo()
        self.assertEqual(snap["user:team"], "platform")
        self.assertEqual(RECORD["rewind"]["snapshots"]["after turn 1"]["user:team"], "payments")

    def test_rewound_events_stay_in_the_log(self):
        snap = self.run_demo()
        self.assertEqual(snap["events"], RECORD["rewind"]["snapshots"]["after turn 2"]["events"] + 1)

    def test_an_unknown_invocation_id_raises(self):
        from google.adk.errors._invocation_not_found_error import InvocationNotFoundError
        with self.assertRaisesRegex(InvocationNotFoundError, "Invocation ID not found: e-0000"):
            self.run_demo(rewind_to_unknown=True)


class Migrate(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tmp = tempfile.TemporaryDirectory()
        cls.dir = Path(cls.tmp.name)
        asyncio.run(legacy.build(str(cls.dir / "legacy.db")))

    @classmethod
    def tearDownClass(cls):
        cls.tmp.cleanup()

    def setUp(self):
        warnings.simplefilter("ignore")

    def migrate(self, source: str, dest: str, *extra: str) -> subprocess.CompletedProcess:
        return subprocess.run([ADK, "migrate", "session", "--source_db_url", f"sqlite:///{source}",
                               "--dest_db_url", f"sqlite:///{dest}", *extra], cwd=self.dir, capture_output=True, text=True)

    def deltas(self, name: str) -> list[str]:
        return [l for l in asyncio.run(check_migrated.check(str(self.dir / name))) if l.startswith("event e-2 docs")]

    def test_a_restricted_migration_drops_an_event_it_cannot_unpickle(self):
        proc = self.migrate("legacy.db", "migrated.db")
        self.assertEqual(proc.returncode, 0)
        self.assertIn("Refusing to load collections.Counter", proc.stderr)
        self.assertEqual(self.deltas("migrated.db"), ["event e-2 docs_assistant  state_delta={}"])

    def test_allow_unsafe_unpickling_keeps_the_event(self):
        proc = self.migrate("legacy.db", "migrated-full.db", "--allow-unsafe-unpickling")
        self.assertEqual(proc.returncode, 0)
        self.assertIn("page_views", self.deltas("migrated-full.db")[0])

    def test_in_place_migration_is_refused(self):
        proc = self.migrate("legacy.db", "legacy.db")
        self.assertEqual(proc.returncode, 1)
        self.assertIn("In-place migration is not supported", proc.stdout + proc.stderr)

    def test_a_misspelled_source_path_reports_no_migration_needed(self):
        proc = self.migrate("legacy-db.db", "out.db")
        self.assertEqual(proc.returncode, 0)
        self.assertIn("already at latest version 1. No migration needed.", proc.stderr)
        self.assertTrue((self.dir / "legacy-db.db").exists())
        self.assertFalse((self.dir / "out.db").exists())


if __name__ == "__main__":
    unittest.main()

"""Record the rewind demo on Gemini and the session-database migration commands.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes run-output.txt and recorded-runs.json. Databases are created in a temporary folder.
"""
import asyncio
import json
import logging
import os
import re
import subprocess
import sys
import tempfile
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

TMP = Path(tempfile.mkdtemp())
os.environ["DOCS_DB"] = f"sqlite+aiosqlite:///{TMP / 'sessions.db'}"

import check_migrated  # noqa: E402
import legacy  # noqa: E402
import pins  # noqa: E402

HERE = Path(__file__).resolve().parent
ADK = str(Path(sys.executable).with_name("adk"))
TURNS = ["I'm on the payments team. Pin the rollback page.",
         "Pin the tokens page too. And I moved to the platform team.",
         "Which pages are pinned, and which team am I on?"]


async def rewind_demo() -> dict:
    runner = pins.build_runner()
    session = await runner.session_service.create_session(app_name=pins.APP, user_id=pins.USER)
    rec = {"snapshots": {}, "turns": TURNS}
    try:
        await pins.ask(runner, session.id, TURNS[0])
        rec["snapshots"]["after turn 1"] = await pins.snapshot(runner, session.id)
        second = await pins.ask(runner, session.id, TURNS[1])
        rec["snapshots"]["after turn 2"] = await pins.snapshot(runner, session.id)
        rec["rewound_invocation"] = second[0].invocation_id
        await runner.rewind_async(user_id=pins.USER, session_id=session.id, rewind_before_invocation_id=second[0].invocation_id)
        rec["snapshots"]["after rewind"] = await pins.snapshot(runner, session.id)
        third = await pins.ask(runner, session.id, TURNS[2])
        rec["answer"] = third[-1].content.parts[-1].text
        stored = await runner.session_service.get_session(app_name=pins.APP, user_id=pins.USER, session_id=session.id)
        rec["rewind_event"] = next(e.actions.model_dump(mode="json", exclude_none=True, exclude_defaults=True)
                                   for e in stored.events if e.actions.rewind_before_invocation_id)
        rec["model_turns"] = [{"content": {"role": "model", "parts": [
            {k: v for k, v in p.model_dump(mode="json", exclude_none=True).items() if k != "thought_signature"}
            for p in e.content.parts]}} for e in stored.events if e.content and e.content.role == "model"]
        try:
            await runner.rewind_async(user_id=pins.USER, session_id=session.id, rewind_before_invocation_id="e-0000")
        except Exception as exc:  # recorded failure
            rec["unknown_invocation"] = f"{type(exc).__name__}: {exc}"
    finally:
        await runner.close()
    return rec


def sh(cmd: list[str]) -> dict:
    proc = subprocess.run(cmd, cwd=TMP, capture_output=True, text=True, timeout=120)
    lines = [re.sub(r"^\d{4}-\d\d-\d\d [\d:,]+ - ", "", l) for l in (proc.stderr + proc.stdout).splitlines()
             if l.strip() and not l.startswith(" ")]
    lines = [l[:260] + (" …" if len(l) > 260 else "") for l in lines]
    shown = ["adk", *cmd[1:]] if cmd[0] == ADK else ["python", *[Path(a).name if a.endswith(".py") else a for a in cmd[1:]]]
    return {"command": " ".join(shown), "lines": lines, "exit": proc.returncode}


def migrations() -> dict:
    py = sys.executable
    runs = {
        "build": sh([py, str(HERE / "legacy.py"), "legacy.db"]),
        "migrate": sh([ADK, "migrate", "session", "--source_db_url", "sqlite:///legacy.db", "--dest_db_url", "sqlite:///migrated.db"]),
        "check": sh([py, str(HERE / "check_migrated.py"), "migrated.db"]),
        "migrate_unsafe": sh([ADK, "migrate", "session", "--source_db_url", "sqlite:///legacy.db",
                              "--dest_db_url", "sqlite:///migrated-full.db", "--allow-unsafe-unpickling"]),
        "check_unsafe": sh([py, str(HERE / "check_migrated.py"), "migrated-full.db"]),
        "in_place": sh([ADK, "migrate", "session", "--source_db_url", "sqlite:///legacy.db", "--dest_db_url", "sqlite:///legacy.db"]),
        "typo": sh([ADK, "migrate", "session", "--source_db_url", "sqlite:///legacy-db.db", "--dest_db_url", "sqlite:///out.db"]),
    }
    runs["typo"]["created"] = sorted(p.name for p in TMP.glob("*.db"))
    return runs


def main() -> None:
    rec = {"model": pins.MODEL, "rewind": asyncio.run(rewind_demo()), "migrate": migrations()}
    r = rec["rewind"]
    out = ["$ python pins.py", *[f"{k}: {v}" for k, v in r["snapshots"].items()], f"turn 3: {r['answer']}", "",
           "# runner.rewind_async(..., rewind_before_invocation_id='e-0000')", r["unknown_invocation"], ""]
    for run in rec["migrate"].values():
        out += [f"$ {run['command']}", *run["lines"], f"exit {run['exit']}", ""]
    out.append(f"files now: {rec['migrate']['typo']['created']}")
    text = "\n".join(out)
    (HERE / "run-output.txt").write_text(text + "\n", encoding="utf-8")
    (HERE / "recorded-runs.json").write_text(json.dumps(rec, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()