arborist/tests/test_wallet_spv.py

"""SPV-style wallet for arborist (`#wallet-in-cloud`).

End-to-end: corpus + server in-process → client query → Merkle verify
against the wallet's trust anchor. Plus the two failure modes a wallet
must catch — tampered chunk bytes and wrong trust anchor.
"""
from __future__ import annotations

from pathlib import Path
from typing import Iterator

import pytest

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.qa.client import StubClient
from arborist.snapshot import compute_snapshot_root
from arborist.source import Source
from arborist.store import connect
from arborist.wallet import (
    AnswerBundle,
    VerificationError,
    build_answer_bundle,
    verify_bundle,
)


class _FakeSource(Source):
    source_type = "test"

    def __init__(self, docs):
        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.split("/")[-1])


CORPUS_DOCS = [
    _doc(
        "test://doc/anarchism",
        (
            "Anarchism is a political philosophy that promotes a stateless society. "
            * 8
            + "It seeks to diminish or abolish authority in the conduct of human relations. "
            * 8
        ),
    ),
    _doc(
        "test://doc/capital",
        (
            "The eight forms of capital include living, social, and intellectual capital. "
            * 8
            + "Merkle providence proves that an answer derives from a specific source. "
            * 8
        ),
    ),
]


@pytest.fixture
def corpus_db(tmp_path: Path) -> Path:
    db = tmp_path / "corpus.db"
    conn = connect(db)
    try:
        ingest_source(conn, _FakeSource(CORPUS_DOCS))
    finally:
        conn.close()
    return db


def _make_fake_result(conn, uri: str, *, used: bool = True) -> dict:
    """Construct a minimal query-result dict pointing at one ingested doc.

    Bypasses the LLM + retrieval pipeline so we can unit-test the
    proof bundle in isolation. The shape matches what `query()`
    returns for a lattice-mode call with one cited source.
    """
    row = conn.execute(
        "SELECT document_root, title FROM documents WHERE document_uri = ?",
        (uri,),
    ).fetchone()
    assert row is not None
    return {
        "answer_text": f"[stub] this answer cites {uri}",
        "audit_mode": "STRICT",
        "cache_key": "stub-cache-key",
        "context_root": "stub-context-root",
        "sources": [
            {
                "document_root": row["document_root"],
                "document_uri": uri,
                "title": row["title"],
                "score": 1.0,
                "chunk_idx": 0,
                "shard": "single.db",
                "source_role": "primary_answer_source",
                "used": used,
                "used_pointer_ids": ["E1"] if used else [],
            }
        ],
    }


# --- proof builder + verifier (unit, no server) -----------------------------


def test_proof_bundle_verifies_against_correct_anchor(corpus_db: Path):
    conn = connect(corpus_db)
    try:
        snap, doc_count = compute_snapshot_root(conn)
        assert doc_count == len(CORPUS_DOCS)
        result = _make_fake_result(conn, "test://doc/anarchism")
        bundle = build_answer_bundle(conn, result=result, snapshot_root=snap)
    finally:
        conn.close()
    # Bundle ships chunks of the cited doc (no exception on verify).
    assert len(bundle.chunks) > 0
    assert all(c["document_root"] == result["sources"][0]["document_root"]
               for c in bundle.chunks)
    verify_bundle(bundle, trust_anchor=snap)


def test_proof_bundle_roundtrips_through_dict(corpus_db: Path):
    """Bundle → dict → AnswerBundle.from_dict → verify still passes.

    The wire format is JSON; this asserts the round-trip preserves
    every field the verifier reads."""
    conn = connect(corpus_db)
    try:
        snap, _ = compute_snapshot_root(conn)
        result = _make_fake_result(conn, "test://doc/capital")
        bundle = build_answer_bundle(conn, result=result, snapshot_root=snap)
    finally:
        conn.close()
    round_tripped = AnswerBundle.from_dict(bundle.to_dict())
    verify_bundle(round_tripped, trust_anchor=snap)


def test_proof_fails_on_tampered_chunk_body(corpus_db: Path):
    """Server lies: substitute different bytes for one chunk body.
    Wallet recomputes hash_leaf(body), sees it doesn't match the
    declared leaf_hash, and rejects.
    """
    conn = connect(corpus_db)
    try:
        snap, _ = compute_snapshot_root(conn)
        result = _make_fake_result(conn, "test://doc/anarchism")
        bundle = build_answer_bundle(conn, result=result, snapshot_root=snap)
    finally:
        conn.close()
    tampered = AnswerBundle.from_dict(bundle.to_dict())
    # Mutate chunk[0].body — leaf_hash stays the same, so the body
    # hash check fails first.
    tampered.chunks[0]["body"] = "ATTACKER PAYLOAD " + tampered.chunks[0]["body"]
    with pytest.raises(VerificationError, match="body hash"):
        verify_bundle(tampered, trust_anchor=snap)


def test_proof_fails_on_wrong_trust_anchor(corpus_db: Path):
    """Wallet has the wrong snapshot_root → verify immediately fails."""
    conn = connect(corpus_db)
    try:
        snap, _ = compute_snapshot_root(conn)
        result = _make_fake_result(conn, "test://doc/anarchism")
        bundle = build_answer_bundle(conn, result=result, snapshot_root=snap)
    finally:
        conn.close()
    bogus_anchor = "ff" * 32
    with pytest.raises(VerificationError, match="snapshot_root mismatch"):
        verify_bundle(bundle, trust_anchor=bogus_anchor)


def test_proof_fails_on_forged_leaf_hash(corpus_db: Path):
    """Server claims a body has a different leaf_hash than the bundle
    declares (mismatch between chunks[i].leaf_hash and
    chunk_proofs[i].leaf_hash). Wallet catches the mismatch before
    even hashing."""
    conn = connect(corpus_db)
    try:
        snap, _ = compute_snapshot_root(conn)
        result = _make_fake_result(conn, "test://doc/anarchism")
        bundle = build_answer_bundle(conn, result=result, snapshot_root=snap)
    finally:
        conn.close()
    tampered = AnswerBundle.from_dict(bundle.to_dict())
    tampered.chunks[0]["leaf_hash"] = "ab" * 32
    with pytest.raises(VerificationError, match="leaf_hash"):
        verify_bundle(tampered, trust_anchor=snap)


# --- HTTP server + thin client (end-to-end) ---------------------------------


def test_server_healthz_and_snapshot_root(corpus_db: Path, tmp_path: Path):
    """Server returns /healthz and /snapshot_root over real HTTP."""
    from arborist.wallet.client import WalletClient
    from arborist.wallet.server import WalletServer, serve_in_thread

    wallet_server = WalletServer(
        qa_db=tmp_path / "qa.db",
        single_db=corpus_db,
    )
    httpd, _thread, base_url = serve_in_thread(server=wallet_server, port=0)
    try:
        conn = connect(corpus_db)
        try:
            snap, _ = compute_snapshot_root(conn)
        finally:
            conn.close()

        client = WalletClient(server_url=base_url, trust_anchor=snap)
        assert client.healthz() == {"status": "ok"}
        resp = client.snapshot_root()
        assert resp["snapshot_root"] == snap
        assert resp["doc_count"] == len(CORPUS_DOCS)
    finally:
        httpd.shutdown()


def test_client_local_verify_quotes_runs_on_authenticated_chunks(
    corpus_db: Path, tmp_path: Path
):
    """verify_locally=True re-grounds the answer on the bundle's chunks
    (which are now cryptographically authenticated). A verbatim quote
    answer reaches STRICT locally; a fabricated answer falls to
    UNGROUNDED locally even if the server claims STRICT — that's the
    'wallet doesn't have to trust the server's verifier' property.
    """
    from arborist.wallet.client import WalletClient
    from arborist.wallet.server import WalletServer, serve_in_thread

    verbatim = (
        '"eight forms of capital include living, social, and intellectual"'
    )

    def _factory():
        return StubClient(answer=f"According to the source: {verbatim}.")

    wallet_server = WalletServer(
        qa_db=tmp_path / "qa.db",
        single_db=corpus_db,
        chat_client_factory=_factory,
    )
    httpd, _t, base_url = serve_in_thread(server=wallet_server, port=0)
    try:
        conn = connect(corpus_db)
        try:
            snap, _ = compute_snapshot_root(conn)
        finally:
            conn.close()
        client = WalletClient(server_url=base_url, trust_anchor=snap)
        verified = client.ask("What are the forms of capital?")
        # Local verifier ran and reached STRICT on the authenticated chunks.
        assert verified.local_audit_mode == "STRICT"
        assert verified.local_n_verified is not None
        assert verified.local_n_verified >= 1
        assert verified.local_verifier_method in {"quote", "span", "entity", "paraphrase"}
    finally:
        httpd.shutdown()


def test_server_ask_end_to_end_returns_verifiable_bundle(
    corpus_db: Path, tmp_path: Path
):
    """Wallet asks a question over HTTP; server runs query() against a
    real corpus; bundle comes back; wallet verifies it.

    Uses a StubClient that quotes a verbatim phrase from the corpus
    so the verifier reaches STRICT — gives us a clean STRICT path
    to bundle and verify."""
    from arborist.wallet.client import WalletClient
    from arborist.wallet.server import WalletServer, serve_in_thread

    # Verbatim phrase that appears in CORPUS_DOCS[1].
    verbatim = (
        '"eight forms of capital include living, social, and intellectual"'
    )

    def _factory():
        return StubClient(answer=f"According to the source: {verbatim}.")

    wallet_server = WalletServer(
        qa_db=tmp_path / "qa.db",
        single_db=corpus_db,
        chat_client_factory=_factory,
        default_top_k=4,
    )
    httpd, _thread, base_url = serve_in_thread(server=wallet_server, port=0)
    try:
        conn = connect(corpus_db)
        try:
            snap, _ = compute_snapshot_root(conn)
        finally:
            conn.close()

        client = WalletClient(server_url=base_url, trust_anchor=snap)
        verified = client.ask("What are the forms of capital?")
        # End-to-end success: verifies, no exception. Chunks fetched.
        assert verified.snapshot_root == snap
        assert verified.n_chunks_verified > 0
        assert verified.answer_text.startswith("According to the source:")
        # audit_mode is whatever query() produced; the test cares about
        # cryptographic verification, not the verifier verdict.
        assert verified.audit_mode in {"STRICT", "HYBRID", "UNGROUNDED"}
    finally:
        httpd.shutdown()