rune/scripts/check_memory_selector_parity.py

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

  1. FROZEN — the gap truth corpus (src/rust/memory_selector/frozen_truth.json,
     generated by scripts/gen_memory_selector_truth.py under the
     scripts/ms_freeze_seams rig against the LIVE oracle with
     RUNE_MSE_CORE=python pinned): the kernel-shaped rows (ms_add_source)
     replayed through the ctypes FFI, case by case; the pin-contract row
     asserted in both directions.
  2. FUZZ — 200 seeded random registry ops (seed 2479): the oracle's dict
     semantics vs the FFI kernel, exact compare.
  3. LIVE — dual-mode facade battery: the same facade calls (the frozen
     _source_store E2E rows + registry order via the real add_source + one
     full select_for_task) executed with RUNE_MSE_CORE=python (oracle floor)
     and with the Rust core live; results must be identical.

Exit 0 only at 100% parity on the gap-closing functions.
"""
import copy
import functools
import json
import os
import random
import sys
from pathlib import Path

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

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

from rune import memory_store as store  # noqa: E402
from rune import memory_selector as ms  # noqa: E402
from rune import mse_loader  # noqa: E402

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

TRUTH = json.loads(
    (REPO / "src" / "rust" / "memory_selector" / "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):
    """Flip the facade pin at runtime (mse_loader.core reads env per call)."""
    if mode == "python":
        os.environ["RUNE_MSE_CORE"] = "python"
    else:
        os.environ.pop("RUNE_MSE_CORE", None)


def seed_store():
    """The rig's seed (scripts/ms_freeze_seams/sitecustomize.py), verbatim
    rows, at the parity scratch — both pins see byte-identical stores."""
    import shutil

    if os.path.isdir(PARITY_HOME):
        shutil.rmtree(PARITY_HOME)
    os.makedirs(os.path.join(PARITY_HOME, "store"), exist_ok=True)
    store.DB_PATH = os.path.join(PARITY_HOME, "store", "memory.db")
    store.STORE_DIR = os.path.join(PARITY_HOME, "store")
    store._embed = lambda texts: {}  # lexical-only, deterministic
    pinned = os.path.join(PARITY_HOME, "pinned-store.db")
    db = store.connect(pinned)
    try:
        store.add_entry(db, "Kubernetes pods restart automatically on failure",
                        importance=0.5, source="truth")
        store.add_entry(db, "Kubernetes pods restart policy for the cluster",
                        importance=0.5, source="truth")
        store.add_entry(db, "Voice identity and speaking style stay consistent",
                        importance=0.5, source="truth", memory_type="self")
        db.commit()
    finally:
        db.close()
    return pinned


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

for row in TRUTH:
    fn = row["fn"]
    if fn == "ms_add_source":
        names, name = row["args"]
        check("frozen", f"ffi_add_source:{name}",
              row["expected"], mse_loader.add_source(names, name))

pin_row = next(r for r in TRUTH if r["fn"] == "_mse_core")
pin("python")
check("frozen", "pin_contract:python_floor", None, ms._mse_core())
pin("rust")
check("frozen", "pin_contract:core_reachable", True, ms._mse_core() is not None)

# ── layer 2: FUZZ (registry semantics: oracle dict vs FFI kernel) ──────────

rng = random.Random(2479)
DEFAULTS = ["store", "neurons", "graphs", "abstraction", "consolidated", "persona"]
for trial in range(200):
    regs: dict[str, object] = {}
    names = DEFAULTS[: rng.randrange(0, 4)]
    for n in names:
        regs[n] = object()
    for _ in range(rng.randrange(1, 5)):
        name = rng.choice(DEFAULTS + ["zz", "custom", "x"])
        regs[name] = object()
        names = mse_loader.add_source(names, name)
        check("fuzz", f"trial{trial}:{name}", list(regs), names)
        if fails:
            break
    if fails:
        break

# ── layer 3: LIVE (dual-mode facade battery) ───────────────────────────────

PINNED = seed_store()


def fresh_selector():
    sel = ms.MemorySelector.__new__(ms.MemorySelector)
    sel.db_path = PINNED
    sel.min_score = 0.42
    sel._sources = {}
    return sel


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


# The frozen _source_store E2E rows, both pins.
for row in TRUTH:
    if row["fn"] != "MemorySelector._source_store":
        continue
    query, types, budget, min_score = row["args"]

    def run(query=query, types=types, budget=budget, min_score=min_score):
        sel = fresh_selector()
        out = ms.MemorySelector._source_store(sel, query, copy.deepcopy(types), budget, min_score)
        return out

    dual(f"_source_store:{query[:24]}:{types}", run)

# Registry order through the REAL add_source, both pins (append + replace).
for name in ("zz_custom", "store", "aa_first"):
    def reg(name=name):
        sel = fresh_selector()
        for n in DEFAULTS:
            ms.MemorySelector.add_source(sel, n, lambda *a: [])
        ms.MemorySelector.add_source(sel, name, lambda *a: [])
        return ms.MemorySelector.source_names(sel)

    dual(f"add_source:{name}", reg)

# Full select_for_task through the store source only (the other sources'
# stores are not seeded here), both pins.
def selection():
    sel = fresh_selector()
    sel._sources = {"store": functools.partial(ms._store_recall, db_path=PINNED)}
    out = ms.MemorySelector.select_for_task(sel, "kubernetes pods restart", budget=4)
    return {
        "blocks": [
            {a: getattr(b, a) for a in
             ("name", "priority", "content", "tokens", "source", "evictable")}
            for b in out.blocks
        ],
        "candidates_evaluated": out.candidates_evaluated,
        "selected_count": out.selected_count,
        "abstained": out.abstained,
        "types_queried": out.types_queried,
    }

dual("select_for_task:store_only", selection)

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)