arborist/tests/test_qa.py

"""Q&A: 8-dim cache_key, falsification-aware lookup, audit chain.

Mock client only — no network. Verifies the v9.8 admissibility invariant
end-to-end.
"""

from __future__ import annotations

import hashlib
from typing import Iterator

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.merkle import proof_from_dict, verify_proof
from arborist.qa import ask
from arborist.qa.client import StubClient
from arborist.qa.keys import (
    cache_key,
    conversation_hash,
    governance_policy_hash,
    model_profile_hash,
    question_hash,
)
from arborist.source import Source
from arborist.store import connect


class FakeSource(Source):
    source_type = "test"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str) -> Document:
    return Document(uri=uri, content=content, source_type="test", title=uri)


LONG = (
    "The eight forms of capital include living, social, and intellectual. " * 10
    + "Merkle providence proves answer derives from a specific source. " * 10
)


def _ingest_one(conn) -> str:
    ingest_source(conn, FakeSource([_doc("test://qa", LONG)]))
    return conn.execute(
        "SELECT document_root FROM documents WHERE document_uri='test://qa'"
    ).fetchone()["document_root"]


def test_ask_writes_record_with_strict_proof(tmp_path):
    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
        # Verbatim quote from LONG → audit_mode=STRICT.
        client = StubClient(
            answer=(
                'The document states: '
                '"eight forms of capital include living, social, and intellectual"'
            )
        )
        result = ask(
            conn,
            document_root=root,
            question="What are the forms of capital?",
            client=client,
            model_id="test-model",
            revision="r1",
            quantization="fp8",
        )
        assert result["status"] == "cache_miss_then_written"
        assert result["audit_mode"] == "STRICT"
        assert result["n_quotes"] == 1
        assert result["n_verified"] == 1
        assert result["unverified_quotes"] == []

        # Stored proof reconstructs the source root.
        proof = proof_from_dict(result["merkle_proof"]["chunk_0_proof"])
        assert verify_proof(proof)
        assert proof.root.hex() == root

        # Providence cache row landed with correct fields.
        row = conn.execute("SELECT * FROM providence_cache").fetchone()
        assert row["source_root"] == root
        assert row["falsification_state"] == "live"
        assert row["chain"] == "private"
        assert row["audit_event_hash"] is not None

        # Audit chain has providence_write event.
        last = conn.execute(
            "SELECT event_type FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()["event_type"]
        assert last == "providence_write"
    finally:
        conn.close()


def test_ask_cache_hit_does_not_call_client(tmp_path):
    db = tmp_path / "hit.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
        client = StubClient(answer="first answer")
        ask(conn, document_root=root, question="Q1?", client=client, model_id="m")
        assert len(client.calls) == 1

        # Second ask with identical inputs -> hit, no client call.
        result = ask(
            conn, document_root=root, question="Q1?", client=client, model_id="m"
        )
        assert result["status"] == "cache_hit"
        assert result["answer_text"] == "first answer"
        assert len(client.calls) == 1  # unchanged

        row = conn.execute(
            "SELECT hit_count FROM providence_cache"
        ).fetchone()
        assert row["hit_count"] == 1
    finally:
        conn.close()


def test_different_model_yields_different_cache_key(tmp_path):
    db = tmp_path / "model.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
        client = StubClient(answer="ans")
        ask(conn, document_root=root, question="Q?", client=client, model_id="A")
        ask(conn, document_root=root, question="Q?", client=client, model_id="B")
        n = conn.execute("SELECT COUNT(*) FROM providence_cache").fetchone()[0]
        assert n == 2
        assert len(client.calls) == 2
    finally:
        conn.close()


def test_falsification_skips_cache_hit(tmp_path):
    """Stale records are ignored even when the 8-dim key matches."""
    db = tmp_path / "stale.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
        client = StubClient(answer="first")
        first = ask(
            conn, document_root=root, question="Q?", client=client, model_id="m"
        )
        ckey = first["cache_key"]
        # Mark stale.
        conn.execute(
            "UPDATE providence_cache SET falsification_state='stale' "
            "WHERE cache_key=?",
            (ckey,),
        )
        # Asking again must MISS and call client; new record inserted (or
        # rejected on PRIMARY KEY collision since cache_key is the PK).
        try:
            ask(conn, document_root=root, question="Q?", client=client, model_id="m")
            # If accepted, we'd have 2 records — but PK collision will throw.
            assert False, "expected PK collision on stale-then-write"
        except Exception:  # IntegrityError from cache_key PK
            pass

        # Falsified record stays stale; no fresh record landed.
        rows = conn.execute(
            "SELECT cache_key, falsification_state FROM providence_cache"
        ).fetchall()
        assert len(rows) == 1
        assert rows[0]["falsification_state"] == "stale"
        assert len(client.calls) == 2  # second call did happen
    finally:
        conn.close()


def test_unknown_document(tmp_path):
    db = tmp_path / "u.db"
    conn = connect(db)
    try:
        result = ask(
            conn,
            document_root="00" * 32,
            question="Q?",
            client=StubClient(),
            model_id="m",
        )
        assert result["status"] == "unknown_document"
    finally:
        conn.close()


def test_cold_source_refuses(tmp_path):
    """Source must be hot — answer derived from evicted content can't be proved."""
    from arborist.evict import evict_to_cold

    db = tmp_path / "cold.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
        evict_to_cold(conn)
        result = ask(
            conn,
            document_root=root,
            question="Q?",
            client=StubClient(),
            model_id="m",
        )
        assert result["status"] == "source_cold"
    finally:
        conn.close()


def test_cache_key_is_pure_function():
    """Hashes are deterministic; manual computation matches the runner."""
    src_root = "a" * 64
    qh = question_hash("What is X?")
    mh = model_profile_hash("m", "r", "q")
    msg = [{"role": "user", "content": "x"}]
    ch = conversation_hash(msg)
    gh = governance_policy_hash({"temperature": 0.1})
    k1 = cache_key(src_root, qh, mh, ch, gh, "v1", "v1", "v1")
    k2 = cache_key(src_root, qh, mh, ch, gh, "v1", "v1", "v1")
    assert k1 == k2
    # Bumping any dim changes the key.
    k3 = cache_key(src_root, qh, mh, ch, gh, "v2", "v1", "v1")
    assert k3 != k1
    # Pure SHA-256 of the joined string.
    expected = hashlib.sha256(
        "|".join([src_root, qh, mh, ch, gh, "v1", "v1", "v1"]).encode()
    ).hexdigest()
    assert k1 == expected


def test_concurrent_ask_same_cache_key_no_unique_crash(tmp_path):
    """Two threads run ``ask()`` on the same question/model concurrently.

    The cache lookup happens *outside* the write transaction, so both miss
    and both reach the ``providence_cache`` INSERT — ``ON CONFLICT(cache_key)
    DO NOTHING`` makes the loser no-op instead of raising
    ``UNIQUE constraint failed: providence_cache.cache_key``. Exactly one
    cache row lands. (Without the fix the losers raise; with ``busy_timeout``
    on every connection the loser waits on the writer's lock then no-ops.)
    """
    import threading

    db = tmp_path / "race-qa.db"
    conn = connect(db)
    try:
        root = _ingest_one(conn)
    finally:
        conn.close()

    n_threads = 6
    barrier = threading.Barrier(n_threads)
    errors: list[BaseException] = []
    lock = threading.Lock()

    def worker() -> None:
        try:
            c = connect(db)
            try:
                barrier.wait(timeout=20)
                ask(
                    c,
                    document_root=root,
                    question="What are the forms of capital?",
                    client=StubClient(answer="The forms include living and social."),
                    model_id="test-model",
                )
            finally:
                c.close()
        except BaseException as exc:  # noqa: BLE001 — surface it
            with lock:
                errors.append(exc)

    threads = [threading.Thread(target=worker) for _ in range(n_threads)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=40)

    assert not errors, f"concurrent ask() raised: {errors!r}"

    conn = connect(db)
    try:
        n_rows = conn.execute("SELECT COUNT(*) FROM providence_cache").fetchone()[0]
    finally:
        conn.close()
    assert n_rows == 1, f"expected exactly one providence_cache row, got {n_rows}"