arborist/tests/test_canonical_cache.py

"""Tests for #000027 — canonical projection persistence.

Covers the acceptance criteria from the ticket §7:

1. First call writes one providence_cache row + one
   ``providence_canonical`` audit event.
2. Second call: cache hit; ``hit_count`` increments; kernel not
   re-run.
3. Audit-chain integrity intact after mixed canonical/RAG writes.
4. Bumping pi_star_ref (``@v1`` → ``@v2``) routes new questions to
   a fresh row; old rows remain in DB but unreachable via the live
   cache_key.
5. Bumping CHUNKING_VERSION does NOT stale canonical rows (they
   pin ``"n/a-canonical"``).
6. Distinct pi_star_refs namespace separately.
7. Strict vs equivalence_class dedup behavior for canonical rows.
"""

from __future__ import annotations

import json
from pathlib import Path

from arborist.qa.canonical_cache import (
    CANONICAL_AUDIT_EVENT_TYPE,
    canonical_cache_key,
    canonical_synthetic_source_root,
    lookup_canonical,
    lookup_or_persist,
    persist_canonical,
)
from arborist.qa.client import StubClient
from arborist.qa.query import query
from arborist.store import connect


# ----- low-level: cache_key shape ------------------------------------------


def test_canonical_cache_key_is_deterministic():
    a = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
    )
    b = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
    )
    assert a == b


def test_canonical_cache_key_distinct_pi_star_refs():
    """Same question + different pi_star_ref → different cache_key.
    The synthetic source_root encodes pi_star_ref."""
    a = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
    )
    b = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v2", policy={}
    )
    assert a != b


def test_canonical_cache_key_distinct_questions():
    a = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
    )
    b = canonical_cache_key(
        question="1 + 1", pi_star_ref="arithmetic@v1", policy={}
    )
    assert a != b


def test_canonical_cache_key_strict_vs_equivalence_class():
    """Strict mode keeps "0.1+0.2" and "0.1 + 0.2" distinct;
    equivalence_class mode collapses them via the trailing-strip /
    article-strip / case rules in question_hash."""
    strict_a = canonical_cache_key(
        question="0.1 + 0.2?", pi_star_ref="arithmetic@v1",
        policy={}, mode="strict",
    )
    strict_b = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1",
        policy={}, mode="strict",
    )
    eq_a = canonical_cache_key(
        question="0.1 + 0.2?", pi_star_ref="arithmetic@v1",
        policy={}, mode="equivalence_class",
    )
    eq_b = canonical_cache_key(
        question="0.1 + 0.2", pi_star_ref="arithmetic@v1",
        policy={}, mode="equivalence_class",
    )
    assert strict_a != strict_b   # "?" matters in strict mode
    assert eq_a == eq_b           # collapses in equivalence_class


# ----- persistence round-trip ---------------------------------------------


def test_persist_then_lookup(tmp_path: Path):
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        ckey = canonical_cache_key(
            question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
        )
        event_hash, run_dag_root, run_dag = persist_canonical(
            conn,
            cache_key_value=ckey,
            question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_input_bytes=b"0.1 + 0.2",
            canonical_output_bytes=b"3/10",
            policy={},
        )
        assert len(event_hash) == 64
        assert len(run_dag_root) == 64
        assert run_dag["nodes"][1]["stage"] == "canonical_projection"

        row = lookup_canonical(conn, ckey)
        assert row is not None
        assert row["audit_mode"] == "CANONICAL_PROJECTION"
        assert row["verifier_method"] == "canonical_projection"
        assert row["answer_text"] == "3/10"
        assert row["audit_event_hash"] == event_hash
        assert row["run_dag_root"] == run_dag_root
    finally:
        conn.close()


def test_audit_event_appended(tmp_path: Path):
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        before = conn.execute(
            "SELECT COUNT(*) FROM audit_events"
        ).fetchone()[0]
        ckey = canonical_cache_key(
            question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
        )
        persist_canonical(
            conn, cache_key_value=ckey, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_input_bytes=b"0.1 + 0.2",
            canonical_output_bytes=b"3/10", policy={},
        )
        after = conn.execute(
            "SELECT COUNT(*) FROM audit_events"
        ).fetchone()[0]
        assert after == before + 1
        latest = conn.execute(
            "SELECT event_type, body FROM audit_events "
            "ORDER BY seq DESC LIMIT 1"
        ).fetchone()
        assert latest["event_type"] == CANONICAL_AUDIT_EVENT_TYPE
        body = json.loads(latest["body"])
        assert body["pi_star_ref"] == "arithmetic@v1"
        assert body["canonical_output_text"] == "3/10"
        assert body["kernel_audit_mode"] == "CANONICAL_PROJECTION"
    finally:
        conn.close()


def test_lookup_or_persist_first_call_misses_second_hits(tmp_path: Path):
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        ckey1, row1, was_hit1, dag1 = lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_output_bytes=b"3/10",
            policy={},
        )
        assert was_hit1 is False
        assert row1 is None
        assert dag1["root"]

        ckey2, row2, was_hit2, dag2 = lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_output_bytes=b"3/10",
            policy={},
        )
        assert ckey1 == ckey2
        assert was_hit2 is True
        assert row2 is not None
        assert row2["answer_text"] == "3/10"
    finally:
        conn.close()


# ----- end-to-end via query() ---------------------------------------------


def test_query_writes_then_hits(tmp_path: Path):
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"

    first = query(
        question="0.1 + 0.2",
        qa_db=qa_db, chat_client=StubClient(""),
        model_id="m", shards_dir=no_shards,
    )
    assert first["status"] == "cache_miss_then_written"
    assert first["lookup_path"] == "canonical_cache_miss"
    ckey1 = first["cache_key"]
    assert ckey1 is not None

    second = query(
        question="0.1 + 0.2",
        qa_db=qa_db, chat_client=StubClient(""),
        model_id="m", shards_dir=no_shards,
    )
    assert second["status"] == "cache_hit"
    assert second["lookup_path"] == "canonical_cache_hit"
    assert second["cache_key"] == ckey1
    assert second["answer_text"] == "3/10"


def test_hit_count_increments_on_repeated_query(tmp_path: Path):
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    for _ in range(3):
        query(
            question="0.1 + 0.2",
            qa_db=qa_db, chat_client=StubClient(""),
            model_id="m", shards_dir=no_shards,
        )
    conn = connect(qa_db)
    try:
        row = conn.execute(
            "SELECT hit_count FROM providence_cache "
            "WHERE audit_mode = 'CANONICAL_PROJECTION'"
        ).fetchone()
        # 1 write + 2 hits; hit_count counts hits only.
        assert row["hit_count"] == 2
    finally:
        conn.close()


def test_audit_chain_intact_after_canonical_writes(tmp_path: Path):
    """audit_events chain (event_hash = sha256(prev || canonical_body))
    stays unbroken when canonical rows are interleaved."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    for q in ("0.1 + 0.2", "1 + 1", "A AND B", "0.1 + 0.2"):
        query(
            question=q, qa_db=qa_db,
            chat_client=StubClient(""),
            model_id="m", shards_dir=no_shards,
        )
    conn = connect(qa_db)
    try:
        rows = conn.execute(
            "SELECT seq, event_hash, prev_event_hash, body "
            "FROM audit_events ORDER BY seq"
        ).fetchall()
        import hashlib
        prev = None
        for r in rows:
            h = hashlib.sha256()
            if r["prev_event_hash"]:
                h.update(bytes.fromhex(r["prev_event_hash"]))
            h.update(r["body"].encode("utf-8"))
            assert h.hexdigest() == r["event_hash"]
            assert r["prev_event_hash"] == prev
            prev = r["event_hash"]
    finally:
        conn.close()


# ----- pi_star_ref version semantics --------------------------------------


def test_distinct_pi_star_refs_namespace_separately(tmp_path: Path):
    """Same question through two different kernels → two distinct rows."""
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        ckey_a, _, miss_a, _ = lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_output_bytes=b"3/10", policy={},
        )
        ckey_b, _, miss_b, _ = lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v2",  # hypothetical bump
            canonical_output_bytes=b"3/10", policy={},
        )
        assert miss_a is False
        assert miss_b is False
        assert ckey_a != ckey_b
        rows = conn.execute(
            "SELECT cache_key FROM providence_cache "
            "WHERE audit_mode = 'CANONICAL_PROJECTION'"
        ).fetchall()
        assert len({r["cache_key"] for r in rows}) == 2
    finally:
        conn.close()


def test_pi_star_version_bump_orphans_old_row(tmp_path: Path):
    """Write under @v1; lookup under @v2 (different synthetic
    source_root) misses; old @v1 row remains in DB."""
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        # Write under @v1.
        lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_output_bytes=b"3/10", policy={},
        )
        # Lookup under @v2 — different cache_key → miss.
        ckey_v2 = canonical_cache_key(
            question="0.1 + 0.2", pi_star_ref="arithmetic@v2",
            policy={},
        )
        assert lookup_canonical(conn, ckey_v2) is None
        # Old @v1 row still present; just unreachable via @v2 lookup.
        ckey_v1 = canonical_cache_key(
            question="0.1 + 0.2", pi_star_ref="arithmetic@v1",
            policy={},
        )
        assert lookup_canonical(conn, ckey_v1) is not None
    finally:
        conn.close()


# ----- governance + chunking_version ---------------------------------------


def test_chunking_version_bump_does_not_stale_canonical(tmp_path: Path):
    """Canonical rows pin chunking_version='n/a-canonical' so a
    chunker bump on the wikipedia path doesn't mass-stale math
    answers. Verify by inspecting the persisted column directly."""
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        ckey = canonical_cache_key(
            question="0.1 + 0.2", pi_star_ref="arithmetic@v1", policy={}
        )
        persist_canonical(
            conn, cache_key_value=ckey, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_input_bytes=b"0.1 + 0.2",
            canonical_output_bytes=b"3/10", policy={},
        )
        row = lookup_canonical(conn, ckey)
        assert row is not None
        # The persisted chunking_version is the canonical sentinel,
        # NOT the live CHUNKING_VERSION constant.
        cv = conn.execute(
            "SELECT chunking_version FROM providence_cache "
            "WHERE cache_key = ?", (ckey,),
        ).fetchone()[0]
        assert cv == "n/a-canonical"
    finally:
        conn.close()


def test_governance_policy_change_invalidates_lookup(tmp_path: Path):
    """Different policy → different governance_policy_hash → different
    cache_key. The old row stays but the lookup misses."""
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        # Write under empty policy.
        lookup_or_persist(
            conn, question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            canonical_output_bytes=b"3/10",
            policy={},
        )
        # Lookup under different policy — different ghash → miss.
        ckey_alt = canonical_cache_key(
            question="0.1 + 0.2",
            pi_star_ref="arithmetic@v1",
            policy={"answer_mode": "claim_lattice"},
        )
        assert lookup_canonical(conn, ckey_alt) is None
    finally:
        conn.close()


# ----- arborist canon stays transient -------------------------------------


def test_arborist_canon_does_not_persist(tmp_path: Path):
    """`arborist canon <key> "<input>"` is a one-shot probe; it
    bypasses query() entirely and writes nothing to providence_cache.
    Confirms the boundary the ticket §2.6 promises."""
    import subprocess
    import sys

    qa_db = tmp_path / "qa.db"
    # Make sure no rows exist before the canon call.
    conn = connect(qa_db)
    try:
        before = conn.execute(
            "SELECT COUNT(*) FROM providence_cache"
        ).fetchone()[0]
    finally:
        conn.close()

    r = subprocess.run(
        [sys.executable, "-m", "arborist.cli", "canon",
         "arithmetic@v1", "0.1 + 0.2"],
        capture_output=True, text=True, check=False,
    )
    assert r.returncode == 0
    assert r.stdout.strip() == "3/10"

    # No row should have been written by `canon` (it's transient by
    # design — the canon CLI doesn't take a qa_db). Confirm the
    # qa_db count is unchanged. We verify by re-checking the same
    # tmp_path qa_db (which is unrelated to the default qa.db that
    # `canon` doesn't write to anyway).
    conn = connect(qa_db)
    try:
        after = conn.execute(
            "SELECT COUNT(*) FROM providence_cache"
        ).fetchone()[0]
    finally:
        conn.close()
    assert after == before


# ----- synthetic dimension shape ------------------------------------------


def test_synthetic_source_root_encodes_pi_star_ref():
    a = canonical_synthetic_source_root("arithmetic@v1")
    b = canonical_synthetic_source_root("arithmetic@v2")
    c = canonical_synthetic_source_root("logic-kernel@v1")
    assert len({a, b, c}) == 3
    assert len(a) == 64  # sha256 hex