arborist/tests/test_witness.py

"""Tests for the multi-modality witness (Ticket #000028).

Three modalities — kernel, cache, alien (LLM) — fan out in parallel
over a canonical-shape question. Every modality's answer is
canonicalized through the same kernel for byte-comparison. The
witness records cross-modality agreement as a label.

These tests exercise ``arborist.qa.witness`` directly with
:class:`StubClient`. Live-LLM coverage lives under the bench harness;
this file stays offline.
"""

from __future__ import annotations

import time

import pytest

from arborist.qa.client import StubClient
from arborist.qa.witness import (
    Witness,
    _canonicalize_via_kernel,
    _classify_agreement,
    run_witness,
)


# ----- _canonicalize_via_kernel ------------------------------------------


@pytest.mark.parametrize(
    "text,ref,expected",
    [
        ("3/10", "arithmetic@v1", b"3/10"),
        ("0.3", "arithmetic@v1", b"3/10"),
        ("0.1 + 0.2", "arithmetic@v1", b"3/10"),
        ("  3/10  ", "arithmetic@v1", b"3/10"),
        ("TRUE", "logic-kernel@v1", b"TRUE"),
    ],
)
def test_canonicalize_idempotent(text, ref, expected):
    assert _canonicalize_via_kernel(text, ref) == expected


@pytest.mark.parametrize(
    "text,ref",
    [
        ("the answer is 0.3", "arithmetic@v1"),
        ("approximately 0.3", "arithmetic@v1"),
        ("", "arithmetic@v1"),
        (None, "arithmetic@v1"),
    ],
)
def test_canonicalize_rejects_non_canonical(text, ref):
    assert _canonicalize_via_kernel(text, ref) is None


# ----- _classify_agreement ------------------------------------------------


def _mk_modality(name, ok, bytes_=None, error=None):
    from arborist.qa.witness import ModalityResult
    return ModalityResult(
        modality=name,
        raw_answer=bytes_.decode() if bytes_ else None,
        canonical_bytes=bytes_ if ok else None,
        error=error if not ok else None,
        elapsed_ms=0.0,
    )


def test_classify_strict_witnessed():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", True, b"3/10"),
        "llm": _mk_modality("llm", True, b"3/10"),
    }
    assert _classify_agreement(mods, b"3/10") == "STRICT-WITNESSED"


def test_classify_llm_diverged():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", True, b"3/10"),
        "llm": _mk_modality("llm", True, b"2/5"),
    }
    assert _classify_agreement(mods, b"3/10") == "LLM-DIVERGED"


def test_classify_cache_drift():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", True, b"4/10"),
        "llm": _mk_modality("llm", True, b"3/10"),
    }
    assert _classify_agreement(mods, b"3/10") == "CACHE-DRIFT"


def test_classify_kernel_llm_agree_no_cache():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", False, error="ABSENT"),
        "llm": _mk_modality("llm", True, b"3/10"),
    }
    assert _classify_agreement(mods, b"3/10") == "KERNEL-LLM-AGREE"


def test_classify_kernel_llm_diverged_no_cache():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", False, error="ABSENT"),
        "llm": _mk_modality("llm", True, b"5/10"),
    }
    assert _classify_agreement(mods, b"3/10") == "KERNEL-LLM-DIVERGED"


def test_classify_kernel_cache_agree_llm_absent():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", True, b"3/10"),
        "llm": _mk_modality("llm", False, error="PIS_REJECT"),
    }
    assert _classify_agreement(mods, b"3/10") == "KERNEL-CACHE-AGREE"


def test_classify_kernel_only():
    mods = {
        "kernel": _mk_modality("kernel", True, b"3/10"),
        "cache": _mk_modality("cache", False, error="ABSENT"),
        "llm": _mk_modality("llm", False, error="ABSENT"),
    }
    assert _classify_agreement(mods, b"3/10") == "KERNEL-ONLY"


# ----- run_witness end-to-end --------------------------------------------


def test_witness_kernel_llm_agree():
    """LLM emits the same canonical bytes — no cache."""
    client = StubClient(answer="3/10")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=client,
        model_id="stub",
    )
    assert isinstance(w, Witness)
    assert w.agreement_label == "KERNEL-LLM-AGREE"
    assert w.modalities["kernel"].ok is True
    assert w.modalities["cache"].error == "ABSENT"
    assert w.modalities["llm"].ok is True
    assert w.modalities["llm"].canonical_bytes == b"3/10"
    assert w.modalities["llm"].raw_answer == "3/10"


def test_witness_llm_lexically_different_but_canonically_same():
    """LLM says '0.3'; kernel canonicalizes to '3/10' — agreement."""
    client = StubClient(answer="0.3")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=client,
        model_id="stub",
    )
    assert w.agreement_label == "KERNEL-LLM-AGREE"
    assert w.modalities["llm"].canonical_bytes == b"3/10"


def test_witness_llm_diverged():
    """LLM emits a wrong canonical-shape answer — divergence recorded."""
    client = StubClient(answer="0.4")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=client,
        model_id="stub",
    )
    assert w.agreement_label == "KERNEL-LLM-DIVERGED"
    assert w.modalities["llm"].canonical_bytes == b"2/5"  # 0.4 → 2/5
    assert w.modalities["llm"].raw_answer == "0.4"


def test_witness_llm_not_canonical_text():
    """LLM emits prose; kernel rejects → LLM modality absent."""
    client = StubClient(answer="the answer is 0.3")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=client,
        model_id="stub",
    )
    # No cache, LLM rejected — kernel-only result.
    assert w.agreement_label == "KERNEL-ONLY"
    assert w.modalities["llm"].ok is False
    assert w.modalities["llm"].error == "PIS_REJECT"
    assert w.modalities["llm"].raw_answer == "the answer is 0.3"


def test_witness_cache_present_and_agrees():
    """Cache closure returns the canonical bytes — STRICT-WITNESSED."""
    client = StubClient(answer="3/10")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        cache_lookup=lambda: b"3/10",
        chat_client=client,
        model_id="stub",
    )
    assert w.agreement_label == "STRICT-WITNESSED"
    assert w.modalities["cache"].ok is True
    assert w.modalities["cache"].canonical_bytes == b"3/10"


def test_witness_cache_drift_caught():
    """Cache returns stale bytes; kernel + LLM agree → CACHE-DRIFT."""
    client = StubClient(answer="3/10")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        cache_lookup=lambda: b"4/10",  # someone bumped the kernel without invalidating
        chat_client=client,
        model_id="stub",
    )
    assert w.agreement_label == "CACHE-DRIFT"
    assert w.modalities["cache"].canonical_bytes == b"2/5"  # 4/10 → 2/5


def test_witness_no_chat_client_kernel_only():
    """No chat client provided — LLM modality reports ABSENT."""
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=None,
        model_id="",
    )
    assert w.agreement_label == "KERNEL-ONLY"
    assert w.modalities["llm"].error == "ABSENT"


def test_witness_llm_raises_recorded_as_error():
    """chat_completion raises — modality records LLM_ERROR, not crash."""

    class RaisingClient:
        def chat_completion(self, *a, **kw):
            raise RuntimeError("connection refused")

    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=RaisingClient(),
        model_id="stub",
    )
    assert w.agreement_label == "KERNEL-ONLY"
    assert w.modalities["llm"].error == "LLM_ERROR:RuntimeError"


def test_witness_to_dict_round_trip():
    """to_dict() emits a JSON-serializable shape."""
    import json

    client = StubClient(answer="3/10")
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        chat_client=client,
        model_id="stub",
    )
    d = w.to_dict()
    blob = json.dumps(d)
    assert "STRICT" in blob or "AGREE" in blob
    assert d["pi_star_ref"] == "arithmetic@v1"
    assert d["canonical_answer_text"] == "3/10"
    assert "kernel" in d["modalities"]
    assert "llm" in d["modalities"]


def test_witness_parallel_not_sequential():
    """Total wall time ≤ slowest modality + overhead, not sum."""

    class SlowClient:
        def chat_completion(self, *a, **kw):
            time.sleep(0.30)
            return "3/10"

    def slow_cache():
        time.sleep(0.30)
        return None

    t0 = time.monotonic()
    w = run_witness(
        question="0.1 + 0.2",
        pi_star_ref="arithmetic@v1",
        canonical_answer_bytes=b"3/10",
        cache_lookup=slow_cache,
        chat_client=SlowClient(),
        model_id="stub",
    )
    elapsed = time.monotonic() - t0
    # Sequential would be ~0.60s; parallel should be ~0.30s + small overhead.
    # Allow generous slack to avoid CI flake.
    assert elapsed < 0.55, f"witness ran sequentially? wall={elapsed:.3f}s"
    assert w.agreement_label == "KERNEL-LLM-AGREE"


# ----- query() integration ------------------------------------------------


def test_query_canonical_path_off_by_default(tmp_path):
    """Default policy: canonical_witness_enabled is False; no LLM call.
    Persistence is also default-on (#000027), so first call writes
    a row → status='cache_miss_then_written'."""
    from arborist.qa.query import query

    client = StubClient(answer="<should-not-be-called>")
    result = query(
        question="0.1 + 0.2",
        qa_db=tmp_path / "qa.db",
        chat_client=client,
        model_id="stub",
    )
    assert result["status"] == "cache_miss_then_written"
    assert result["audit_mode"] == "CANONICAL_PROJECTION"
    assert result.get("witness") is None
    # Critical: no LLM round-trip happened.
    assert client.calls == []


def test_query_canonical_with_witness_calls_llm(tmp_path):
    """policy['canonical_witness_enabled']=True → witness runs."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    result = query(
        question="0.1 + 0.2",
        qa_db=tmp_path / "qa.db",
        chat_client=client,
        model_id="stub",
        policy=policy,
    )
    assert result["status"] == "cache_miss_then_written"
    witness = result.get("witness")
    assert witness is not None
    # First call: no prior cache row to compare against → cache leg
    # is ABSENT; agreement is kernel↔LLM only.
    assert witness["agreement_label"] == "KERNEL-LLM-AGREE"
    # The LLM was actually called.
    assert len(client.calls) == 1


def test_query_canonical_witness_reaches_strict_after_persist(tmp_path):
    """Post-#000027 + cache-leg wire: a SECOND witness call (after the
    first writes a row) compares kernel + cache + LLM, all three
    byte-equal → STRICT-WITNESSED. This was structurally unreachable
    before #000027 landed — the cache leg always returned None."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    qa_db = tmp_path / "qa.db"

    # First call: persists row; cache leg ABSENT; KERNEL-LLM-AGREE.
    first = query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    assert first["status"] == "cache_miss_then_written"
    assert first["witness"]["agreement_label"] == "KERNEL-LLM-AGREE"

    # Second call: cache hits (no LLM run from query() path; witness
    # still calls LLM separately when enabled). Cache leg now
    # populated with the persisted bytes → STRICT-WITNESSED.
    second = query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    assert second["status"] == "cache_hit"
    assert second["witness"]["agreement_label"] == "STRICT-WITNESSED"


# ---------- #000028 follow-ups: capital ledger + sample-rate -------------


def test_query_canonical_witness_sample_rate_zero_skips(tmp_path):
    """sample_rate=0.0 → witness never fires (passive-calibration
    config); LLM is not called even when canonical_witness_enabled=True."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    policy["canonical_witness_sample_rate"] = 0.0
    qa_db = tmp_path / "qa.db"
    result = query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    assert result["status"] == "cache_miss_then_written"
    assert result.get("witness") is None  # skipped — sampled out
    assert client.calls == []  # critical: no LLM round-trip


def test_query_canonical_witness_sample_rate_one_always_fires(tmp_path):
    """sample_rate=1.0 (default when enabled) → witness fires every call."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    policy["canonical_witness_sample_rate"] = 1.0
    qa_db = tmp_path / "qa.db"
    result = query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    assert result.get("witness") is not None
    assert len(client.calls) == 1


def test_query_canonical_witness_capital_ledger_records_cost(tmp_path):
    """When witness fires + persistence writes a row, the witness LLM
    cost lands in capital_ledger as op_type='canonical_witness'.
    ForkScore can then compare witness-on vs witness-off forks
    honestly via the ledger summary."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query
    from arborist.store import connect

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    qa_db = tmp_path / "qa.db"
    query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    conn = connect(qa_db)
    try:
        rows = conn.execute(
            "SELECT op_type, estimator_inputs_blob FROM capital_ledger "
            "WHERE op_type = 'canonical_witness'"
        ).fetchall()
    finally:
        conn.close()
    assert len(rows) == 1
    import json as _json
    inputs = _json.loads(rows[0]["estimator_inputs_blob"])
    assert inputs["pi_star_ref"] == "arithmetic@v1"
    assert inputs["agreement_label"] in (
        "STRICT-WITNESSED", "KERNEL-LLM-AGREE",
        "CACHE-DRIFT", "LLM-DIVERGED", "KERNEL-LLM-DIVERGED",
    )


def test_query_canonical_witness_no_ledger_when_sampled_out(tmp_path):
    """Sampled-out witness must not record capital cost (no work happened)."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query
    from arborist.store import connect

    client = StubClient(answer="3/10")
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    policy["canonical_witness_sample_rate"] = 0.0
    qa_db = tmp_path / "qa.db"
    query(
        question="0.1 + 0.2", qa_db=qa_db,
        chat_client=client, model_id="stub", policy=policy,
    )
    conn = connect(qa_db)
    try:
        n = conn.execute(
            "SELECT COUNT(*) FROM capital_ledger "
            "WHERE op_type = 'canonical_witness'"
        ).fetchone()[0]
    finally:
        conn.close()
    assert n == 0