Files
2026-07-13 13:12:33 +08:00

284 lines
10 KiB
Python

"""RouterDecisionWriter: inserts, sanitization, pruning, purge, rehydration reads.
The sanitization tests enforce the V017 privacy bar: ``trail``/``flags``
columns may contain only enum-like tokens and numbers — never prompt text.
"""
from __future__ import annotations
import json
import sqlite3
from pathlib import Path
from opensquilla.persistence.migrator import apply_pending
from opensquilla.persistence.router_decision_writer import (
RouterDecisionWriter,
open_router_decision_writer,
sanitize_flags,
sanitize_probs,
sanitize_trail,
)
MIGRATIONS_DIR = Path(__file__).resolve().parents[2] / "migrations"
PROMPT_SENTINEL = "please summarize my confidential quarterly report"
def _make_writer(tmp_path: Path, **kwargs) -> tuple[RouterDecisionWriter, str]:
db = str(tmp_path / "sessions.sqlite")
apply_pending(db, MIGRATIONS_DIR)
conn = sqlite3.connect(db, check_same_thread=False, isolation_level=None)
conn.row_factory = sqlite3.Row
return RouterDecisionWriter(conn, **kwargs), db
def _base_record(**overrides) -> dict:
record = {
"decision_id": "a" * 32,
"session_key": "agent:main:webchat:s1",
"turn_index": 0,
"ts_ms": 1_000_000,
"classifier": "v4_phase3",
"proposed_tier": "c1",
"confidence": 0.91,
"probs": [0.05, 0.91, 0.03, 0.01],
"flags": ["code", "multi_step"],
"final_tier": "c2",
"provider": "openrouter",
"model": "deepseek/deepseek-chat",
"thinking_level": "medium",
"source": "v4_phase3",
"trail": [
{"stage": "classify", "tier": "c1", "route_class": "R1"},
{"stage": "anti_downgrade", "applied": True, "previous_tier": "c2"},
{"stage": "final", "tier": "c2", "route_class": "R2"},
],
"baseline_model": "anthropic/claude-sonnet",
"savings_pct": 42.5,
"executed_kind": "single",
"ensemble_profile": None,
"fallback_hops": 0,
}
record.update(overrides)
return record
def _rows(db: str) -> list[sqlite3.Row]:
conn = sqlite3.connect(db)
conn.row_factory = sqlite3.Row
try:
return conn.execute("SELECT * FROM router_decisions ORDER BY ts_ms").fetchall()
finally:
conn.close()
def test_open_router_decision_writer_boot_constructor(tmp_path: Path) -> None:
"""The boot-path constructor produces a working writer over a real DB."""
db = str(tmp_path / "sessions.sqlite")
apply_pending(db, MIGRATIONS_DIR)
writer = open_router_decision_writer(db, retention_days=7)
try:
assert writer.record_decision(_base_record()) is True
assert len(_rows(db)) == 1
finally:
writer.close()
def test_record_decision_inserts_row(tmp_path: Path) -> None:
writer, db = _make_writer(tmp_path)
assert writer.record_decision(_base_record()) is True
rows = _rows(db)
assert len(rows) == 1
row = rows[0]
assert row["session_key"] == "agent:main:webchat:s1"
assert row["proposed_tier"] == "c1"
assert row["final_tier"] == "c2"
assert row["executed_kind"] == "single"
assert row["savings_pct"] == 42.5
assert json.loads(row["probs"]) == [0.05, 0.91, 0.03, 0.01]
assert json.loads(row["flags"]) == ["code", "multi_step"]
assert json.loads(row["trail"])[0] == {"stage": "classify", "tier": "c1", "route_class": "R1"}
writer.close()
def test_record_decision_is_best_effort_after_close(tmp_path: Path) -> None:
writer, _db = _make_writer(tmp_path)
writer.close()
assert writer.record_decision(_base_record()) is False # no raise
def test_record_decision_rejects_missing_identity(tmp_path: Path) -> None:
writer, db = _make_writer(tmp_path)
assert writer.record_decision(_base_record(decision_id=None)) is False
assert writer.record_decision(_base_record(session_key="")) is False
assert _rows(db) == []
writer.close()
def test_no_prompt_text_reaches_any_column(tmp_path: Path) -> None:
"""V017 privacy bar: free text in any field is dropped before insert."""
writer, db = _make_writer(tmp_path)
dirty = _base_record(
classifier=PROMPT_SENTINEL,
model=PROMPT_SENTINEL,
flags=["code", PROMPT_SENTINEL, {"nested": PROMPT_SENTINEL}],
probs=[0.5, PROMPT_SENTINEL, float("nan")],
trail=[
{"stage": "classify", "tier": "c1", "note": PROMPT_SENTINEL},
{PROMPT_SENTINEL: "value"},
{"stage": "final", "tier": PROMPT_SENTINEL},
PROMPT_SENTINEL,
],
)
assert writer.record_decision(dirty) is True
row = _rows(db)[0]
serialized = json.dumps({key: row[key] for key in row.keys()})
assert PROMPT_SENTINEL not in serialized
assert json.loads(row["flags"]) == ["code"]
assert json.loads(row["probs"]) == [0.5]
trail = json.loads(row["trail"])
assert trail == [
{"stage": "classify", "tier": "c1"},
{"stage": "final"},
]
writer.close()
def test_trail_and_flags_contain_only_enum_tokens_and_numbers() -> None:
"""Schema test for the sanitizers themselves (no-prompt-text bar)."""
token_ok = "c2"
text_bad = "user said: my ssn is 123-45-6789"
flags = sanitize_flags([token_ok, text_bad, 7, None, "R1"])
assert flags == ["c2", "R1"]
trail = sanitize_trail(
[
{
"stage": "confidence_gate",
"applied": False,
"threshold": 0.5,
"echo": text_bad,
}
]
)
assert trail == [{"stage": "confidence_gate", "applied": False, "threshold": 0.5}]
for entry in trail:
for value in entry.values():
assert isinstance(value, (bool, int, float)) or (
isinstance(value, str) and " " not in value and len(value) <= 128
)
assert sanitize_probs([0.1, "text", None, True, 0.9]) == [0.1, 0.9]
def test_write_time_opportunistic_pruning(tmp_path: Path) -> None:
now_ms = 10_000_000_000_000
writer, db = _make_writer(
tmp_path,
retention_days=30,
prune_every=4,
clock=lambda: now_ms,
)
stale_ts = now_ms - 31 * 24 * 60 * 60 * 1000
fresh_ts = now_ms - 1000
writer.record_decision(_base_record(decision_id="old1", ts_ms=stale_ts))
writer.record_decision(_base_record(decision_id="old2", ts_ms=stale_ts))
writer.record_decision(_base_record(decision_id="new1", ts_ms=fresh_ts))
assert len(_rows(db)) == 3 # prune not yet due
writer.record_decision(_base_record(decision_id="new2", ts_ms=fresh_ts))
remaining = {row["decision_id"] for row in _rows(db)}
assert remaining == {"new1", "new2"}
writer.close()
def test_purge_for_session(tmp_path: Path) -> None:
writer, db = _make_writer(tmp_path)
writer.record_decision(_base_record(decision_id="d1", session_key="agent:a"))
writer.record_decision(_base_record(decision_id="d2", session_key="agent:a"))
writer.record_decision(_base_record(decision_id="d3", session_key="agent:b"))
assert writer.purge_for_session("agent:a") == 2
assert {row["session_key"] for row in _rows(db)} == {"agent:b"}
writer.close()
def test_load_recent_history_bounds_window_and_per_session(tmp_path: Path) -> None:
now_ms = 5_000_000_000_000
writer, _db = _make_writer(tmp_path, clock=lambda: now_ms)
# 7 recent decisions in session a -> only the last 5 come back.
for index in range(7):
writer.record_decision(
_base_record(
decision_id=f"a{index}",
session_key="agent:a",
turn_index=index,
ts_ms=now_ms - (7 - index) * 1000,
final_tier="c2" if index % 2 else "c1",
)
)
# One stale decision outside the 1800s window -> excluded.
writer.record_decision(
_base_record(decision_id="b0", session_key="agent:b", ts_ms=now_ms - 2000 * 1000)
)
# One fresh decision in session c.
writer.record_decision(
_base_record(decision_id="c0", session_key="agent:c", ts_ms=now_ms - 500)
)
grouped = writer.load_recent_history(window_seconds=1800, per_session=5)
assert set(grouped) == {"agent:a", "agent:c"}
assert [row["turn_index"] for row in grouped["agent:a"]] == [2, 3, 4, 5, 6]
assert grouped["agent:a"][-1]["final_tier"] == "c1" # index 6 -> even -> c1
assert len(grouped["agent:c"]) == 1
writer.close()
def test_list_decisions_orders_newest_first_and_parses_json(tmp_path: Path) -> None:
writer, _db = _make_writer(tmp_path)
for index in range(3):
writer.record_decision(
_base_record(decision_id=f"d{index}", ts_ms=1_000 * (index + 1))
)
rows = writer.list_decisions()
assert [row["decision_id"] for row in rows] == ["d2", "d1", "d0"]
top = rows[0]
# Whitelisted columns only, JSON columns parsed back into structures.
assert set(top) == {
"decision_id", "session_key", "turn_index", "ts_ms", "classifier",
"proposed_tier", "confidence", "probs", "flags", "final_tier",
"provider", "model", "thinking_level", "source", "trail",
"baseline_model", "savings_pct", "executed_kind", "ensemble_profile",
"fallback_hops",
}
assert top["probs"] == [0.05, 0.91, 0.03, 0.01]
assert top["flags"] == ["code", "multi_step"]
assert top["trail"][0] == {"stage": "classify", "tier": "c1", "route_class": "R1"}
assert top["savings_pct"] == 42.5
writer.close()
def test_list_decisions_filters_pages_and_limits(tmp_path: Path) -> None:
writer, _db = _make_writer(tmp_path)
for index in range(4):
writer.record_decision(
_base_record(decision_id=f"a{index}", session_key="agent:a", ts_ms=1_000 + index)
)
writer.record_decision(
_base_record(decision_id="b0", session_key="agent:b", ts_ms=9_000)
)
only_a = writer.list_decisions(session_key="agent:a")
assert [row["decision_id"] for row in only_a] == ["a3", "a2", "a1", "a0"]
page = writer.list_decisions(session_key="agent:a", before_ts_ms=1_002, limit=1)
assert [row["decision_id"] for row in page] == ["a1"]
assert len(writer.list_decisions(limit=2)) == 2
writer.close()
def test_list_decisions_is_best_effort_after_close(tmp_path: Path) -> None:
writer, _db = _make_writer(tmp_path)
writer.record_decision(_base_record())
writer.close()
assert writer.list_decisions() == [] # no raise