Read a window of the story per turn instead of all of it
Two reads still grew without bound after the snapshot fix. `Action.variants` holds every discarded retry attempt, but a list response only needs how many there are — so each retry permanently added ~5 KB to every later load of that adventure. Defer the column and keep the count beside it (migration 37, backfilled server-side), with set_variants() as the one write path that keeps the two in step. `story_actions()` walked adventure.actions, then every caller threw almost all of it away: the builder concatenates the story and immediately cuts it back to the token budget, the NPC check looks at the last 6, retrieval at the last 4, the cursor clamp only wants a count. A turn on a 200-action adventure read 839 KB to use ~70 KB, and grew with every turn played. app/context/ history.py serves those shapes from SQL; window_covering() measures the actions it fetched and projects how many more it needs, fetching only the part it does not already hold. Memorybank cursors move to position_of_index() and settled_count()/settled_slice() — same arithmetic, no full list. The scripting pipeline still receives the whole history per AI Dungeon's API, and every helper reuses adventure.actions when it is already loaded, so a scripted adventure pays what it always did and never twice. Measured at production shape: retry tax 5.1 KB -> 0; turn 200 839 KB -> 129 KB and flat from ~turn 50; a 200-turn playthrough 84.5 MB -> 23.0 MB; a delete 115 KB -> 5 KB. Verified the window builds a byte-identical prompt to the full story across budgets from 1K to 100K tokens, with and without the retry exclusion - this is a cost change and nothing else. Cursor helpers checked against the old list arithmetic, including after deleting a middle action. Counts are real SELECT count(...): Query.count() wraps the entity select in a subquery, so the SQL named every deferred column and the egress guard could not tell it apart from a bulk fetch. 139 tests pass; the four new guards verified by sabotage. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UeQVy5bEjLhfgWNc27Efet
This commit is contained in:
parththakkar106
co-authored by
Claude Opus 5
Claude Opus 5
parent
f1bd099ec8
commit
47e33fa311
@@ -22,6 +22,7 @@ from fastapi.testclient import TestClient
|
||||
from sqlalchemy import event, text
|
||||
|
||||
from app import auth, limits, migrations, models
|
||||
from app.context import history
|
||||
from app.database import Base, SessionLocal, engine, get_db
|
||||
from app.main import app
|
||||
|
||||
@@ -37,6 +38,14 @@ BIG_SNAPSHOT = {
|
||||
},
|
||||
}
|
||||
|
||||
# Retry history: each discarded attempt keeps its full narration, so an action
|
||||
# retried a few times carries several KB that a list response only ever counts.
|
||||
BIG_VARIANTS = [
|
||||
{"text": "z" * 4_000, "reasoning": None, "script_state": {},
|
||||
"created_at": "2026-01-01T00:00:00"}
|
||||
for _ in range(3)
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sql_log():
|
||||
@@ -71,6 +80,10 @@ def client(monkeypatch):
|
||||
context_snapshot=BIG_SNAPSHOT,
|
||||
world_delta={"delta": {"player.hp": -15},
|
||||
"applied": [{"path": "player.hp", "old": 100, "new": 85}]},
|
||||
# Every AI action has been retried twice, so `variants` is carrying
|
||||
# weight the list response must not pay for.
|
||||
variants=BIG_VARIANTS if i % 2 else None,
|
||||
variant_count=len(BIG_VARIANTS) if i % 2 else 0,
|
||||
))
|
||||
setup.commit()
|
||||
adv_id, user_id = adventure.id, user.id
|
||||
@@ -129,6 +142,48 @@ def test_world_changes_still_works_without_the_snapshot(client):
|
||||
]
|
||||
|
||||
|
||||
def test_loading_an_adventure_does_not_fetch_variants(client, sql_log):
|
||||
"""Same failure as context_snapshot, one size down: the payload carries
|
||||
only `variant_count`, but loading the column to compute it made every retry
|
||||
a permanent tax on every later load of that adventure."""
|
||||
r = client.get(f"/api/adventures/{client.adv_id}")
|
||||
assert r.status_code == 200, r.text
|
||||
|
||||
selects = action_selects(sql_log)
|
||||
assert selects, "expected at least one SELECT against actions"
|
||||
offenders = [s for s in selects if "variants" in s]
|
||||
assert offenders == [], f"variants was fetched in bulk:\n{offenders[0][:400]}"
|
||||
|
||||
|
||||
def test_variant_count_survives_variants_being_deferred(client):
|
||||
"""The pager reads this number; it has to be right without the column."""
|
||||
r = client.get(f"/api/adventures/{client.adv_id}")
|
||||
by_type = {}
|
||||
for action in r.json()["actions"]:
|
||||
by_type.setdefault(action["type"], []).append(action)
|
||||
assert all(a["variant_count"] == len(BIG_VARIANTS) for a in by_type["ai"])
|
||||
assert all(a["variant_count"] == 0 for a in by_type["do"])
|
||||
|
||||
|
||||
def test_counting_actions_does_not_name_the_deferred_columns(client, sql_log):
|
||||
"""A count that wraps the entity select in a subquery names every column in
|
||||
the emitted SQL — no bytes come back, but the database still reads them and
|
||||
the guard above cannot tell it apart from a real bulk fetch."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
adventure = db.get(models.Adventure, client.adv_id)
|
||||
sql_log.clear()
|
||||
assert history.count(adventure) == 12
|
||||
counts = [s for s in sql_log if "count" in s.lower()]
|
||||
assert counts, "expected a COUNT to be emitted"
|
||||
for column in ("context_snapshot", "state_before", "world_state_before", "variants"):
|
||||
assert not any(column in s for s in counts), (
|
||||
f"{column} is named by the count query:\n{counts[0][:400]}"
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
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}")
|
||||
@@ -163,6 +218,27 @@ def test_backfill_populates_world_delta_from_existing_snapshots(client):
|
||||
db.close()
|
||||
|
||||
|
||||
def test_backfill_populates_variant_count_from_existing_variants(client):
|
||||
"""Migration 37 counts the lists server-side — reading them into Python to
|
||||
count them would mean pulling the column across the wire once to stop
|
||||
pulling it across forever."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
db.execute(text("UPDATE actions SET variant_count = 0"))
|
||||
db.commit()
|
||||
|
||||
with engine.begin() as conn:
|
||||
migrations._backfill_variant_count(conn)
|
||||
|
||||
db.expire_all()
|
||||
actions = db.query(models.Action).order_by(models.Action.index).all()
|
||||
for action in actions:
|
||||
expected = len(BIG_VARIANTS) if action.type == "ai" else 0
|
||||
assert action.variant_count == expected, f"action {action.index}"
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def test_backfill_leaves_actions_without_world_state_alone(client):
|
||||
db = SessionLocal()
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,249 @@
|
||||
"""The context builder reads a window of the story, not all of it.
|
||||
|
||||
Walking `adventure.actions` every turn made a turn cost O(story length), so a
|
||||
long adventure read hundreds of KB to use the tail of it — and the cost grew
|
||||
with every turn played. `app.context.history` serves tails, slices and counts
|
||||
from SQL instead.
|
||||
|
||||
Two things have to hold, and both are easy to break by accident:
|
||||
|
||||
* the window must produce **exactly** the prompt the full story produced, or
|
||||
this is a behaviour change wearing an optimization's clothes;
|
||||
* the helpers must agree with the old list arithmetic, because memorybank's
|
||||
cursors are *positions* in that list and a cursor off by one silently
|
||||
summarizes the wrong actions.
|
||||
|
||||
python -m pytest tests/test_history_window.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 sqlalchemy import event
|
||||
|
||||
from app import memorybank, models
|
||||
from app.context import builder, history
|
||||
from app.database import Base, SessionLocal, engine
|
||||
|
||||
# Long enough that a window is much smaller than the whole story.
|
||||
ACTION_COUNT = 200
|
||||
NARRATION = (
|
||||
"The scrub gives way to a shallow bowl of land where woodsmoke hangs in "
|
||||
"flat grey layers, and somewhere behind the largest tent a woman is "
|
||||
"arguing, low and fast. "
|
||||
) * 3
|
||||
|
||||
SCHEMA = {
|
||||
"player": {"hp": {"min": 0, "max": 100, "initial": 100, "desc": "Health"}},
|
||||
"npcs": {
|
||||
"gwen": {"name": "Gwen", "keys": ["gwen"], "desc": "A scout.",
|
||||
"stats": {"trust": {"min": 0, "max": 100, "initial": 30}}},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def story():
|
||||
"""An adventure with ACTION_COUNT actions, plus its settings."""
|
||||
Base.metadata.create_all(bind=engine)
|
||||
db = SessionLocal()
|
||||
user = models.User(is_guest=False, email="window@example.com")
|
||||
db.add(user)
|
||||
db.flush()
|
||||
settings = models.Settings(user_id=user.id, api_key="enc:dummy", model="m")
|
||||
db.add(settings)
|
||||
scenario = models.Scenario(user_id=user.id, title="S", stat_schema=SCHEMA,
|
||||
prompt="A long road." * 50)
|
||||
db.add(scenario)
|
||||
db.flush()
|
||||
adventure = models.Adventure(
|
||||
user_id=user.id, title="Long", scenario_id=scenario.id, script_state={},
|
||||
memory="The hero is hunting bandits. " * 20,
|
||||
world_state={"player": {"hp": 100}, "npc": {"gwen": {"trust": 30}},
|
||||
"milestones": {}, "flags": {}, "_meta": {"last_changed": {}}},
|
||||
)
|
||||
db.add(adventure)
|
||||
db.flush()
|
||||
db.add(models.StoryCard(adventure_id=adventure.id, name="Gwen", keys="gwen",
|
||||
entry="A scout with sharp eyes.", type="lore"))
|
||||
for i in range(ACTION_COUNT):
|
||||
db.add(models.Action(
|
||||
adventure_id=adventure.id, index=i,
|
||||
type="ai" if i % 2 else "do",
|
||||
text=f"[{i}] {NARRATION}",
|
||||
world_delta={"delta": {"player.hp": -1},
|
||||
"applied": [{"path": "player.hp", "old": 100, "new": 99}]},
|
||||
))
|
||||
db.commit()
|
||||
db.expire_all()
|
||||
adventure = db.get(models.Adventure, adventure.id)
|
||||
settings = db.get(models.Settings, settings.id)
|
||||
try:
|
||||
yield db, adventure, settings
|
||||
finally:
|
||||
db.close()
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
def full_window(adventure, budget_tokens, token_counter, exclude_action_id=None):
|
||||
"""Stand-in for window_covering that hands back the entire story, i.e. the
|
||||
behaviour this module replaced."""
|
||||
return history.story_actions(adventure, exclude_action_id)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def actions_loaded():
|
||||
"""Counts Action rows the ORM materializes, i.e. how much of the story was
|
||||
actually fetched. rowcount is meaningless for SELECT on SQLite, so count
|
||||
the objects the mapper builds instead."""
|
||||
loaded = {"n": 0}
|
||||
|
||||
def on_load(target, context):
|
||||
loaded["n"] += 1
|
||||
|
||||
event.listen(models.Action, "load", on_load)
|
||||
try:
|
||||
yield loaded
|
||||
finally:
|
||||
event.remove(models.Action, "load", on_load)
|
||||
|
||||
|
||||
# ------------------------------------------------------- the prompt is equal
|
||||
|
||||
@pytest.mark.parametrize("budget", [1024, 4096, 8192, 16384, 65536])
|
||||
def test_window_builds_the_same_prompt_as_the_whole_story(story, budget, monkeypatch):
|
||||
db, adventure, settings = story
|
||||
settings.context_token_budget = budget
|
||||
|
||||
windowed = builder.build_context(adventure, settings)
|
||||
monkeypatch.setattr(builder.history, "window_covering", full_window)
|
||||
db.expire(adventure)
|
||||
everything = builder.build_context(adventure, settings)
|
||||
|
||||
assert windowed[0] == everything[0], "system prompt differs"
|
||||
assert windowed[1] == everything[1], "story prompt differs"
|
||||
assert windowed[2]["history"] == everything[2]["history"]
|
||||
assert windowed[2]["cards"] == everything[2]["cards"]
|
||||
|
||||
|
||||
def test_window_matches_on_the_retry_shape(story, monkeypatch):
|
||||
"""Retry excludes the action being regenerated; the exclusion has to reach
|
||||
the window query, not just the in-memory filter."""
|
||||
db, adventure, settings = story
|
||||
last = history.tail(adventure, 1)[0]
|
||||
|
||||
windowed = builder.build_context(adventure, settings, exclude_action_id=last.id)
|
||||
assert f"[{last.index}]" not in windowed[1]
|
||||
|
||||
monkeypatch.setattr(builder.history, "window_covering", full_window)
|
||||
db.expire(adventure)
|
||||
everything = builder.build_context(adventure, settings, exclude_action_id=last.id)
|
||||
assert windowed[1] == everything[1]
|
||||
|
||||
|
||||
def test_reported_total_is_the_whole_story_not_the_window(story):
|
||||
"""Insights says "N of M actions included"; M must not become the window."""
|
||||
db, adventure, settings = story
|
||||
settings.context_token_budget = 4096
|
||||
report = builder.build_context(adventure, settings)[2]
|
||||
assert report["history"]["total"] == ACTION_COUNT
|
||||
assert report["history"]["included"] < ACTION_COUNT
|
||||
|
||||
|
||||
# ------------------------------------------------------------ it is bounded
|
||||
|
||||
def test_building_context_reads_far_less_than_the_whole_story(story, actions_loaded):
|
||||
db, adventure, settings = story
|
||||
# Expire first: expiring afterwards would discard the unflushed change and
|
||||
# silently put the budget back to its default.
|
||||
db.expire_all()
|
||||
# Small enough that the budget, not the length of the story, decides.
|
||||
settings.context_token_budget = 4096
|
||||
actions_loaded["n"] = 0
|
||||
|
||||
report = builder.build_context(adventure, settings)[2]
|
||||
included = report["history"]["included"]
|
||||
|
||||
assert included < ACTION_COUNT, "fixture is too short to prove anything"
|
||||
# The window aims a margin past the budget and re-asks if it fell short, so
|
||||
# it reads somewhat more than it includes. What matters is that the read is
|
||||
# a function of the token budget, not of how long the story has got.
|
||||
assert actions_loaded["n"] < ACTION_COUNT // 2, (
|
||||
f"read {actions_loaded['n']} action rows out of {ACTION_COUNT} to "
|
||||
f"include {included} — the window is not bounding the read"
|
||||
)
|
||||
|
||||
|
||||
def test_window_is_ordered_and_free_of_duplicates(story):
|
||||
"""The window grows by fetching only what it does not already hold, so an
|
||||
off-by-one in the offset would show up as a repeated or missing action."""
|
||||
db, adventure, settings = story
|
||||
window = history.window_covering(adventure, 16384, builder.count_tokens)
|
||||
ids = [a.id for a in window]
|
||||
assert len(ids) == len(set(ids)), "the same action appeared twice in the window"
|
||||
assert ids == sorted(ids), "window must be oldest-first"
|
||||
|
||||
|
||||
# --------------------------------------------- the cursor arithmetic agrees
|
||||
|
||||
def test_helpers_agree_with_the_full_list(story):
|
||||
db, adventure, settings = story
|
||||
actions = history.story_actions(adventure)
|
||||
assert len(actions) == ACTION_COUNT
|
||||
|
||||
assert history.count(adventure) == len(actions)
|
||||
assert history.max_action_index(adventure) == max(a.index for a in actions)
|
||||
assert [a.id for a in history.tail(adventure, 4)] == [a.id for a in actions[-4:]]
|
||||
assert [a.id for a in history.slice_(adventure, 10, 6)] == [a.id for a in actions[10:16]]
|
||||
assert [a.id for a in history.tail_range(adventure, 5, 3)] == \
|
||||
[a.id for a in actions[-8:-5]]
|
||||
assert memorybank.settled_count(adventure) == len(actions) - 1
|
||||
|
||||
for probe in (0, 1, ACTION_COUNT // 2, ACTION_COUNT - 1):
|
||||
target = actions[probe]
|
||||
expected = next(i for i, a in enumerate(actions) if a.index >= target.index)
|
||||
assert history.position_of_index(adventure, target.index) == expected
|
||||
|
||||
|
||||
def test_positions_still_line_up_after_a_middle_action_is_deleted(story):
|
||||
"""The gap in Action.index is exactly what makes positions and indexes
|
||||
diverge — the case that has broken the cursors twice before."""
|
||||
db, adventure, settings = story
|
||||
actions = history.story_actions(adventure)
|
||||
victim = actions[50]
|
||||
db.delete(victim)
|
||||
db.commit()
|
||||
db.expire(adventure)
|
||||
|
||||
remaining = history.story_actions(adventure)
|
||||
assert len(remaining) == ACTION_COUNT - 1
|
||||
assert history.count(adventure) == ACTION_COUNT - 1
|
||||
for probe in (0, 49, 50, 51, ACTION_COUNT - 2):
|
||||
target = remaining[probe]
|
||||
expected = next(i for i, a in enumerate(remaining) if a.index >= target.index)
|
||||
assert history.position_of_index(adventure, target.index) == expected, probe
|
||||
|
||||
|
||||
def test_blank_actions_are_excluded_the_same_way_in_sql_and_python(story):
|
||||
"""SQL and Python must agree on membership or a cursor points elsewhere."""
|
||||
db, adventure, settings = story
|
||||
for blank in ("", " ", "\n", "\t\n "):
|
||||
db.add(models.Action(adventure_id=adventure.id,
|
||||
index=history.max_action_index(adventure) + 1,
|
||||
type="story", text=blank))
|
||||
db.commit()
|
||||
db.expire(adventure)
|
||||
|
||||
# SQL path (relationship not loaded)
|
||||
from_sql = history.count(adventure)
|
||||
# Python path (relationship loaded)
|
||||
adventure.actions # noqa: B018 — force the collection into memory
|
||||
from_python = history.count(adventure)
|
||||
|
||||
assert from_sql == from_python == ACTION_COUNT
|
||||
Reference in New Issue
Block a user