wallet: SPV-style cloud-only consumer (server + thin client + Merkle bundle)
A wallet client holds only a snapshot_root (trust anchor) and verifies
Merkle proofs on every answer. No SQLite, no FTS, no chunks locally.
Same shape as Bitcoin SPV (Electrum / mobile wallet): server can DOS
but cannot forge content whose hash chains up to the trusted anchor.
New module `arborist/wallet/`:
- proof.py: AnswerBundle + build_answer_bundle (server) +
verify_bundle (client). Two proof legs per chunk:
chunk_body → leaf_hash → document_root via in-doc
Merkle proof, then document_root → snapshot_root via
the corpus-wide sorted-doc-roots tree (mirrors
snapshot.compute_snapshot_root). Single-doc corpus
degenerates to "document_root IS snapshot_root" and is
handled with an explicit `degenerate_single_doc` flag.
- server.py: WalletServer + http.server.ThreadingHTTPServer wrapper.
Pure stdlib. GET /healthz, GET /snapshot_root, POST /ask.
Each request opens its own DB connection so SQLite's
single-writer model never bites.
- client.py: WalletClient: urllib + json + arborist.wallet.proof.
Returns a VerifiedAnswer or raises VerificationError /
WalletError. No corpus dependency.
New CLI:
- `arborist serve` — start the wallet server. ARBORIST_WALLET_STUB=1
swaps the LLM for StubClient (lets ops sanity-check verification
without burning tokens).
- `arborist wallet anchor` — fetch the server's current snapshot_root.
- `arborist wallet ask` — submit a question, verify the AnswerBundle
against --trust-anchor, exit 3 on VerificationError.
Tests (tests/test_wallet_spv.py, 7 cases):
- happy: bundle → verify pass against correct anchor
- dict round-trip via to_dict/from_dict still verifies
- tamper: rewrite a chunk body → body hash check fails
- forged leaf_hash: chunks[i].leaf_hash != chunk_proofs[i].leaf_hash
fails before any hashing
- wrong trust_anchor: bundle.snapshot_root != anchor fails immediately
- /healthz and /snapshot_root over real HTTP
- end-to-end ask: corpus → in-process server → urllib client → verify
This commit is contained in:
parent
d43714a503
commit
7e160584b6
6 changed files with 1048 additions and 0 deletions
268
tests/test_wallet_spv.py
Normal file
268
tests/test_wallet_spv.py
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
"""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_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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue