interactive-story/backend/tests/test_context_performance.py

"""M6: the read paths this milestone touches must not grow a query per row.

M5 spent a review finding on an N+1 in the action list. M6 adds three things
that could each reintroduce one — a memory's provenance, a summary's source
coordinates, and the derived-work status — so each is measured here rather than
argued about.

The assertions are on *growth*, not on an exact count. A fixed number would
break on any unrelated query and teach the next person to raise the number; what
matters is that doubling the rows does not double the queries.

    python -m pytest tests/test_context_performance.py -v
"""

import asyncio

import pytest
from fastapi import Depends
from fastapi.testclient import TestClient
from sqlalchemy import event

from app import auth, limits, memorybank, models, summaries
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 StubEmbedder:
    async def embed(self, texts):
        return [[1.0, 0.0, 0.0] for _ in texts]


@pytest.fixture()
def sql_log():
    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)
    memorybank._vector_cache.clear()
    setup = SessionLocal()
    user = models.User(is_guest=False, email="perf@example.com")
    setup.add(user)
    setup.flush()
    setup.add(models.Settings(
        user_id=user.id, api_key="enc:dummy", model="test-model",
        embedding_model="embed-test", context_token_budget=8000,
        max_output_tokens=400, memory_top_k=5,
    ))
    adventure = models.Adventure(
        user_id=user.id, title="Perf", memory_bank_enabled=True, auto_summarize=True,
    )
    setup.add(adventure)
    setup.flush()
    setup.add(models.Action(adventure_id=adventure.id, type="start", text="A road."))
    setup.commit()
    adv_id, user_id = adventure.id, user.id
    setup.close()

    monkeypatch.setattr(limits, "check_row_cap", lambda *a, **k: None)
    monkeypatch.setattr(adventures.turns, "OpenAICompatibleProvider", ScriptedProvider)
    monkeypatch.setattr(memorybank, "embedding_provider", lambda s: StubEmbedder())
    app.dependency_overrides[auth.get_current_user] = (
        lambda db=Depends(get_db): db.get(models.User, user_id)
    )
    test_client = TestClient(app)
    test_client.adv_id = adv_id
    test_client.user_id = user_id
    try:
        yield test_client
    finally:
        app.dependency_overrides.clear()
        memorybank._vector_cache.clear()
        Base.metadata.drop_all(bind=engine)


def _grow(client, *, turns, memories, summary_rows):
    with SessionLocal() as db:
        adventure = db.get(models.Adventure, client.adv_id)
        for i in range(turns):
            db.add(models.Action(adventure_id=adventure.id,
                                 type="ai" if i % 2 else "do",
                                 text=f"[{i}] The road bends onward. " * 6))
        db.commit()
    with SessionLocal() as db:
        adventure = db.get(models.Adventure, client.adv_id)
        for i in range(memories):
            memory = models.Memory(
                adventure_id=adventure.id, text=f"Memory {i}: something happened.",
                branch_id=adventure.head_branch_id, depth=adventure.head_depth,
                source_start=0, source_end=adventure.head_depth,
            )
            memorybank.set_vector(memory, [1.0, 0.0, 0.0])
            db.add(memory)
        for i in range(summary_rows):
            summaries.record(db, adventure, f"Summary {i}.")
        db.commit()


def _count(sql_log, client) -> int:
    sql_log.clear()
    r = client.get(f"/api/adventures/{client.adv_id}/context")
    assert r.status_code == 200, r.text[:200]
    return len(sql_log)


def test_assembling_context_does_not_cost_a_query_per_memory(client, sql_log):
    """A memory's provenance is fetched in the same read as its text, so more
    memories must not mean more queries."""
    _grow(client, turns=10, memories=5, summary_rows=1)
    small = _count(sql_log, client)
    _grow(client, turns=0, memories=25, summary_rows=0)
    large = _count(sql_log, client)

    assert large <= small + 2, (
        f"{small} queries with 5 memories, {large} with 30 — "
        "the context read is paying per memory"
    )


def test_assembling_context_does_not_cost_a_query_per_summary(client, sql_log):
    """Only the eligible summary is read, however many are retained."""
    _grow(client, turns=10, memories=2, summary_rows=2)
    small = _count(sql_log, client)
    _grow(client, turns=0, memories=0, summary_rows=30)
    large = _count(sql_log, client)

    assert large <= small + 2, (
        f"{small} queries with 2 summaries, {large} with 32 — "
        "the context read is paying per summary"
    )


def test_assembling_context_does_not_cost_a_query_per_turn(client, sql_log):
    """The history window is one read, not one per action."""
    _grow(client, turns=10, memories=2, summary_rows=1)
    small = _count(sql_log, client)
    _grow(client, turns=60, memories=0, summary_rows=0)
    large = _count(sql_log, client)

    assert large <= small + 2, (
        f"{small} queries at 10 turns, {large} at 70 — "
        "the context read is paying per turn"
    )


def test_the_derived_status_endpoint_does_not_pay_per_summary(client, sql_log):
    """The listing resolves the eligible summary once, not once per row."""
    _grow(client, turns=6, memories=1, summary_rows=3)
    sql_log.clear()
    assert client.get(f"/api/adventures/{client.adv_id}/derived").status_code == 200
    small = len(sql_log)

    _grow(client, turns=0, memories=0, summary_rows=30)
    sql_log.clear()
    assert client.get(f"/api/adventures/{client.adv_id}/derived").status_code == 200
    large = len(sql_log)

    assert large <= small + 1, (
        f"{small} queries with 3 summaries, {large} with 33"
    )


def test_the_context_size_stops_growing_once_the_budget_is_reached(client):
    """The companion to the query counts: more story, not more prompt.

    Measured from a story that already fills the budget. Comparing a short story
    to a long one only shows that the prompt grew, which it is supposed to do
    until it reaches the ceiling; what F03 is about is that it stops there.
    """
    _grow(client, turns=140, memories=3, summary_rows=1)
    filled = client.get(f"/api/adventures/{client.adv_id}/context").json()
    budget = filled["tokens"]["budget"]
    assert filled["tokens"]["total"] > budget * 0.5, (
        "the fixture never filled the budget, so this proves nothing"
    )

    _grow(client, turns=280, memories=0, summary_rows=0)
    doubled = client.get(f"/api/adventures/{client.adv_id}/context").json()

    assert doubled["history"]["total"] > filled["history"]["total"] * 2, "fixture too small"
    assert doubled["tokens"]["total"] <= budget
    # Three times the story, and the prompt does not move.
    assert doubled["tokens"]["total"] <= filled["tokens"]["total"] + 50, (
        f"{filled['tokens']['total']} -> {doubled['tokens']['total']} tokens "
        f"while the story went from {filled['history']['total']} to "
        f"{doubled['history']['total']} actions"
    )
    # And it is bounded by the budget rather than by the length of the story.
    assert doubled["history"]["included"] < doubled["history"]["total"]