6krrt/tests/test_prefix_probe.py

"""The prefix-stability probe, and the asymmetry it exists to measure.

Wave 1 item 1.3 of ``plans/token-waste-waves.md``. The probe is only worth
reading on live traffic if it reports the right answer on a case where the
answer is already known, so the load-bearing test here replays the exact
scenario the direct investigation used -- a ten-turn agent session with one
tool result appended and a *literally identical* relevance order on both
turns, so embedding jitter cannot be the explanation -- and pins both halves:

  relevance path: diverges at message 19 of 82, 75% of tokens re-billed
  uniform path:   diverges only at the appended message, nothing rewritten

Everything else here guards the properties that make the probe safe to leave
on: it stores no content, it is bounded in memory, and it never raises on the
request path.

Offline throughout: temp SQLite from ``config/schema.sql``, never the live
router.db.
"""

import json
import sqlite3
from pathlib import Path

import dispatcher
import metrics
import prefix_probe
from context_prune import estimate_tokens, extract_text, prune_context, trim_candidates

ROOT = Path(__file__).resolve().parent.parent
SCHEMA_SQL = (ROOT / "config" / "schema.sql").read_text()

# Production pinch values, so the scenario below is the one that actually runs.
BUDGET, KEEP, MAX_SUM, PROTECTED = 50000, 4, 4000, 20000


def _tool(i: int, size: int = 7000) -> dict:
    return {"role": "tool", "name": f"t{i % 5}", "content": f"r{i} " + ("x" * size)}


def _session(turns: int, tools_per_turn: int = 6) -> list[dict]:
    """A long agent session: user, assistant, then a run of tool results."""
    messages = [{"role": "system", "content": "agent"}]
    k = 0
    for i in range(turns):
        messages.append({"role": "user", "content": f"req {i}"})
        messages.append({"role": "assistant", "content": f"calling {i}"})
        for _ in range(tools_per_turn):
            messages.append(_tool(k))
            k += 1
    return messages


def _prune(messages: list[dict], relevance_order):
    pruned, _stats = prune_context(
        list(messages),
        budget_tokens=BUDGET,
        keep_last_turns=KEEP,
        max_summarize_chars=MAX_SUM,
        relevance_order=relevance_order,
        protected_max_chars=PROTECTED,
    )
    return pruned


def _two_turns(relevance_order_for):
    """Turn N and turn N+1 of the same session, one tool result apart."""
    turn_n = _session(10)
    turn_n1 = turn_n + [_tool(999)]
    candidates, _, _ = trim_candidates(turn_n, KEEP)
    # The candidate set must be identical across the two turns, or the
    # comparison would be measuring the fixture rather than prune_context.
    candidates_next, _, _ = trim_candidates(turn_n1, KEEP)
    assert candidates == candidates_next
    order = relevance_order_for(candidates)
    return _prune(turn_n, order), _prune(turn_n1, order)


def _observe_pair(pruned_n, pruned_n1):
    prefix_probe.reset()
    first = prefix_probe.observe("session-under-test", pruned_n)
    second = prefix_probe.observe("session-under-test", pruned_n1)
    return first, second


# --- the measurement the item exists for ------------------------------------


def test_relevance_path_divergence_matches_the_direct_investigation():
    """The known result, reproduced through the probe rather than by hand.

    A *fixed* relevance order on both turns, so nothing about embeddings can
    be blamed: the growing ``target_save`` alone pulls one more candidate into
    the compressed set, at an arbitrary message POSITION, and everything from
    there on stops matching last turn.
    """
    # Identical on both turns, by construction.
    pruned_n, pruned_n1 = _two_turns(lambda c: list(range(len(c)))[::-1])
    _first, divergence = _observe_pair(pruned_n, pruned_n1)

    assert len(pruned_n) == 81 and len(pruned_n1) == 82
    assert divergence.index == 19, (
        "the relevance path's first divergence moved; the direct test in "
        "plans/token-waste-waves.md measured message index 19 of 82"
    )
    assert divergence.prev_message_count == 81
    assert divergence.destructive, (
        "index 19 is far below the previous turn's 81 messages: this is "
        "rewritten history, not an append"
    )

    total = sum(estimate_tokens(extract_text(m)) for m in pruned_n1)
    share = divergence.tokens_after / total
    assert 0.70 < share < 0.80, (
        f"~75% of the payload's tokens should sit at or after the divergence; "
        f"got {share:.0%}"
    )


def test_uniform_path_diverges_only_where_the_conversation_grew():
    """The control: a contiguous positional region cannot be disturbed.

    ``index == prev_message_count`` is the signature of a payload that only
    grew -- the first position the previous turn did not have. This is the
    property Wave 3 wants the relevance path to gain.
    """
    pruned_n, pruned_n1 = _two_turns(lambda c: None)
    _first, divergence = _observe_pair(pruned_n, pruned_n1)

    assert divergence.index == divergence.prev_message_count == len(pruned_n)
    assert not divergence.destructive
    total = sum(estimate_tokens(extract_text(m)) for m in pruned_n1)
    assert divergence.tokens_after / total < 0.05, (
        "only the appended tool result should be uncached on the uniform path"
    )


def test_the_two_paths_disagree_by_an_order_of_magnitude():
    """Stated as one assertion, because the asymmetry IS the finding."""
    relevance = _observe_pair(*_two_turns(lambda c: list(range(len(c)))[::-1]))[1]
    uniform = _observe_pair(*_two_turns(lambda c: None))[1]
    assert relevance.tokens_after > 10 * uniform.tokens_after


# --- what the probe refuses to do -------------------------------------------


def test_nothing_recoverable_is_retained():
    """The store holds digests and counts -- never a byte of any message.

    This router never stores raw task text anywhere. The probe is the closest
    thing to a counterexample in the codebase, so the invariant is asserted
    rather than trusted: no remembered value may contain any distinctive
    substring of the payload, and every remembered value must be a short hex
    digest or an integer.
    """
    prefix_probe.reset()
    secret = "correct-horse-battery-staple"
    messages = [
        {"role": "system", "content": f"prompt {secret}"},
        {"role": "user", "content": f"please do the thing with {secret}"},
        {"role": "tool", "name": "read", "content": secret * 100},
    ]
    prefix_probe.observe("privacy-session", messages)

    remembered = prefix_probe._store["privacy-session"]
    blob = json.dumps([list(remembered.hashes), list(remembered.tokens)])
    assert secret not in blob
    assert "horse" not in blob and "staple" not in blob
    for digest in remembered.hashes:
        assert len(digest) == 2 * prefix_probe._DIGEST_BYTES
        int(digest, 16)  # hex only; raises otherwise
    for count in remembered.tokens:
        assert isinstance(count, int)


def test_a_first_turn_reports_nothing_rather_than_zero():
    """None means "no comparison", which a stored 0 would misread as total loss."""
    prefix_probe.reset()
    assert prefix_probe.observe("fresh", [{"role": "user", "content": "hi"}]) is None


def test_a_session_without_a_key_is_skipped():
    """session_fingerprint returns None for a textless conversation."""
    prefix_probe.reset()
    assert prefix_probe.observe(None, [{"role": "user", "content": "hi"}]) is None
    assert prefix_probe._store == {}


def test_an_unserializable_message_never_breaks_a_dispatch():
    """Best-effort by contract: this runs before a paid token ships."""
    prefix_probe.reset()

    class Opaque:
        def __repr__(self):
            return "<opaque>"

    messages = [{"role": "user", "content": Opaque()}]
    assert prefix_probe.observe("hostile", messages) is None
    # And the turn after it still compares cleanly.
    assert prefix_probe.observe("hostile", messages) is not None


def test_the_store_is_bounded():
    """An entry is O(messages), unlike session_cache's fixed-size dataclass."""
    prefix_probe.reset()
    for i in range(prefix_probe.MAX_SESSIONS + 10):
        prefix_probe.observe(f"s{i}", [{"role": "user", "content": f"m{i}"}])
    assert len(prefix_probe._store) == prefix_probe.MAX_SESSIONS
    assert "s0" not in prefix_probe._store, "oldest session must be evicted"
    assert f"s{prefix_probe.MAX_SESSIONS + 9}" in prefix_probe._store


def test_a_shrinking_payload_reads_as_destructive():
    """A client that trims its own history loses the cached tail, correctly."""
    long_turn = prefix_probe.fingerprint(
        [{"role": "user", "content": f"m{i}"} for i in range(10)]
    )
    short_turn = prefix_probe.fingerprint(
        [{"role": "user", "content": f"m{i}"} for i in range(4)]
    )
    divergence = prefix_probe.compare(long_turn, short_turn)
    assert divergence.index == 4
    assert divergence.prev_message_count == 10
    assert divergence.destructive
    assert divergence.tokens_after == 0


def test_key_order_is_not_a_divergence():
    """The same message built in a different key order is the same message."""
    a = prefix_probe.message_digest({"role": "tool", "name": "read", "content": "x"})
    b = prefix_probe.message_digest({"content": "x", "name": "read", "role": "tool"})
    assert a == b


# --- storage ----------------------------------------------------------------


def _schema_without_probe_columns() -> str:
    """schema.sql as it stood before this item, for the migration tests."""
    lines = []
    inside = False
    skipping = False
    for line in SCHEMA_SQL.splitlines():
        stripped = line.strip()
        if stripped.startswith("CREATE TABLE IF NOT EXISTS route_decisions"):
            inside = True
        if inside:
            if stripped.startswith("profile "):
                lines.append("    profile                  TEXT")
                skipping = True
                continue
            if skipping:
                if stripped == ");":
                    skipping = False
                    inside = False
                    lines.append(line)
                continue
        lines.append(line)
    return "\n".join(lines)


def test_ensure_route_decisions_adds_the_probe_columns(tmp_path):
    """Additive, idempotent, NULL on every pre-existing row."""
    conn = sqlite3.connect(tmp_path / "migrate.db")
    conn.executescript(_schema_without_probe_columns())
    before = {r[1] for r in conn.execute("PRAGMA table_info(route_decisions)")}
    assert not before & set(metrics._PREFIX_PROBE_COLUMNS)

    conn.execute(
        "INSERT INTO route_decisions (observed_at, kind) VALUES ('t', 'chat')"
    )
    conn.commit()

    dispatcher.ensure_route_decisions(conn)
    dispatcher.ensure_route_decisions(conn)  # idempotent

    columns = [r[1] for r in conn.execute("PRAGMA table_info(route_decisions)")]
    for column in metrics._PREFIX_PROBE_COLUMNS:
        assert columns.count(column) == 1
    row = conn.execute(
        f"SELECT {', '.join(metrics._PREFIX_PROBE_COLUMNS)} FROM route_decisions"
    ).fetchone()
    assert all(value is None for value in row), (
        "a row written before the probe existed cannot say what its prefix "
        "did, and nothing may be inferred for it"
    )
    conn.close()


def test_recent_decisions_tolerates_a_pre_migration_database(tmp_path):
    """metrics is imported by admin, which can open a DB that never booted.

    The live router.db is exactly this until its next restart: the ALTER has
    not run, and an OperationalError here would 500 the whole /metrics
    payload rather than leaving three fields NULL.
    """
    conn = sqlite3.connect(tmp_path / "old.db")
    conn.executescript(_schema_without_probe_columns())
    conn.execute(
        "INSERT INTO route_decisions (observed_at, kind, selected_model) "
        "VALUES ('2026-09-13T00:00:00+00:00', 'chat', 'm')"
    )
    conn.commit()
    conn.row_factory = sqlite3.Row

    rows = metrics.recent_decisions(conn, limit=5)
    conn.close()
    assert len(rows) == 1
    for column in metrics._PREFIX_PROBE_COLUMNS:
        assert column in rows[0] and rows[0][column] is None


def test_persist_route_decision_stores_the_divergence(tmp_path, monkeypatch):
    """The three columns round-trip, and are NULL together when unmeasured."""
    db_path = tmp_path / "persist.db"
    conn = sqlite3.connect(db_path)
    conn.executescript(SCHEMA_SQL)
    conn.commit()
    conn.close()
    monkeypatch.setattr(dispatcher.cfg.database, "path", str(db_path))
    monkeypatch.setattr(dispatcher.cfg.logging, "log_route_decisions", True)

    dispatcher.persist_route_decision(
        "chat",
        selected_model="m",
        selected_provider="neuralwatt",
        prefix_divergence=prefix_probe.Divergence(
            index=19, tokens_after=82595, prev_message_count=81
        ),
    )
    dispatcher.persist_route_decision(
        "chat", selected_model="m", selected_provider="neuralwatt"
    )

    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT prefix_divergence_index, prefix_tokens_after_divergence, "
        "prefix_prev_message_count FROM route_decisions ORDER BY id"
    ).fetchall()
    conn.close()
    assert tuple(rows[0]) == (19, 82595, 81)
    assert tuple(rows[1]) == (None, None, None), (
        "no previous turn must store NULL, not a zero that reads as total "
        "cache loss at message 0"
    )