arborist/bench/batteries/b_5r.py

"""5R battery — React / Rearrange / Restore / Replicate / Resonate.

Workspace-operator battery. Per SQD whitepaper §9.3, the 5R battery
tests operators applied to a workspace state. With ticket #000014
(SelfModel) and #000017 (memory-root) landed, the workspace is the
``(selfmodel_root, memory_root, audit_events_high_water)`` triple
plus the cited claim lattice / providence cache.

Five sub-batteries:

- **React** — incorporate new facts/constraints into the workspace.
  Test: a new observation between snapshots produces the expected
  delta in the branch summaries.
- **Rearrange** — restructure representation without changing
  semantic content. Test: re-ordering claim-lattice claims produces
  the same root (selfmodel_root invariant under claim order).
- **Restore** — retrieve a prior fact and re-assert it cleanly.
  Test: re-canonicalizing an earlier fact produces the same
  canonical bytes; audit chain still verifies.
- **Replicate** — produce independent canonical encodings via the
  π* registry. Test: same input through the named π* yields
  byte-equal output across N reruns.
- **Resonate** — detect low variance across encodings/reruns. Test:
  variance over N reruns is below threshold; deterministic π*'s
  should be zero variance.

All runners deterministic, no LLM-as-judge. Phase 1a fixtures embed
the workspace state directly; Phase 1b.2 will read live
selfmodel_records / memory_records from a shard.
"""

from __future__ import annotations

import statistics
from pathlib import Path

from bench.batteries.base import (
    BatteryResult,
    TaskResult,
    fixture_digest,
    fixture_meta,
    iter_tasks,
    validate_carrier,
)
from bench.batteries.b_5s import _runtime_digest


def _carrier_check(task: dict) -> tuple[bool, str]:
    reason = validate_carrier(task)
    if reason is not None:
        return False, reason
    return True, ""


def _build_result(
    sub_battery: str,
    fixtures_path: Path,
    per_task: list[TaskResult],
    metrics: dict[str, float],
) -> BatteryResult:
    pass_count = sum(1 for t in per_task if t.passed)
    fail_count = len(per_task) - pass_count
    return BatteryResult(
        battery="5r",
        sub_battery=sub_battery,
        fixture_path=str(fixtures_path),
        fixture_digest=fixture_digest(fixtures_path),
        pass_count=pass_count,
        fail_count=fail_count,
        metrics=metrics,
        per_task=per_task,
        runtime_digest=_runtime_digest(),
    )


# ---------------------------------------------------------------------
# React — incorporate new fact/constraint
# ---------------------------------------------------------------------


def _live_workspace_apply(facts: list[str], event_type: str = "observation") -> list[str]:
    """Apply each fact as an ``observation`` audit event on a fresh
    temp shard; return the bodies of the resulting audit_events as
    a list of JSON strings the runner can substring-match against.

    Used by React and Restore live modes — both treat the workspace
    as the audit chain itself.
    """
    import os
    import tempfile

    from arborist.store import append_audit, connect, transaction

    tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
    tmp.close()
    try:
        conn = connect(tmp.name)
        try:
            with transaction(conn):
                for f in facts:
                    append_audit(
                        conn, event_type=event_type,
                        subject_root=None, body={"fact": f},
                    )
            rows = conn.execute(
                "SELECT body FROM audit_events ORDER BY seq"
            ).fetchall()
            return [row["body"] for row in rows]
        finally:
            conn.close()
    finally:
        try:
            os.unlink(tmp.name)
        except OSError:
            pass


def run_react(fixtures_path: Path) -> BatteryResult:
    """Each task: workspace state at t, an observation at t+1, and an
    expected_delta describing how the workspace should change.

    Two fixture modes:

    - **Embedded** (Phase 1a): runner diffs ``snapshot_t1.facts``
      against ``expected_delta`` lists. Pure list-set check, no
      arborist surface call.
    - **Live** (Phase 1b.2): each fact in ``snapshot_t1.facts`` is
      applied as a real ``observation`` audit event on a temp
      shard; runner queries the resulting audit_events bodies and
      asserts ``expected_delta.added_facts`` substrings are
      present, ``expected_delta.removed_facts`` substrings are NOT.
      Tests that the audit chain actually integrates observations.
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            delta = task["expected_delta"]
            if task.get("live"):
                source = "live"
                t1_facts = task["snapshot_t1"].get("facts", [])
                bodies = _live_workspace_apply(t1_facts)
                bodies_text = "\n".join(bodies)
                added_ok = all(
                    f in bodies_text for f in delta.get("added_facts", [])
                )
                removed_ok = all(
                    f not in bodies_text
                    for f in delta.get("removed_facts", [])
                )
            else:
                source = "embedded"
                t1_set = set(task["snapshot_t1"].get("facts", []))
                added_ok = all(f in t1_set for f in delta.get("added_facts", []))
                removed_ok = all(
                    f not in t1_set for f in delta.get("removed_facts", [])
                )
            passed_all = added_ok and removed_ok
            expected = task.get("expected", "pass")
            observed = "pass" if passed_all else "fail"
            passed = observed == expected
            detail = {
                "source": source,
                "added_ok": added_ok,
                "removed_ok": removed_ok,
                "expected": expected,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "react"),
        fixtures_path,
        per_task,
        {"react_integration_rate": rate},
    )


# ---------------------------------------------------------------------
# Rearrange — restructure without semantic shift
# ---------------------------------------------------------------------


def run_rearrange(fixtures_path: Path) -> BatteryResult:
    """Each task: two ordered fact lists (alt_a, alt_b) that should
    canonicalize to the same bytes through the named π* (or hashing
    rule). Pass = canonical encodings are byte-equal.

    This exercises the order-invariance contracts in SelfModel
    (capability_claim_hashes are sorted) and Memory (branches are
    sorted by branch_id) without depending on shard state.
    """
    from arborist.pi_star import get

    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            pi_star = get(task["pi_star_ref"])
            alt_a = task["alt_a"].encode("utf-8")
            alt_b = task["alt_b"].encode("utf-8")
            canon_a = pi_star.canonicalize(alt_a)
            canon_b = pi_star.canonicalize(alt_b)
            equivalent = canon_a == canon_b
            expected = task.get("expected_equivalent", True)
            passed = equivalent == bool(expected)
            detail = {
                "canon_match": equivalent,
                "expected_equivalent": expected,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "rearrange"),
        fixtures_path,
        per_task,
        {"rearrange_invariance_rate": rate},
    )


# ---------------------------------------------------------------------
# Restore — retrieve prior fact and re-assert
# ---------------------------------------------------------------------


def run_restore(fixtures_path: Path) -> BatteryResult:
    """Each task: a prior_fact + history + current_facts. Pass =
    the prior_fact is retrievable.

    Two fixture modes:

    - **Embedded** (Phase 1a): list-membership check across history
      snapshots and current_facts. No arborist call.
    - **Live** (Phase 1b.2): all historical + current facts are
      applied as ``observation`` audit_events on a temp shard;
      retrievability = the prior_fact appears in any audit_event
      body. Tests that the audit chain actually preserves prior
      observations.
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            prior_fact = task["prior_fact"]
            history = task.get("history", [])
            current = task.get("current_facts", [])
            if task.get("live"):
                source = "live"
                all_facts: list[str] = []
                for s in history:
                    all_facts.extend(s.get("facts", []))
                all_facts.extend(current)
                bodies = _live_workspace_apply(all_facts)
                bodies_text = "\n".join(bodies)
                retrievable = prior_fact in bodies_text
                in_current = prior_fact in current
                in_history = any(
                    prior_fact in s.get("facts", []) for s in history
                )
            else:
                source = "embedded"
                in_current = prior_fact in current
                in_history = any(
                    prior_fact in s.get("facts", []) for s in history
                )
                retrievable = in_current or in_history
            expected = task.get("expected", "pass")
            observed = "pass" if retrievable else "fail"
            passed = observed == expected
            detail = {
                "source": source,
                "in_current": in_current,
                "in_history": in_history,
                "retrievable": retrievable,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "restore"),
        fixtures_path,
        per_task,
        {"restore_retrievability_rate": rate},
    )


# ---------------------------------------------------------------------
# Replicate — independent canonical encodings
# ---------------------------------------------------------------------


def run_replicate(fixtures_path: Path) -> BatteryResult:
    """Each task: an input + π* + expected_replicas count. Run π*
    canonicalize N times; pass = all N outputs are byte-equal.

    Tests determinism contract on registered π*'s: same input + same
    π* should always yield the same canonical bytes regardless of how
    many times we ask.
    """
    from arborist.pi_star import get

    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            pi_star = get(task["pi_star_ref"])
            n = int(task.get("replicas", 5))
            input_bytes = task["input"].encode("utf-8")
            outputs = [pi_star.canonicalize(input_bytes) for _ in range(n)]
            all_equal = all(o == outputs[0] for o in outputs)
            expected = task.get("expected", "pass")
            observed = "pass" if all_equal else "fail"
            passed = observed == expected
            detail = {
                "replicas": n,
                "all_equal": all_equal,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "replicate"),
        fixtures_path,
        per_task,
        {"replicate_determinism_rate": rate},
    )


# ---------------------------------------------------------------------
# Resonate — variance across encodings/reruns
# ---------------------------------------------------------------------


def run_resonate(fixtures_path: Path) -> BatteryResult:
    """Each task: an input + π* + max_variance threshold. Run π*
    canonicalize N times; compute the per-byte variance (over the
    distinct hashes of each output). Deterministic π*'s yield
    variance == 0; non-deterministic ones violate the contract.

    The metric is ``distinct_output_count``: how many unique outputs
    in N runs. Pass = distinct_count <= ``expected_max_distinct``.
    """
    from arborist.pi_star import get

    per_task: list[TaskResult] = []
    distinct_counts: list[int] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            pi_star = get(task["pi_star_ref"])
            n = int(task.get("runs", 10))
            input_bytes = task["input"].encode("utf-8")
            outputs = {pi_star.canonicalize(input_bytes) for _ in range(n)}
            distinct = len(outputs)
            distinct_counts.append(distinct)
            expected_max = int(task.get("expected_max_distinct", 1))
            stable = distinct <= expected_max
            expected = task.get("expected", "pass")
            observed = "pass" if stable else "fail"
            passed = observed == expected
            detail = {
                "runs": n,
                "distinct_outputs": distinct,
                "expected_max_distinct": expected_max,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    mean_distinct = statistics.mean(distinct_counts) if distinct_counts else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "resonate"),
        fixtures_path,
        per_task,
        {
            "resonate_stability_rate": rate,
            "mean_distinct_outputs": mean_distinct,
        },
    )


SUB_BATTERIES = {
    "react": run_react,
    "rearrange": run_rearrange,
    "restore": run_restore,
    "replicate": run_replicate,
    "resonate": run_resonate,
}