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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UeQVy5bEjLhfgWNc27Efet
This commit is contained in:
parththakkar106
2026-08-03 20:35:54 +05:30
co-authored by Claude Opus 5
parent 7c538c8235
commit f1bd099ec8
5 changed files with 276 additions and 19 deletions
+8 -6
View File
@@ -62,13 +62,15 @@ def _history_text(action: models.Action) -> str:
delta). The block is stripped before storage/UI, so without this every past 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 AI turn would look like one that emitted nothing — biasing the model, by
imitation, to stop emitting too. Player turns and blockless turns are 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 text = action.text
snap = action.context_snapshot if isinstance(action.context_snapshot, dict) else None wd = action.world_delta if isinstance(action.world_delta, dict) else None
if snap: if wd:
ws = snap.get("world_state") block = worldstate.render_delta_block(wd.get("delta") or {})
if isinstance(ws, dict):
block = worldstate.render_delta_block(ws.get("delta") or {})
if block: if block:
text = f"{text}\n{block}" text = f"{text}\n{block}"
return text return text
+41
View File
@@ -110,10 +110,49 @@ MIGRATIONS: list[tuple[int, str]] = [
# predates this column. # predates this column.
(34, "ALTER TABLE actions ADD COLUMN variants JSON"), (34, "ALTER TABLE actions ADD COLUMN variants JSON"),
(35, "ALTER TABLE actions ADD COLUMN variant_index INTEGER NOT NULL DEFAULT 0"), (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) 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: def _get_version(conn) -> int:
if conn.dialect.name == "sqlite": if conn.dialect.name == "sqlite":
@@ -153,6 +192,8 @@ def bootstrap(engine: Engine) -> None:
for version, sql in MIGRATIONS: for version, sql in MIGRATIONS:
if version > current: if version > current:
conn.execute(text(sql)) conn.execute(text(sql))
if version == WORLD_DELTA_VERSION:
_backfill_world_delta(conn)
current = version current = version
_set_version(conn, current) _set_version(conn, current)
_encrypt_plaintext_api_keys(conn) _encrypt_plaintext_api_keys(conn)
+31 -11
View File
@@ -205,13 +205,30 @@ class Action(Base):
text: Mapped[str] = mapped_column(Text, default="") text: Mapped[str] = mapped_column(Text, default="")
# Reasoning-model "thinking" that preceded the text (AI actions only). # Reasoning-model "thinking" that preceded the text (AI actions only).
reasoning: Mapped[str | None] = mapped_column(Text, nullable=True) 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 # 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. # script hooks ran, so undo/retry can roll the shared scoreboard back.
# NULL for actions created before this column existed. # NULL for actions created before this column existed. Deferred: only ever
state_before: Mapped[dict | None] = mapped_column(JSON, nullable=True) # 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. # 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 # Retry history (AI actions): every attempt made for this turn, oldest
# first, INCLUDING the active one. NULL/empty means never retried — the row # first, INCLUDING the active one. NULL/empty means never retried — the row
# is its own only version. `variant_index` says which entry `text`, # is its own only version. `variant_index` says which entry `text`,
@@ -225,14 +242,17 @@ class Action(Base):
@property @property
def world_changes(self) -> list[dict]: def world_changes(self) -> list[dict]:
"""Compact per-turn RPG state changes (Phase 12), derived from the """Compact per-turn RPG state changes (Phase 12), for the inline summary
stored snapshot, for the inline summary under an AI message. Labels are under an AI message. Labels are path-based (no schema needed):
path-based (no schema needed): `npc.gwen.trust` -> "gwen trust".""" `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 Reads `world_delta`, never `context_snapshot` — this runs for every
if not isinstance(ws, dict): 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 [] return []
applied = (ws.get("report") or {}).get("applied") or [] applied = wd.get("applied") or []
out: list[dict] = [] out: list[dict] = []
for entry in applied: for entry in applied:
parts = str(entry.get("path", "")).split(".") parts = str(entry.get("path", "")).split(".")
+18
View File
@@ -333,6 +333,22 @@ def snapshot_world_state(adventure: models.Adventure) -> dict:
VARIANT_SNAPSHOT_KEYS = ("world_state", "script", "raw_output") 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: def variant_of(action: models.Action, adventure: models.Adventure) -> dict:
"""Freeze an action's *current* content as a variant entry. """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: else:
snapshot.pop(key, None) snapshot.pop(key, None)
action.context_snapshot = snapshot action.context_snapshot = snapshot
action.world_delta = world_delta_of(snapshot)
action.variant_index = index action.variant_index = index
if isinstance(entry.get("script_state"), dict): if isinstance(entry.get("script_state"), dict):
adventure.script_state = copy.deepcopy(entry["script_state"]) adventure.script_state = copy.deepcopy(entry["script_state"])
@@ -637,6 +654,7 @@ async def _generate_turn(
text=text, text=text,
reasoning=reasoning, reasoning=reasoning,
context_snapshot=snapshot, context_snapshot=snapshot,
world_delta=world_delta_of(snapshot),
state_before=state_before, state_before=state_before,
world_state_before=world_state_before, world_state_before=world_state_before,
) )
+176
View File
@@ -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()