interactive-story/backend/tests/test_v11_context_reserve.py

"""v1.1 WP-A1: a deliberate safety reserve, and a turn the server cut is not silent.

M11 made the verified window a ceiling. It did not make the application's count
the server's count. The application counts with `cl100k_base`, the narrator with
its own tokenizer, and the v1 evidence left 23-42 real tokens between the largest
prompt and the edge of a 16,384 window. Past that edge Ollama does not refuse.
Measured against the reference CPU host (Ollama 0.33, a 4,096 window), a
6,316-token prompt came back 200 with `prompt_tokens` 2,050: the front of the
prompt, which in this design is the narrator's rules and the canon, was gone.

So the tests below are in three halves.

**The reserve.** `max(256, ceil(5% of the effective window))`, taken from the
budget before any history is chosen, on top of an exact reply allocation.

**The arithmetic.** The assembled prompt, plus the application text the provider
adds to every request, plus the reply allocation, plus the reserve, fits the
effective window. Protected context that cannot fit that way fails before the
model is called.

**The accounting.** Where the server reports how many prompt tokens it read, the
turn records `fits`, `exceeded` or `truncation_suspected`. Where it reports
nothing, the turn says `unknown`, never `fits`. A discrepancy found after the
reply is recorded and shown; it never costs the reader an accepted turn.

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

import pytest
from fastapi import Depends
from fastapi.testclient import TestClient
from sqlalchemy.orm import undefer

from app import auth, contextwindow, limits, models
from app.context import builder
from app.database import Base, SessionLocal, engine, get_db
from app.main import app
from app.providers.openai_compatible import CHAT_CONTINUE_HINT, OpenAICompatibleProvider
from app.routers import adventures

from fakes import ScriptedProvider

ENDPOINT = "http://127.0.0.1:11434/v1"


@pytest.fixture(autouse=True)
def _clear_window_cache():
    contextwindow.cache_clear()
    yield
    contextwindow.cache_clear()


# ------------------------------------------------------------- the reserve

@pytest.mark.parametrize("window, reserve", [
    (1024, 256),
    (4096, 256),     # 5% is 204.8, so the floor holds
    (5120, 256),     # exactly 5% is the floor
    (5121, 257),     # 256.05 rounds up
    (8192, 410),     # 409.6 rounds up
    (16384, 820),    # 819.2 rounds up
    (32768, 1639),   # 1638.4 rounds up
])
def test_the_reserve_is_the_larger_of_the_floor_and_five_percent_rounded_up(window, reserve):
    assert contextwindow.safety_reserve(window) == reserve


def test_the_reserve_is_far_larger_than_the_v1_margin_at_the_evidence_window():
    """The v1 evidence left 23-42 tokens at 16,384. 64 tokens of slack was all
    the arithmetic kept for drift and separators together."""
    assert contextwindow.safety_reserve(16384) >= 10 * 64


# ---------------------------------------------------------- the arithmetic

@pytest.fixture()
def client(monkeypatch):
    Base.metadata.create_all(bind=engine)
    setup = SessionLocal()
    user = models.User(is_guest=False, email="v11reserve@example.com")
    setup.add(user)
    setup.flush()
    setup.add(models.Settings(
        user_id=user.id, model="qwen2.5:3b-instruct", endpoint_url=ENDPOINT,
        embedding_model="", context_token_budget=16384, max_output_tokens=500,
    ))
    adventure = models.Adventure(
        user_id=user.id, title="Reserved",
        campaign_canon={"rules": [
            "The abbey seal has never been broken.",
            "The sealed crypt is named CANON-SENTINEL-RESERVE-5120.",
        ]},
    )
    setup.add(adventure)
    setup.flush()
    setup.add(models.Action(
        adventure_id=adventure.id, type="start", text="Rain over Westhaven."))
    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)
    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
    try:
        yield test_client
    finally:
        app.dependency_overrides.clear()
        adventures.turns._active_turns.clear()
        Base.metadata.drop_all(bind=engine)


def _long_story(adv_id, turns=120):
    from app import tree

    with SessionLocal() as db:
        adventure = db.get(models.Adventure, adv_id)
        for i in range(turns):
            for kind, text in (
                ("do", f"I search the {i}th chamber of the undercroft."),
                ("ai", "The lantern gutters. " + ("Cold stone, and older dust. " * 40)),
            ):
                action = models.Action(adventure_id=adv_id, type=kind, text=text)
                db.add(action)
                db.flush()
                tree.place_action(db, adventure, action)
        db.commit()


def _settings(**changes):
    with SessionLocal() as db:
        settings = db.query(models.Settings).first()
        for key, value in changes.items():
            setattr(settings, key, value)
        db.commit()


def _build(client, window):
    with SessionLocal() as db:
        adventure = db.get(models.Adventure, client.adv_id)
        settings = db.query(models.Settings).first()
        return builder.build_context(adventure, settings, window=window)


def _sent(system, story) -> int:
    """What the provider actually sends in chat mode, by the application's count."""
    return (builder.count_tokens(system) + builder.count_tokens(story)
            + builder.count_tokens(CHAT_CONTINUE_HINT))


CONFIGURATIONS = {
    "verified 4,096": (dict(context_token_budget=16384),
                       contextwindow.Window(4096, contextwindow.LOADED)),
    "verified 8,192": (dict(context_token_budget=16384),
                       contextwindow.Window(8192, contextwindow.PARAMETERS)),
    "verified 16,384": (dict(context_token_budget=16384),
                        contextwindow.Window(16384, contextwindow.LOADED)),
    "declared 6,000": (dict(context_token_budget=16384),
                       contextwindow.Window(6000, contextwindow.DECLARED)),
    "unverified, configured 12,000": (dict(context_token_budget=12000),
                                      contextwindow.UNVERIFIED),
}


@pytest.mark.parametrize("name", list(CONFIGURATIONS))
def test_the_prompt_leaves_the_reply_and_the_reserve_free(client, name):
    """A1-2, on the assembled text rather than the builder's own arithmetic."""
    changes, window = CONFIGURATIONS[name]
    _settings(**changes)
    _long_story(client.adv_id, turns=120)
    system, story, report = _build(client, window)
    tokens = report["tokens"]
    budget = tokens["budget"]

    assert tokens["safety_reserve"] == contextwindow.safety_reserve(budget)
    assert tokens["output_reserve"] == 500
    sent = _sent(system, story)
    assert sent + tokens["output_reserve"] + tokens["safety_reserve"] <= budget, (
        name, sent, tokens)
    # The history is what gave way, not the canon.
    assert "CANON-SENTINEL-RESERVE-5120" in system
    assert report["history"]["included"] < report["history"]["total"]


def test_the_report_prices_the_text_the_provider_adds(client):
    """The chat hint rides on every request and was never counted."""
    _, _, report = _build(client, contextwindow.Window(4096, contextwindow.LOADED))
    tokens = report["tokens"]
    assert tokens["transport"] >= builder.count_tokens(CHAT_CONTINUE_HINT)
    assert tokens["estimate"] == tokens["total"] + builder.count_tokens(CHAT_CONTINUE_HINT)


def test_the_reserve_follows_the_effective_window_not_the_setting(client):
    """5% of a 4,096 server, not 5% of a 16,384 setting it will never read."""
    _, _, capped = _build(client, contextwindow.Window(4096, contextwindow.LOADED))
    _, _, full = _build(client, contextwindow.Window(16384, contextwindow.LOADED))
    assert capped["tokens"]["safety_reserve"] == 256
    assert full["tokens"]["safety_reserve"] == 820


def test_protected_context_that_only_fits_without_the_reserve_fails_explicitly(client):
    """A1-3 in the builder. Before v1.1 this prompt would have been built.

    The canon is sized so that protected text plus the reply fits a 4,096 window
    with room to spare, and does not fit once the 256-token reserve is taken.
    """
    small = contextwindow.Window(4096, contextwindow.LOADED)
    # Measured with a window large enough never to overflow, because repeated
    # text merges tokens at its seams and cannot be priced by multiplication.
    roomy = contextwindow.Window(32768, contextwindow.LOADED)
    rules = None
    with SessionLocal() as db:
        adventure = db.get(models.Adventure, client.adv_id)
        settings = db.query(models.Settings).first()
        base_rules = list(adventure.campaign_canon["rules"])
        filler = "The bell tolls once for every name in the ledger."
        copies = 1
        while True:
            candidate = base_rules + [" ".join([filler] * copies)]
            adventure.campaign_canon = {"rules": candidate}
            _, _, measured = builder.build_context(adventure, settings, window=roomy)
            t = measured["tokens"]
            # What protected context costs at 4,096, without the reserve.
            without_reserve = t["protected"] + t["transport"] + t["output_reserve"]
            if without_reserve + 64 >= 4096 - 60:
                break
            copies += 1
        rules = candidate
        adventure.campaign_canon = {"rules": rules}
        db.commit()

        # The case this test is about: v1's arithmetic, with its 64-token margin,
        # would have built this prompt. v1.1's reserve does not fit.
        assert without_reserve + 64 < 4096
        assert without_reserve + contextwindow.safety_reserve(4096) >= 4096

        with pytest.raises(builder.ContextOverflow) as caught:
            builder.build_context(adventure, settings, window=small)
    message = str(caught.value)
    assert "safety" in message
    assert "load the model with a larger window" in message


def test_an_overflowing_turn_never_reaches_the_model(client, monkeypatch):
    """A1-3 end to end: the refusal happens before the provider is called."""
    async def verified(endpoint, model, declared=None, use_cache=True):
        return contextwindow.Window(1024, contextwindow.LOADED)

    monkeypatch.setattr(adventures.turns.contextwindow, "probe", verified)
    ScriptedProvider.replies = ["This must never be generated."]
    response = client.post(f"/api/adventures/{client.adv_id}/actions",
                           json={"type": "do", "text": "open the crypt"})
    assert response.status_code == 200
    assert "safety" in response.text
    assert ScriptedProvider.calls == 0
    with SessionLocal() as db:
        assert db.query(models.Action).filter_by(
            adventure_id=client.adv_id, type="ai").count() == 0


# ---------------------------------------------------------- the accounting

def _classify(prompt_tokens=None, *, usage=None, estimate=3500, budget=4096,
              output=500, verified=True):
    if usage is None and prompt_tokens is not None:
        usage = {"prompt_tokens": prompt_tokens, "completion_tokens": 40}
    return contextwindow.classify_usage(
        usage, estimate=estimate, budget=budget, max_output_tokens=output,
        window_verified=verified,
    )


def test_a_prompt_the_server_read_in_full_fits():
    # The 13-token chat-template overhead measured against the real server.
    result = _classify(3513)
    assert result["status"] == contextwindow.FITS
    assert result["server_prompt_tokens"] == 3513
    assert result["difference"] == 13
    assert result["safety_reserve"] == 256
    assert result["observed_margin"] == 4096 - 500 - 3513


def test_a_server_that_counts_more_than_the_reserve_allows_is_exceeded():
    """The prompt plus the reply allocation no longer fits the window."""
    result = _classify(3700)
    assert result["status"] == contextwindow.EXCEEDED
    assert result["observed_margin"] < 0


def test_a_server_that_read_far_less_than_was_sent_is_suspected_of_truncating():
    """The real shape: 6,316 sent, 2,050 read, HTTP 200, no error."""
    result = _classify(2050, estimate=6316)
    assert result["status"] == contextwindow.TRUNCATION_SUSPECTED
    assert result["difference"] == 2050 - 6316


def test_a_small_undercount_is_tokenizer_drift_not_truncation():
    """A tokenizer thriftier than `cl100k_base` reads fewer tokens honestly. Only
    a shortfall larger than the reserve is called truncation."""
    assert _classify(3500 - 255)["status"] == contextwindow.FITS
    assert _classify(3500 - 257)["status"] == contextwindow.TRUNCATION_SUSPECTED


@pytest.mark.parametrize("usage", [
    None,
    {},
    {"completion_tokens": 40},
    {"prompt_tokens": 0},
    {"prompt_tokens": "3500"},
    {"prompt_tokens": -1},
])
def test_no_usable_count_is_unknown_never_fits(usage):
    result = _classify(usage=usage)
    assert result["status"] == contextwindow.UNKNOWN
    assert result["server_prompt_tokens"] is None
    assert result["observed_margin"] is None


def test_the_accounting_says_when_the_window_itself_was_not_verified():
    result = _classify(3513, verified=False)
    assert result["status"] == contextwindow.FITS
    assert result["window_verified"] is False
    assert "not verified" in result["detail"]


def test_the_stream_asks_the_server_to_report_its_usage():
    """Measured: Ollama 0.33 sends no usage in a stream unless asked."""
    provider = OpenAICompatibleProvider(ENDPOINT, "m")
    from app.providers.base import PromptParts

    for mode in ("chat", "completion"):
        provider.api_mode = mode
        _url, body = provider._request(PromptParts(system="s", story="t"), 0.7, 50)
        assert body["stream"] is True
        assert body["stream_options"] == {"include_usage": True}


def _latest_ai(adv_id):
    with SessionLocal() as db:
        return (
            db.query(models.Action)
            .filter(models.Action.adventure_id == adv_id, models.Action.type == "ai")
            .options(undefer(models.Action.context_snapshot))
            .order_by(models.Action.id.desc()).first()
        )


def _play(client, monkeypatch, usage, window=4096, reply="The crypt is still sealed."):
    async def verified(endpoint, model, declared=None, use_cache=True):
        return contextwindow.Window(window, contextwindow.LOADED, 32768, "fake")

    monkeypatch.setattr(adventures.turns.contextwindow, "probe", verified)
    monkeypatch.setattr(ScriptedProvider, "last_usage", usage)
    ScriptedProvider.replies = [reply]
    return client.post(f"/api/adventures/{client.adv_id}/actions",
                       json={"type": "do", "text": "look at the seal"})


def test_a_turn_records_what_the_server_read(client, monkeypatch):
    response = _play(client, monkeypatch, None)
    assert response.status_code == 200, response.text[:300]
    estimate = _latest_ai(client.adv_id).context_snapshot["tokens"]["estimate"]

    response = _play(client, monkeypatch,
                     {"prompt_tokens": estimate + 13, "completion_tokens": 9})
    assert response.status_code == 200, response.text[:300]
    snapshot = _latest_ai(client.adv_id).context_snapshot
    accounting = snapshot["accounting"]
    assert accounting["status"] == contextwindow.FITS
    assert accounting["server_prompt_tokens"] == estimate + 13
    assert accounting["estimate"] == snapshot["tokens"]["estimate"]
    assert '"accounting"' in response.text
    assert contextwindow.FITS in response.text


def test_a_turn_with_no_reported_usage_is_unknown(client, monkeypatch):
    response = _play(client, monkeypatch, None)
    assert response.status_code == 200
    assert _latest_ai(client.adv_id).context_snapshot["accounting"]["status"] == (
        contextwindow.UNKNOWN)


def test_a_suspected_truncation_keeps_the_turn_and_says_so(client, monkeypatch, caplog):
    """A1-7 and A1-8. The reader watched the narration arrive; it stays."""
    response = _play(client, monkeypatch, {"prompt_tokens": 12, "completion_tokens": 9},
                     reply="The seal holds, and the rain goes on.")
    assert response.status_code == 200, response.text[:300]
    action = _latest_ai(client.adv_id)
    assert action is not None
    assert action.text == "The seal holds, and the rain goes on."
    accounting = action.context_snapshot["accounting"]
    assert accounting["status"] == contextwindow.TRUNCATION_SUSPECTED
    assert contextwindow.TRUNCATION_SUSPECTED in response.text
    assert any(contextwindow.TRUNCATION_SUSPECTED in r.getMessage() for r in caplog.records)

    # Inspectable afterwards through the same route the context panel reads.
    context = client.get(
        f"/api/adventures/{client.adv_id}/actions/{action.id}/context")
    assert context.status_code == 200
    assert context.json()["accounting"]["status"] == contextwindow.TRUNCATION_SUSPECTED


def test_each_attempt_keeps_its_own_accounting_when_the_live_flag_moves():
    """Found by the A2 long run. Accounting belongs to one API call, not to the
    turn's shared prompt. A retry demotes the old attempt, and a take selection
    hands the prompt from one attempt to another. Neither may drop an attempt's
    accounting or give it another attempt's."""
    from app import attempts

    class Node:
        def __init__(self, snapshot):
            self.context_snapshot = snapshot

    shared = {"tokens": {"estimate": 3000}, "sections": [], "window": {"verified": True}}
    first = Node(shared | {"raw_output": "one", "usage": {"prompt_tokens": 3015},
                           "accounting": {"status": contextwindow.FITS, "server_prompt_tokens": 3015}})
    second = Node({"raw_output": "two", "usage": {"prompt_tokens": 12},
                   "accounting": {"status": contextwindow.TRUNCATION_SUSPECTED,
                                  "server_prompt_tokens": 12}})

    # Superseded by a retry: the old attempt keeps only its own slices.
    attempts.keep_own_slices(Node(dict(first.context_snapshot)))
    demoted = Node(dict(first.context_snapshot))
    attempts.keep_own_slices(demoted)
    assert demoted.context_snapshot["accounting"]["server_prompt_tokens"] == 3015
    assert "tokens" not in demoted.context_snapshot

    # The prompt moves to the second attempt; each keeps its own accounting.
    attempts.hand_over_the_prompt(first, second)
    assert second.context_snapshot["tokens"] == {"estimate": 3000}
    assert second.context_snapshot["accounting"]["status"] == contextwindow.TRUNCATION_SUSPECTED
    assert second.context_snapshot["accounting"]["server_prompt_tokens"] == 12
    assert first.context_snapshot["accounting"]["status"] == contextwindow.FITS
    assert "tokens" not in first.context_snapshot


def test_a_retry_leaves_each_take_with_its_own_accounting(client, monkeypatch):
    """End to end, through the real retry route. Before the fix the live take
    inherited the superseded take's accounting, so the inspector could show one
    call's server count as another's."""
    response = _play(client, monkeypatch, None, reply="The first take.")
    assert response.status_code == 200, response.text[:300]
    estimate = _latest_ai(client.adv_id).context_snapshot["tokens"]["estimate"]
    # Replay the first take with a real count, so it has accounting of its own.
    with SessionLocal() as db:
        first = (db.query(models.Action)
                 .filter(models.Action.adventure_id == client.adv_id,
                         models.Action.type == "ai")
                 .options(undefer(models.Action.context_snapshot)).first())
        snapshot = dict(first.context_snapshot)
        snapshot["accounting"] = contextwindow.classify_usage(
            {"prompt_tokens": estimate + 15}, estimate=estimate, budget=4096,
            max_output_tokens=500, window_verified=True)
        first.context_snapshot = snapshot
        db.commit()
        first_id = first.id

    monkeypatch.setattr(ScriptedProvider, "last_usage",
                        {"prompt_tokens": 12, "completion_tokens": 9})
    ScriptedProvider.replies = ["The second take."]
    retried = client.post(f"/api/adventures/{client.adv_id}/retry")
    assert retried.status_code == 200, retried.text[:300]

    with SessionLocal() as db:
        rows = (db.query(models.Action)
                .filter(models.Action.adventure_id == client.adv_id,
                        models.Action.type == "ai")
                .options(undefer(models.Action.context_snapshot))
                .order_by(models.Action.id).all())
        by_id = {row.id: row for row in rows}
        old = by_id[first_id]
        new = [row for row in rows if row.id != first_id][-1]
        assert new.text == "The second take."
        assert new.live and not old.live
        # The superseded take keeps its own accounting and gives up the prompt.
        assert old.context_snapshot["accounting"]["status"] == contextwindow.FITS
        assert old.context_snapshot["accounting"]["server_prompt_tokens"] == estimate + 15
        assert "tokens" not in old.context_snapshot
        # The live take carries the prompt and its own accounting, not the old one's.
        assert "tokens" in new.context_snapshot
        assert new.context_snapshot["accounting"]["status"] == (
            contextwindow.TRUNCATION_SUSPECTED)
        assert new.context_snapshot["accounting"]["server_prompt_tokens"] == 12


def test_an_exceeded_turn_is_also_kept(client, monkeypatch):
    response = _play(client, monkeypatch, {"prompt_tokens": 3900, "completion_tokens": 9})
    assert response.status_code == 200
    action = _latest_ai(client.adv_id)
    assert action.text == "The crypt is still sealed."
    assert action.context_snapshot["accounting"]["status"] == contextwindow.EXCEEDED