rune/scripts/check_token_optimizer_parity.py

#!/usr/bin/env python3
"""Live parity: token_optimizer Rust kernel vs Python oracle (dichotomy wave,
close-PARTIAL gap pass). Layers:

  1. FROZEN — the gap truth corpus (src/rust/token_optimizer/frozen_truth.json,
     generated by scripts/gen_token_optimizer_truth.py under the
     scripts/to_freeze_seams rig against the LIVE oracle with both core
     pins python): the get_optimizer construction rows replayed through the
     ctypes FFI op (envelope state vs the frozen oracle observation,
     last_save_time excluded — the facade stamps the monotonic clock after
     _load); the pin contracts asserted in BOTH directions.
  2. FUZZ — 40 seeded optimizer instances (seed 2479), each running a
     random op script (static/dynamic/optimize/status/stats/save_if_dirty/
     flush/mark_dirty) with a PATCHED module clock (frozen datetime +
     counter monotonic): the full observation tape (per-op result + state +
     storage file) must be IDENTICAL with RUNE_TOKEN_OPTIMIZER_CORE=python
     (oracle floor) and with the Rust kernel live through the facade's
     _run_kernel delegation.
  3. LIVE — dual-mode E2E on the gap fns: the real module-level
     get_optimizer()/optimize_prompt() API, the 8 storage-file construction
     variants, and the _write_snapshot conversion, both pins.

Exit 0 only at 100% parity on the gap-closing functions.
"""
import json
import os
import random
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
os.environ["RUNE_TOKEN_OPTIMIZER_CORE"] = "python"  # oracle floor to start
PARITY_HOME = "/tmp/to_parity"

sys.path.insert(0, str(REPO / "src"))

from rune import to_loader  # noqa: E402
from rune import token_optimizer as to  # noqa: E402

if not to_loader.available():
    sys.exit("FATAL: Rust core not built — cargo build --release in src/rust/token_optimizer")

TRUTH = json.loads(
    (REPO / "src" / "rust" / "token_optimizer" / "frozen_truth.json").read_text())

fails: list[str] = []
counts = {"frozen": 0, "fuzz": 0, "live": 0}


def fail(layer, name, detail):
    fails.append(f"[{layer}] {name}: {detail}")
    print(f"FAIL[{layer}] {name}: {detail}"[:400])


def j(x):
    return json.dumps(x, sort_keys=True, ensure_ascii=True)


def check(layer, name, want, got):
    counts[layer] += 1
    if j(want) != j(got):
        fail(layer, name, f"want={j(want)[:300]} got={j(got)[:300]}")


def pin(mode):
    if mode == "python":
        os.environ["RUNE_TOKEN_OPTIMIZER_CORE"] = "python"
    else:
        os.environ.pop("RUNE_TOKEN_OPTIMIZER_CORE", None)


# ── layer 1: FROZEN (construction rows through the FFI) ────────────────────

for row in TRUTH:
    if row["fn"] != "to_get_optimizer_state":
        continue
    file_state, seed = row["args"]
    exp = row["expected"]
    env = to_loader.run_op(
        "get_optimizer", {},
        {"storage_text": None if file_state == "absent" else seed})
    got_state = env["value"]
    want_state = {
        "cache": exp["cache"],
        "stats": exp["stats"],
        "dirty": exp["dirty"],
        "changes_since_save": exp["changes_since_save"],
    }
    got_cmp = {k: got_state[k] for k in want_state}
    check("frozen", f"ffi_get_optimizer:{file_state}", want_state, got_cmp)

# pin contracts, both directions (probes live in the facade seam).
pin("python")
check("frozen", "pin_contract:to_core_python_floor", None, to._to_core())
pin("rust")
check("frozen", "pin_contract:to_core_reachable", True, to._to_core() is not None)
pin("python")

# ── the patched clock (both pins see identical time) ───────────────────────

FROZEN_NOW = datetime(2026, 10, 14, 6, 0, 0, tzinfo=timezone.utc)


class _FrozenDatetime(datetime):
    @classmethod
    def now(cls, tz=None):
        return FROZEN_NOW if tz is not None else FROZEN_NOW.replace(tzinfo=None)


class _Clock:
    def __init__(self):
        self.t = 1000.0

    def monotonic(self):
        return self.t

    def advance(self, dt):
        self.t += dt


def reset_scratch():
    if os.path.isdir(PARITY_HOME):
        shutil.rmtree(PARITY_HOME)
    os.makedirs(PARITY_HOME, exist_ok=True)


STORAGE = os.path.join(PARITY_HOME, "token_optimizer.json")
PERSONA = "P" * 80  # deterministic long persona (bypasses persona_definition)

TEXTS = [
    "fix the login bug",
    "kubernetes pods restart on failure",
    "  spaced   out   text  with  runs  ",
    "multi\nline\ncontext\nwith\nnewlines",
    "",
    "unicode ✓ snowman ☃ and dashes — oh my",
]
CONTEXTS = ["", "recent work on the parser; recent work on the parser;",
            "noise noise noise noise noise noise noise noise noise noise",
            "one two three four five six seven eight nine ten"]


def drive(script, seed_file=None):
    """Run one op script through the facade; return the observation tape."""
    reset_scratch()
    if seed_file is not None:
        with open(STORAGE, "w") as f:
            f.write(seed_file)
    clock = _Clock()
    to.time = clock  # the facade's `import time` binding — module-local
    to.datetime = _FrozenDatetime
    to._optimizer = None
    opt = to.TokenOptimizer(storage_path=STORAGE)
    opt._load_persona_definition = lambda: PERSONA
    tape = []

    def state_obs():
        return {
            "cache": [e.to_dict() for e in opt.cache.values()],
            "stats": dict(opt.stats),
            "dirty": opt._dirty,
            "changes": opt._changes_since_save,
        }

    def file_obs():
        if not os.path.exists(STORAGE):
            return None
        with open(STORAGE) as f:
            return f.read()

    for kind, arg in script:
        clock.advance(0.5)
        if kind == "dt":
            clock.advance(arg)  # jump the clock (debounce boundary probes)
        elif kind == "get_static_cache":
            tape.append({"r": opt.get_static_cache(), "s": state_obs()})
        elif kind == "get_dynamic_cache":
            tape.append({"r": opt.get_dynamic_cache(arg), "s": state_obs()})
        elif kind == "optimize_prompt":
            tape.append({"r": opt.optimize_prompt(arg[0], arg[1]), "s": state_obs()})
        elif kind == "get_status":
            tape.append({"r": opt.get_status(), "s": state_obs()})
        elif kind == "get_cache_stats":
            tape.append({"r": opt.get_cache_stats(), "s": state_obs()})
        elif kind == "save_if_dirty":
            opt._save_if_dirty()
            tape.append({"r": None, "s": state_obs(), "f": file_obs()})
        elif kind == "flush":
            opt.flush()
            tape.append({"r": None, "s": state_obs(), "f": file_obs()})
        elif kind == "mark_dirty":
            opt._mark_dirty()
            tape.append({"r": None, "s": state_obs()})
    final = {"tape": tape, "file": file_obs(), "status": opt.get_status()}
    to.time, to.datetime = __import__("time"), datetime
    return final


# ── layer 2: FUZZ (dual-mode full state machine) ──────────────────────────

rng = random.Random(2479)
SEEDS = [None,
         '{"cache": {"static": {"key": "static", "content": "seed persona",'
         ' "tokens": 3, "last_used": "2026-10-14T00:00:00+00:00",'
         ' "use_count": 9, "hit_count": 4}}, "stats": {"total_calls": 2,'
         ' "total_tokens_saved": 10, "cache_hits": 1, "cache_misses": 1}}']
for trial in range(40):
    script = []
    for _ in range(rng.randrange(4, 9)):
        roll = rng.random()
        if roll < 0.10:
            script.append(("dt", rng.choice([0.0, 10.0, 31.0, 65.0])))
        elif roll < 0.30:
            script.append(("get_static_cache", None))
        elif roll < 0.50:
            script.append(("get_dynamic_cache", rng.choice(CONTEXTS)))
        elif roll < 0.70:
            script.append(("optimize_prompt", (rng.choice(TEXTS), rng.choice(CONTEXTS))))
        elif roll < 0.78:
            script.append(("get_status", None))
        elif roll < 0.86:
            script.append(("get_cache_stats", None))
        elif roll < 0.92:
            script.append(("save_if_dirty", None))
        elif roll < 0.97:
            script.append(("flush", None))
        else:
            script.append(("mark_dirty", None))
    seed_file = rng.choice(SEEDS)

    pin("python")
    want = drive(script, seed_file)
    pin("rust")
    got = drive(script, seed_file)
    check("fuzz", f"trial{trial}", want, got)
    if fails:
        break
pin("python")

# ── layer 3: LIVE (dual-mode E2E on the gap fns) ───────────────────────────


def dual(label, call):
    pin("python")
    want = call()
    pin("rust")
    got = call()
    check("live", label, want, got)


# module-level API through the REAL global singleton (get_optimizer path).
def module_api():
    reset_scratch()
    to._optimizer = None
    clock = _Clock()
    to.time, to.datetime = clock, _FrozenDatetime
    os.environ["RUNE_MEMORY_DIR"] = PARITY_HOME
    try:
        r1 = to.optimize_prompt("hello world", "some context here")
        opt = to.get_optimizer()
        r2 = {"status": opt.get_status(), "stats": opt.get_cache_stats(),
              "singleton": to.get_optimizer() is opt}
        return {"r1": r1, "r2": r2}
    finally:
        to.time, to.datetime = __import__("time"), datetime
        to._optimizer = None


dual("module_api:get_optimizer_singleton", module_api)

# the 8 construction variants through the REAL get_optimizer (both pins).
CONS = [r for r in TRUTH if r["fn"] == "to_get_optimizer_state"]
for row in CONS:
    file_state, seed = row["args"]

    def construct(file_state=file_state, seed=seed):
        reset_scratch()
        if file_state != "absent":
            with open(STORAGE, "w") as f:
                f.write(seed)
        to._optimizer = None
        os.environ["RUNE_MEMORY_DIR"] = PARITY_HOME
        try:
            opt = to.get_optimizer()
            second = to.get_optimizer()
            out = {
                "cache": [e.to_dict() for e in opt.cache.values()],
                "stats": dict(opt.stats),
                "dirty": opt._dirty,
                "changes": opt._changes_since_save,
                "singleton_reused": second is opt,
            }
            return out
        finally:
            to._optimizer = None

    dual(f"construct:{file_state}", construct)

# _write_snapshot conversion through the REAL facade method (both pins).
SNAP = [r for r in TRUTH if r["fn"] == "to_write_snapshot_doc"][0]["args"][0]


def snap_write():
    reset_scratch()
    path = os.path.join(PARITY_HOME, "snap.json")
    opt = to.TokenOptimizer.__new__(to.TokenOptimizer)
    opt.storage_path = path
    to.TokenOptimizer._write_snapshot(opt, SNAP)
    with open(path) as f:
        return json.load(f)


dual("_write_snapshot:doc", snap_write)

print(f"parity: {sum(counts.values())} checks "
      f"(frozen {counts['frozen']}, fuzz {counts['fuzz']}, live {counts['live']}) "
      f"— {len(fails)} failures")
sys.exit(1 if fails else 0)