tests/test_memory_budget.py

"""Memory is a budget: assembled context must respect it, and compaction must trigger when
recent turns pile up."""

from __future__ import annotations

from visualnovel import config, memory
from visualnovel.schemas import Character, GameState, Scene, Turn


def _state(n_turns: int = 0) -> GameState:
    s = GameState(
        seed=1,
        vibe="v",
        style_guide="s",
        scene=Scene(id="a", place="A", description="d", present=["moth"]),
        characters={"moth": Character(id="moth", name="Moth", one_line="a moth")},
    )
    for i in range(n_turns):
        s.recent_turns.append(Turn(player=f"p{i}", speaker="moth", dialogue=f"d{i}"))
    return s


def test_context_includes_present_and_input():
    ctx = memory.assemble_context(_state(2), "where am I?")
    assert "Moth" in ctx
    assert "where am I?" in ctx
    assert "SPIRITS ON STAGE" in ctx


def test_context_respects_budget(monkeypatch):
    monkeypatch.setattr(config, "CONTEXT_TOKEN_BUDGET", 50)  # 50 tokens ≈ 200 chars
    ctx = memory.assemble_context(_state(40), "x" * 500)
    assert len(ctx) <= 50 * 4 + len("…(earlier dream trimmed)…\n\n")


def test_player_name_in_context():
    s = _state(1)
    s.flags["player_name"] = "Will"
    ctx = memory.assemble_context(s, "hello")
    assert 'The player is called "Will"' in ctx
    assert 'WILL NOW SAYS: "hello"' in ctx
    # without a name, the classic label is kept
    ctx_anon = memory.assemble_context(_state(1), "hello")
    assert 'THE WANDERER NOW SAYS: "hello"' in ctx_anon


def test_internal_flags_hidden_from_facts():
    s = _state(1)
    s.flags.update(
        {
            "player_name": "Will",
            "beat_started_turn": "3",
            "milestone50_moth": "done",
            "situation_intro": "intro",
            "found_key": "true",  # a real story fact stays visible
        }
    )
    ctx = memory.assemble_context(s, "x")
    assert "FACTS: found_key=true" in ctx
    for hidden in ("beat_started_turn", "milestone50_moth", "situation_intro=", "player_name="):
        assert hidden not in ctx


def test_pacing_nudge_after_lingering_beat():
    s = _state(1)
    s.turn_index = 8  # 8 turns since beat_started_turn=0
    assert "PACING:" in memory.assemble_context(s, "x")
    s.turn_index = 3
    assert "PACING:" not in memory.assemble_context(s, "x")
    s.turn_index = 12
    s.beat = "ended"
    assert "PACING:" not in memory.assemble_context(s, "x")


def test_should_compact():
    assert memory.should_compact(_state(2 * config.RECENT_TURNS_K + 1)) is True
    assert memory.should_compact(_state(2 * config.RECENT_TURNS_K)) is False
    assert memory.should_compact(_state(1)) is False


def test_trim_drops_oldest_recent_first(monkeypatch):
    # Budget large enough for the head + a few turns, too small for all 30 turns:
    # the oldest exchanges must go first, the head and the player line must survive.
    monkeypatch.setattr(config, "CONTEXT_TOKEN_BUDGET", 300)  # ≈ 1200 chars
    ctx = memory.assemble_context(_state(30), "where am I?")
    assert "WORLD STYLE" in ctx
    assert "SPIRITS ON STAGE" in ctx
    assert "where am I?" in ctx
    assert "…(earlier exchange trimmed)…" in ctx
    assert '"d0"' not in ctx  # oldest turn dropped
    assert '"d29"' in ctx  # newest turn kept