From f1bd099ec8909a9c9475a58ce98754b888540fb3 Mon Sep 17 00:00:00 2001 From: parththakkar106 Date: Mon, 3 Aug 2026 20:35:54 +0530 Subject: [PATCH] Cut database egress 189x by deferring the prompt snapshot The free-tier 5 GB/month network transfer allowance ran out, which blocks connections outright. The database is only ~55 MB, so 5 GB meant the whole thing was being pulled roughly 90 times over. Cause: actions is 39 MB of that 55 MB -- 541 rows at ~74 KB each, almost entirely context_snapshot, which stores the whole assembled prompt for a turn. Every adventure load and every turn fetched all of it in order to read two small things out of it: the world-change chips under an AI message (Action.world_changes) and the emit block re-attached when replaying history to the model (_history_text). The Insights viewer is the only consumer that wants the whole snapshot, and it asks for one action at a time. Lifts that slice into its own small actions.world_delta column (migration 36) and marks context_snapshot, state_before and world_state_before deferred, so they load only when something touches the attribute -- Insights, undo and retry, all single-action paths. The backfill runs server-side, dialect-specific (json_extract on SQLite, #> on Postgres), because pulling 39 MB of snapshots into Python to rewrite a slice of each would defeat the purpose. Measured at production shape (541 actions, 72 KB snapshots), one adventure load goes from 38.46 MB to 0.20 MB. The traffic that consumed 5 GB would now be about 27 MB. Deliberately not included: limiting the history query to recent actions, and removing the redundant db.refresh(adventure) calls. Both were sized against the old numbers; against a 0.20 MB load they would take ~27 MB a month down to ~10 MB, which is not worth the complexity. tests/test_egress.py hooks before_cursor_execute and asserts the emitted SQL never names the deferred columns during a bulk load, so this cannot regress silently. 123 tests pass. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01UeQVy5bEjLhfgWNc27Efet --- backend/app/context/builder.py | 18 +-- backend/app/migrations.py | 41 +++++++ backend/app/models.py | 42 +++++-- backend/app/routers/adventures.py | 18 +++ backend/tests/test_egress.py | 176 ++++++++++++++++++++++++++++++ 5 files changed, 276 insertions(+), 19 deletions(-) create mode 100644 backend/tests/test_egress.py diff --git a/backend/app/context/builder.py b/backend/app/context/builder.py index 2d69cf3..43d4d1d 100644 --- a/backend/app/context/builder.py +++ b/backend/app/context/builder.py @@ -62,15 +62,17 @@ def _history_text(action: models.Action) -> str: delta). The block is stripped before storage/UI, so without this every past AI turn would look like one that emitted nothing — biasing the model, by imitation, to stop emitting too. Player turns and blockless turns are - returned unchanged.""" + returned unchanged. + + Reads `world_delta`, not `context_snapshot`: this runs for every action in + the replayed history, and the snapshot is deferred precisely so a turn + never drags the prompt archive out of the database.""" text = action.text - snap = action.context_snapshot if isinstance(action.context_snapshot, dict) else None - if snap: - ws = snap.get("world_state") - if isinstance(ws, dict): - block = worldstate.render_delta_block(ws.get("delta") or {}) - if block: - text = f"{text}\n{block}" + wd = action.world_delta if isinstance(action.world_delta, dict) else None + if wd: + block = worldstate.render_delta_block(wd.get("delta") or {}) + if block: + text = f"{text}\n{block}" return text diff --git a/backend/app/migrations.py b/backend/app/migrations.py index bd01445..c05dd6e 100644 --- a/backend/app/migrations.py +++ b/backend/app/migrations.py @@ -110,10 +110,49 @@ MIGRATIONS: list[tuple[int, str]] = [ # predates this column. (34, "ALTER TABLE actions ADD COLUMN variants JSON"), (35, "ALTER TABLE actions ADD COLUMN variant_index INTEGER NOT NULL DEFAULT 0"), + # Egress: context_snapshot holds the whole assembled prompt (~74 KB/row) and + # was being loaded in bulk for two tiny things — the world-change chips and + # the emit block replayed into history. Lift just that slice into its own + # column so the snapshot can be deferred. Backfilled by _backfill_world_delta. + (36, "ALTER TABLE actions ADD COLUMN world_delta JSON"), ] LATEST_VERSION = max((v for v, _ in MIGRATIONS), default=1) +# Migrations that need a data pass after their DDL, keyed by version. +WORLD_DELTA_VERSION = 36 + + +def _backfill_world_delta(conn) -> None: + """Populate actions.world_delta from the existing context_snapshot. + + Runs entirely server-side: the snapshots are the reason this change exists, + so pulling ~40 MB of them into Python to rewrite a slice would defeat the + point. Dialect-specific because SQLite and Postgres spell JSON access + differently, and both have to work (SQLite locally and in tests). + """ + if conn.dialect.name == "sqlite": + sql = """ + UPDATE actions SET world_delta = json_object( + 'delta', json_extract(context_snapshot, '$.world_state.delta'), + 'applied', json_extract(context_snapshot, '$.world_state.report.applied') + ) + WHERE world_delta IS NULL + AND context_snapshot IS NOT NULL + AND json_extract(context_snapshot, '$.world_state') IS NOT NULL + """ + else: + sql = """ + UPDATE actions SET world_delta = jsonb_build_object( + 'delta', context_snapshot::jsonb #> '{world_state,delta}', + 'applied', context_snapshot::jsonb #> '{world_state,report,applied}' + ) + WHERE world_delta IS NULL + AND context_snapshot IS NOT NULL + AND jsonb_exists(context_snapshot::jsonb, 'world_state') + """ + conn.execute(text(sql)) + def _get_version(conn) -> int: if conn.dialect.name == "sqlite": @@ -153,6 +192,8 @@ def bootstrap(engine: Engine) -> None: for version, sql in MIGRATIONS: if version > current: conn.execute(text(sql)) + if version == WORLD_DELTA_VERSION: + _backfill_world_delta(conn) current = version _set_version(conn, current) _encrypt_plaintext_api_keys(conn) diff --git a/backend/app/models.py b/backend/app/models.py index 11823be..ddb57ac 100644 --- a/backend/app/models.py +++ b/backend/app/models.py @@ -205,13 +205,30 @@ class Action(Base): text: Mapped[str] = mapped_column(Text, default="") # Reasoning-model "thinking" that preceded the text (AI actions only). reasoning: Mapped[str | None] = mapped_column(Text, nullable=True) - context_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True) + # The full assembled prompt for this turn, for the Insights viewer. By far + # the biggest column in the database (~74 KB/row in production), and needed + # by exactly one endpoint, one action at a time — so it is deferred: never + # loaded unless something actually touches the attribute. Bulk readers must + # NOT touch it; that is what `world_delta` below exists for. + context_snapshot: Mapped[dict | None] = mapped_column( + JSON, nullable=True, deferred=True + ) + # The small slice of the snapshot that IS needed in bulk: this turn's RPG + # state changes, for the inline chips under an AI message (world_changes) + # and for re-attaching the emit block when replaying history to the model. + # Mirrors the active variant, same as text/reasoning/context_snapshot. + world_delta: Mapped[dict | None] = mapped_column(JSON, nullable=True) # Copy of Adventure.script_state as it was immediately BEFORE this action's # script hooks ran, so undo/retry can roll the shared scoreboard back. - # NULL for actions created before this column existed. - state_before: Mapped[dict | None] = mapped_column(JSON, nullable=True) + # NULL for actions created before this column existed. Deferred: only ever + # read for the one action being undone or retried. + state_before: Mapped[dict | None] = mapped_column( + JSON, nullable=True, deferred=True + ) # Phase 12: same idea for the RPG world_state, so undo/retry rolls it back too. - world_state_before: Mapped[dict | None] = mapped_column(JSON, nullable=True) + world_state_before: Mapped[dict | None] = mapped_column( + JSON, nullable=True, deferred=True + ) # Retry history (AI actions): every attempt made for this turn, oldest # first, INCLUDING the active one. NULL/empty means never retried — the row # is its own only version. `variant_index` says which entry `text`, @@ -225,14 +242,17 @@ class Action(Base): @property def world_changes(self) -> list[dict]: - """Compact per-turn RPG state changes (Phase 12), derived from the - stored snapshot, for the inline summary under an AI message. Labels are - path-based (no schema needed): `npc.gwen.trust` -> "gwen trust".""" - cs = self.context_snapshot if isinstance(self.context_snapshot, dict) else None - ws = cs.get("world_state") if cs else None - if not isinstance(ws, dict): + """Compact per-turn RPG state changes (Phase 12), for the inline summary + under an AI message. Labels are path-based (no schema needed): + `npc.gwen.trust` -> "gwen trust". + + Reads `world_delta`, never `context_snapshot` — this runs for every + action in a list response, and touching the deferred snapshot here + would drag the whole prompt archive out of the database.""" + wd = self.world_delta if isinstance(self.world_delta, dict) else None + if wd is None: return [] - applied = (ws.get("report") or {}).get("applied") or [] + applied = wd.get("applied") or [] out: list[dict] = [] for entry in applied: parts = str(entry.get("path", "")).split(".") diff --git a/backend/app/routers/adventures.py b/backend/app/routers/adventures.py index 88df0f0..83ea652 100644 --- a/backend/app/routers/adventures.py +++ b/backend/app/routers/adventures.py @@ -333,6 +333,22 @@ def snapshot_world_state(adventure: models.Adventure) -> dict: VARIANT_SNAPSHOT_KEYS = ("world_state", "script", "raw_output") +def world_delta_of(snapshot: dict | None) -> dict | None: + """The bulk-read slice of a context snapshot, for Action.world_delta. + + context_snapshot is deferred (it holds the whole assembled prompt), so the + two things that ARE needed for every action — the world-change chips and + the emit block replayed into history — get their own small column. Keep + this in step with the snapshot wherever one is written.""" + ws = (snapshot or {}).get("world_state") + if not isinstance(ws, dict): + return None + return { + "delta": ws.get("delta") or {}, + "applied": (ws.get("report") or {}).get("applied") or [], + } + + def variant_of(action: models.Action, adventure: models.Adventure) -> dict: """Freeze an action's *current* content as a variant entry. @@ -365,6 +381,7 @@ def apply_variant(action: models.Action, adventure: models.Adventure, index: int else: snapshot.pop(key, None) action.context_snapshot = snapshot + action.world_delta = world_delta_of(snapshot) action.variant_index = index if isinstance(entry.get("script_state"), dict): adventure.script_state = copy.deepcopy(entry["script_state"]) @@ -637,6 +654,7 @@ async def _generate_turn( text=text, reasoning=reasoning, context_snapshot=snapshot, + world_delta=world_delta_of(snapshot), state_before=state_before, world_state_before=world_state_before, ) diff --git a/backend/tests/test_egress.py b/backend/tests/test_egress.py new file mode 100644 index 0000000..8b169bb --- /dev/null +++ b/backend/tests/test_egress.py @@ -0,0 +1,176 @@ +"""Guards on how much the database is asked for. + +context_snapshot holds the entire assembled prompt for a turn (~74 KB/row in +production, 94% of the database). It used to be pulled for every action on +every adventure load and every turn, to read two tiny things out of it. These +tests fail if that regresses. + + python -m pytest tests/test_egress.py -v +""" +import os +import tempfile + +_tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False) +_tmp.close() +os.environ["AIDND_DB_PATH"] = _tmp.name +os.environ.pop("AIDND_DATABASE_URL", None) +os.environ.pop("DATABASE_URL", None) + +import pytest +from fastapi import Depends +from fastapi.testclient import TestClient +from sqlalchemy import event, text + +from app import auth, limits, migrations, models +from app.database import Base, SessionLocal, engine, get_db +from app.main import app + +# A stand-in for the real thing: the assembled prompt, which is what makes the +# column enormous, plus the small world_state slice the UI actually needs. +BIG_SNAPSHOT = { + "system": "x" * 20_000, + "story": "y" * 40_000, + "world_state": { + "delta": {"player.hp": -15}, + "report": {"applied": [{"path": "player.hp", "old": 100, "new": 85}]}, + "state": {"player": {"hp": 85}}, + }, +} + + +@pytest.fixture() +def sql_log(): + """Every statement the ORM sends, for asserting on what was selected.""" + statements: list[str] = [] + + def record(conn, cursor, statement, parameters, context, executemany): + statements.append(statement) + + event.listen(engine, "before_cursor_execute", record) + try: + yield statements + finally: + event.remove(engine, "before_cursor_execute", record) + + +@pytest.fixture() +def client(monkeypatch): + Base.metadata.create_all(bind=engine) + setup = SessionLocal() + user = models.User(is_guest=False, email="egress@example.com") + setup.add(user) + setup.flush() + setup.add(models.Settings(user_id=user.id, api_key="enc:dummy", model="test-model")) + adventure = models.Adventure(user_id=user.id, title="Cave", script_state={}) + setup.add(adventure) + setup.flush() + for i in range(12): + setup.add(models.Action( + adventure_id=adventure.id, index=i, + type="ai" if i % 2 else "do", text=f"Action {i}.", + context_snapshot=BIG_SNAPSHOT, + world_delta={"delta": {"player.hp": -15}, + "applied": [{"path": "player.hp", "old": 100, "new": 85}]}, + )) + setup.commit() + adv_id, user_id = adventure.id, user.id + setup.close() + + monkeypatch.setattr(limits, "rate_limit", lambda *a, **k: None) + monkeypatch.setattr(limits, "check_row_cap", lambda *a, **k: None) + + def _current_user(db=Depends(get_db)): + return db.get(models.User, user_id) + + app.dependency_overrides[auth.get_current_user] = _current_user + c = TestClient(app) + c.adv_id = adv_id + try: + yield c + finally: + app.dependency_overrides.clear() + Base.metadata.drop_all(bind=engine) + + +def action_selects(statements: list[str]) -> list[str]: + return [s for s in statements if "FROM actions" in s and s.lstrip().upper().startswith("SELECT")] + + +# ------------------------------------------------------- the deferred columns + +def test_loading_an_adventure_does_not_fetch_context_snapshot(client, sql_log): + r = client.get(f"/api/adventures/{client.adv_id}") + assert r.status_code == 200, r.text + assert len(r.json()["actions"]) == 12 + + selects = action_selects(sql_log) + assert selects, "expected at least one SELECT against actions" + offenders = [s for s in selects if "context_snapshot" in s] + assert offenders == [], f"context_snapshot was fetched in bulk:\n{offenders[0][:400]}" + + +def test_state_before_and_world_state_before_are_not_fetched_in_bulk(client, sql_log): + """Both are rollback snapshots, only ever needed for the single action + being undone or retried.""" + client.get(f"/api/adventures/{client.adv_id}") + selects = action_selects(sql_log) + for column in ("state_before", "world_state_before"): + offenders = [s for s in selects if column in s] + assert offenders == [], f"{column} was fetched in bulk" + + +def test_world_changes_still_works_without_the_snapshot(client): + """The chips under an AI message must survive the snapshot being deferred.""" + r = client.get(f"/api/adventures/{client.adv_id}") + ai = [a for a in r.json()["actions"] if a["type"] == "ai"] + assert ai, "fixture should have AI actions" + assert ai[0]["world_changes"] == [ + {"kind": "stat", "label": "hp", "delta": -15, "value": 85} + ] + + +def test_snapshot_is_still_reachable_on_demand(client): + """Deferred means lazy, not gone — Insights still gets the full thing.""" + r = client.get(f"/api/adventures/{client.adv_id}") + action_id = r.json()["actions"][0]["id"] + r = client.get(f"/api/adventures/{client.adv_id}/actions/{action_id}/context") + assert r.status_code == 200, r.text + assert r.json()["system"] == "x" * 20_000 + + +# ------------------------------------------------------------------ backfill + +def test_backfill_populates_world_delta_from_existing_snapshots(client): + """Migration 36 lifts the slice out server-side, without reading the + snapshots into Python.""" + db = SessionLocal() + try: + db.execute(text("UPDATE actions SET world_delta = NULL")) + db.commit() + assert db.query(models.Action).filter(models.Action.world_delta.isnot(None)).count() == 0 + + with engine.begin() as conn: + migrations._backfill_world_delta(conn) + + db.expire_all() + actions = db.query(models.Action).all() + assert all(a.world_delta is not None for a in actions) + assert actions[0].world_delta["delta"] == {"player.hp": -15} + assert actions[0].world_delta["applied"] == [ + {"path": "player.hp", "old": 100, "new": 85} + ] + finally: + db.close() + + +def test_backfill_leaves_actions_without_world_state_alone(client): + db = SessionLocal() + try: + db.execute(text("UPDATE actions SET world_delta = NULL, context_snapshot = '{\"story\": \"s\"}'")) + db.commit() + with engine.begin() as conn: + migrations._backfill_world_delta(conn) + db.expire_all() + assert all(a.world_delta is None for a in db.query(models.Action).all()) + finally: + db.close()