"""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 "" 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" )