From 7e160584b6f51d76e24d0f7f5270c424892bfda9 Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Sat, 30 May 2026 07:45:06 -0400 Subject: [PATCH] wallet: SPV-style cloud-only consumer (server + thin client + Merkle bundle) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- arborist/cli.py | 142 +++++++++++++++++ arborist/wallet/__init__.py | 40 +++++ arborist/wallet/client.py | 124 +++++++++++++++ arborist/wallet/proof.py | 301 ++++++++++++++++++++++++++++++++++++ arborist/wallet/server.py | 173 +++++++++++++++++++++ tests/test_wallet_spv.py | 268 ++++++++++++++++++++++++++++++++ 6 files changed, 1048 insertions(+) create mode 100644 arborist/wallet/__init__.py create mode 100644 arborist/wallet/client.py create mode 100644 arborist/wallet/proof.py create mode 100644 arborist/wallet/server.py create mode 100644 tests/test_wallet_spv.py diff --git a/arborist/cli.py b/arborist/cli.py index 2038688..0a1491f 100644 --- a/arborist/cli.py +++ b/arborist/cli.py @@ -7152,9 +7152,151 @@ def build_parser() -> argparse.ArgumentParser: ) recrawl_check_cmd.set_defaults(func=_cmd_crawler_recrawl_check) + # --- wallet-in-cloud (SPV) ------------------------------------------- + serve_cmd = sub.add_parser( + "serve", + help=( + "start an HTTP wallet server: GET /snapshot_root, POST /ask. " + "Wraps query() and returns Merkle-verifiable AnswerBundles for " + "thin clients (`arborist wallet ask`) to verify against a " + "pinned trust anchor — same SPV pattern as Bitcoin/Electrum." + ), + ) + serve_cmd.add_argument("--host", default="127.0.0.1") + serve_cmd.add_argument("--port", type=int, default=8765) + serve_cmd.add_argument( + "--qa-db", default="~/.arborist/qa.db", + help="QA cache DB path (created on first call).", + ) + serve_cmd.add_argument( + "--top-k", dest="serve_top_k", type=int, default=4, + help="default top_k for /ask requests that don't specify one.", + ) + serve_cmd.set_defaults(func=_cmd_serve) + + wallet_cmd = sub.add_parser( + "wallet", + help="SPV wallet client — query a remote arborist server and verify.", + ) + wallet_sub = wallet_cmd.add_subparsers(dest="wallet_cmd", required=True) + + wallet_ask = wallet_sub.add_parser( + "ask", + help=( + "submit a question to the server, verify the returned " + "AnswerBundle against the wallet's trust anchor" + ), + ) + wallet_ask.add_argument("question", type=str) + wallet_ask.add_argument( + "--server-url", required=True, + help="base URL of the arborist wallet server (e.g. http://host:8765)", + ) + wallet_ask.add_argument( + "--trust-anchor", required=True, + help="hex snapshot_root the wallet trusts; out-of-band-set", + ) + wallet_ask.add_argument( + "--top-k", type=int, default=None, + help="optional top_k override on the server.", + ) + wallet_ask.set_defaults(func=_cmd_wallet_ask) + + wallet_anchor = wallet_sub.add_parser( + "anchor", + help=( + "fetch the server's current snapshot_root (advisory — does NOT " + "update the wallet's trust anchor)" + ), + ) + wallet_anchor.add_argument("--server-url", required=True) + wallet_anchor.set_defaults(func=_cmd_wallet_anchor) + return p +def _cmd_serve(args: argparse.Namespace) -> int: + """Start the HTTP wallet server. Blocks until SIGINT/SIGTERM.""" + from arborist.wallet.server import WalletServer, serve + + shards_dir = getattr(args, "shards_dir", None) + single_db = getattr(args, "db", None) + if not shards_dir and not single_db: + print( + "arborist serve needs --shards-dir or --db pointing at the corpus", + file=sys.stderr, + ) + return 2 + qa_db = Path(args.qa_db).expanduser() + qa_db.parent.mkdir(parents=True, exist_ok=True) + + chat_client_factory = None + if os.environ.get("ARBORIST_WALLET_STUB", "0") == "1": + from arborist.qa.client import StubClient + def chat_client_factory(): + return StubClient(answer="[stub] no LLM in --stub mode") + else: + endpoint = os.environ.get("ARBORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1") + model_id = os.environ.get("ARBORIST_LLM_MODEL", "NousResearch/Hermes-3-Llama-3.1-8B-FP8-Dynamic") + from arborist.qa.client import OpenAICompatibleClient + def chat_client_factory(): + return OpenAICompatibleClient(endpoint=endpoint) + + wsrv = WalletServer( + qa_db=qa_db, + shards_dir=Path(shards_dir) if shards_dir else None, + single_db=Path(single_db) if single_db else None, + chat_client_factory=chat_client_factory, + model_id=os.environ.get("ARBORIST_LLM_MODEL", "stub"), + default_top_k=args.serve_top_k, + ) + httpd = serve(server=wsrv, host=args.host, port=args.port) + print( + f"arborist wallet server listening on http://{args.host}:{args.port}", + file=sys.stderr, + ) + try: + httpd.serve_forever() + except KeyboardInterrupt: + print("\nshutting down", file=sys.stderr) + finally: + httpd.shutdown() + return 0 + + +def _cmd_wallet_ask(args: argparse.Namespace) -> int: + from arborist.wallet.client import WalletClient + from arborist.wallet.proof import VerificationError + + client = WalletClient( + server_url=args.server_url, + trust_anchor=args.trust_anchor, + ) + try: + verified = client.ask(args.question, top_k=args.top_k) + except VerificationError as e: + print(f"VERIFICATION FAILED: {e}", file=sys.stderr) + return 3 + print(json.dumps({ + "audit_mode": verified.audit_mode, + "cache_key": verified.cache_key, + "snapshot_root": verified.snapshot_root, + "n_chunks_verified": verified.n_chunks_verified, + "sources": verified.sources, + "answer_text": verified.answer_text, + }, indent=2, ensure_ascii=False)) + return 0 + + +def _cmd_wallet_anchor(args: argparse.Namespace) -> int: + from arborist.wallet.client import WalletClient + # trust_anchor is required by the constructor but unused for + # /snapshot_root (it's an advisory fetch). + client = WalletClient(server_url=args.server_url, trust_anchor="00" * 32) + print(json.dumps(client.snapshot_root(), indent=2)) + return 0 + + def main(argv: list[str] | None = None) -> int: args = build_parser().parse_args(argv) return args.func(args) diff --git a/arborist/wallet/__init__.py b/arborist/wallet/__init__.py new file mode 100644 index 0000000..18623f2 --- /dev/null +++ b/arborist/wallet/__init__.py @@ -0,0 +1,40 @@ +"""SPV-style wallet for arborist (#wallet-in-cloud). + +Pattern: full arborist node lives in the cloud; a thin wallet client +holds only a trust anchor (``snapshot_root``) and verifies Merkle proofs +on each answer. Same shape as Bitcoin SPV (Electrum / mobile wallet): + + Bitcoin SPV Arborist SPV + ---------------------- ----------------------------------- + Block header hash snapshot_root (already exists) + Transaction ID cache_key (already exists) + SPV Merkle branch AnswerBundle.chunk_proofs + doc_proofs + Wallet keys peer identity (already exists) + Full node arborist serve + +The wallet has no SQLite, no FTS, no chunks. It receives an +``AnswerBundle`` from the server containing the answer, the cited +chunk bodies, and Merkle proofs (chunk → document_root, document_root +→ snapshot_root). A malicious server can DOS but cannot forge content +that verifies against the wallet's trust anchor. + +Submodules: + proof build_answer_bundle (server side), verify_bundle (client side) + server HTTP server (GET /snapshot_root, POST /ask) + client WalletClient: thin verifying HTTP client +""" +from __future__ import annotations + +from arborist.wallet.proof import ( + AnswerBundle, + VerificationError, + build_answer_bundle, + verify_bundle, +) + +__all__ = [ + "AnswerBundle", + "VerificationError", + "build_answer_bundle", + "verify_bundle", +] diff --git a/arborist/wallet/client.py b/arborist/wallet/client.py new file mode 100644 index 0000000..4eab6d6 --- /dev/null +++ b/arborist/wallet/client.py @@ -0,0 +1,124 @@ +"""Thin verifying wallet client. No SQLite, no FTS, no corpus. + +The wallet holds: + - ``server_url``: the cloud arborist node URL + - ``trust_anchor``: a ``snapshot_root`` the wallet trusts (set out- + of-band by pinning / pairing / quorum — see CLAUDE.md wallet + bootstrap notes) + +On every ``ask()`` the wallet receives an :class:`AnswerBundle` and +runs the SPV verification path. A passing verification means: the +server returned content that hashes up to the snapshot the wallet +trusts. A failing verification means: server lied, server has a +different corpus, or wallet has the wrong anchor. + +stdlib only — works in a venv with nothing installed. +""" +from __future__ import annotations + +import json +import urllib.error +import urllib.request +from dataclasses import dataclass + +from arborist.wallet.proof import AnswerBundle, verify_bundle + + +class WalletError(Exception): + """Network / protocol error talking to the server.""" + + +@dataclass +class VerifiedAnswer: + """A wallet's verified view of a server's response.""" + + answer_text: str + audit_mode: str + cache_key: str | None + snapshot_root: str + sources: list[dict] + n_chunks_verified: int + + +class WalletClient: + """SPV-style client. Trust = math, not server.""" + + def __init__( + self, + *, + server_url: str, + trust_anchor: str, + timeout_s: float = 30.0, + ): + self.server_url = server_url.rstrip("/") + self.trust_anchor = trust_anchor + self.timeout_s = timeout_s + + # --- raw HTTP helpers --- + + def _get(self, path: str) -> dict: + url = f"{self.server_url}{path}" + try: + with urllib.request.urlopen(url, timeout=self.timeout_s) as resp: + return json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as e: + raise WalletError(f"{path}: HTTP {e.code} {e.reason}") from e + except urllib.error.URLError as e: + raise WalletError(f"{path}: {e.reason}") from e + + def _post(self, path: str, body: dict) -> dict: + url = f"{self.server_url}{path}" + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + url, data=data, headers={"Content-Type": "application/json"} + ) + try: + with urllib.request.urlopen(req, timeout=self.timeout_s) as resp: + return json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as e: + detail = "" + try: + detail = e.read().decode("utf-8", errors="replace") + except Exception: + pass + raise WalletError(f"{path}: HTTP {e.code} {e.reason} {detail}") from e + except urllib.error.URLError as e: + raise WalletError(f"{path}: {e.reason}") from e + + # --- wallet API --- + + def snapshot_root(self) -> dict: + """Fetch the server's current snapshot_root (advisory; does NOT + change the wallet's trust anchor — anchor updates are an + out-of-band ceremony per the bootstrap protocol).""" + return self._get("/snapshot_root") + + def healthz(self) -> dict: + return self._get("/healthz") + + def ask(self, question: str, *, top_k: int | None = None) -> VerifiedAnswer: + """Submit a question, receive an AnswerBundle, verify it against + the wallet's trust anchor. Returns a VerifiedAnswer or raises + ``VerificationError`` (proof) / ``WalletError`` (network). + + A passing return is a cryptographic guarantee that the answer's + cited chunks are part of the corpus committed to by + ``self.trust_anchor``. It does NOT guarantee the LLM phrased + the answer truthfully w.r.t. those chunks — the wallet can run + the existing ``arborist.qa.verify`` verifier against the bundle + locally for that, since the bundle ships the actual chunk bytes. + """ + payload: dict = {"question": question} + if top_k is not None: + payload["top_k"] = top_k + resp = self._post("/ask", payload) + bundle = AnswerBundle.from_dict(resp) + verify_bundle(bundle, trust_anchor=self.trust_anchor) + return VerifiedAnswer( + answer_text=bundle.answer_text, + audit_mode=bundle.audit_mode, + cache_key=bundle.cache_key, + snapshot_root=bundle.snapshot_root, + sources=bundle.sources, + n_chunks_verified=len(bundle.chunks), + ) diff --git a/arborist/wallet/proof.py b/arborist/wallet/proof.py new file mode 100644 index 0000000..2e5b9de --- /dev/null +++ b/arborist/wallet/proof.py @@ -0,0 +1,301 @@ +"""Answer-bundle Merkle proof: chunk → document_root → snapshot_root. + +Server side (``build_answer_bundle``): given a ``query()`` result and a +SQLite connection, gather every cited chunk's body, build the two +proof legs, and serialize a self-contained ``AnswerBundle``. + +Client side (``verify_bundle``): given an ``AnswerBundle`` and a trusted +``snapshot_root``, re-derive every chunk's leaf_hash, walk its proof to +the document_root, walk the document_root proof to snapshot_root, and +check the final root matches the trust anchor. No SQLite required — +pure cryptographic verification on stdlib primitives only. + +Threat model is exactly Bitcoin-SPV: + + Server CAN: refuse to answer (DOS), return UNGROUNDED, censor. + Server CANNOT: forge content whose hash chains up to the wallet's + snapshot_root (SHA-256 preimage-resistance). +""" +from __future__ import annotations + +import sqlite3 +from dataclasses import dataclass, field +from typing import Any + +from arborist.compress import unpack_chunk +from arborist.merkle import ( + MerkleTree, + ProofNode, + hash_combine, + hash_leaf, +) + + +class VerificationError(Exception): + """Raised when an AnswerBundle fails to verify against its trust anchor.""" + + +@dataclass(frozen=True) +class AnswerBundle: + """Self-contained wallet-verifiable answer. + + The bundle is the on-wire shape the server returns and the wallet + consumes. ``snapshot_root`` is the commitment the wallet's trust + anchor must match for the bundle to verify. + + Fields: + answer_text: LLM-produced answer string + audit_mode: verifier verdict (STRICT / HYBRID / UNGROUNDED) + cache_key: 8-dim content-addressed query identity + snapshot_root: hex, the corpus state the bundle commits to + sources: list of source dicts from query() result + chunks: every cited chunk: {document_root, leaf_hash, + leaf_index, body} + chunk_proofs: chunk → document_root inclusion proof per chunk + doc_proofs: document_root → snapshot_root inclusion proof + per cited document + """ + + answer_text: str + audit_mode: str + cache_key: str | None + snapshot_root: str + sources: list[dict] = field(default_factory=list) + chunks: list[dict] = field(default_factory=list) + chunk_proofs: list[dict] = field(default_factory=list) + doc_proofs: list[dict] = field(default_factory=list) + + def to_dict(self) -> dict: + return { + "answer_text": self.answer_text, + "audit_mode": self.audit_mode, + "cache_key": self.cache_key, + "snapshot_root": self.snapshot_root, + "sources": self.sources, + "chunks": self.chunks, + "chunk_proofs": self.chunk_proofs, + "doc_proofs": self.doc_proofs, + } + + @classmethod + def from_dict(cls, d: dict) -> "AnswerBundle": + return cls( + answer_text=d["answer_text"], + audit_mode=d["audit_mode"], + cache_key=d.get("cache_key"), + snapshot_root=d["snapshot_root"], + sources=list(d.get("sources") or []), + chunks=list(d.get("chunks") or []), + chunk_proofs=list(d.get("chunk_proofs") or []), + doc_proofs=list(d.get("doc_proofs") or []), + ) + + +def _all_document_roots_sorted(conn: sqlite3.Connection) -> list[str]: + """Same shape as snapshot._all_document_roots, inlined to keep the + proof module self-contained (snapshot.py imports get heavy).""" + rows = conn.execute( + "SELECT DISTINCT document_root FROM documents ORDER BY document_root ASC" + ).fetchall() + return [r[0] for r in rows] + + +def _doc_root_from_chunks(conn: sqlite3.Connection, droot: str) -> tuple[MerkleTree, list[dict]]: + """Rebuild a document's chunk Merkle tree + return (tree, rows). + + Rows are ordered by ``idx`` so leaf_index alignment matches the + document_root computed at ingest time (``ingest._compute_artifacts``). + """ + rows = conn.execute( + "SELECT idx, leaf_hash, content FROM chunks " + "WHERE document_root = ? AND content IS NOT NULL " + " AND length(content) > 0 " + " AND substr(content, 1, 1) != X'00' " + "ORDER BY idx ASC", + (droot,), + ).fetchall() + leaves = [bytes.fromhex(r[1]) for r in rows] + return MerkleTree.build(leaves), rows + + +def build_answer_bundle( + conn: sqlite3.Connection, + *, + result: dict, + snapshot_root: str, +) -> AnswerBundle: + """Bundle a ``query()`` result into a wallet-verifiable AnswerBundle. + + The bundle ships every chunk of every cited source. For lattice + modes only the ``used=True`` sources are bundled; for non-lattice + modes (quote / span / entity / paraphrase) every source flagged + by retrieval is bundled (no per-chunk used annotation in those + modes). Either way the wallet sees the exact bytes the verifier + grounded against. + """ + sources = result.get("sources") or [] + any_used_flag = any("used" in s for s in sources) + if any_used_flag: + cited = [s for s in sources if s.get("used")] + else: + cited = list(sources) + + all_roots = _all_document_roots_sorted(conn) + if all_roots: + snapshot_tree = MerkleTree.build([bytes.fromhex(r) for r in all_roots]) + else: + snapshot_tree = MerkleTree.build([]) + root_to_idx = {r: i for i, r in enumerate(all_roots)} + + chunks: list[dict] = [] + chunk_proofs: list[dict] = [] + doc_proofs: list[dict] = [] + seen_doc_roots: set[str] = set() + + for src in cited: + droot = src.get("document_root") + if not droot or droot in seen_doc_roots: + continue + seen_doc_roots.add(droot) + + tree, rows = _doc_root_from_chunks(conn, droot) + if not rows: + continue + + for leaf_idx, row in enumerate(rows): + body_text = unpack_chunk(row[2]) or "" + chunks.append({ + "document_root": droot, + "leaf_hash": row[1], + "leaf_index": leaf_idx, + "body": body_text, + }) + p = tree.proof(leaf_idx) + chunk_proofs.append({ + "document_root": droot, + "leaf_index": leaf_idx, + "leaf_hash": row[1], + "siblings": [ + {"hash": s.hash.hex(), "is_left": s.is_left} + for s in p.siblings + ], + }) + + if droot in root_to_idx and all_roots: + if len(all_roots) == 1: + # Degenerate snapshot: doc_root == snapshot_root, no + # siblings needed (matches compute_snapshot_root which + # short-circuits the single-leaf case). + doc_proofs.append({ + "document_root": droot, + "leaf_index": 0, + "siblings": [], + "degenerate_single_doc": True, + }) + else: + doc_idx = root_to_idx[droot] + p = snapshot_tree.proof(doc_idx) + doc_proofs.append({ + "document_root": droot, + "leaf_index": doc_idx, + "siblings": [ + {"hash": s.hash.hex(), "is_left": s.is_left} + for s in p.siblings + ], + }) + + return AnswerBundle( + answer_text=result.get("answer_text", ""), + audit_mode=result.get("audit_mode", "UNGROUNDED"), + cache_key=result.get("cache_key"), + snapshot_root=snapshot_root, + sources=cited, + chunks=chunks, + chunk_proofs=chunk_proofs, + doc_proofs=doc_proofs, + ) + + +def _walk_siblings(start: bytes, siblings: list[dict]) -> bytes: + """Combine ``start`` with sibling hashes per their ``is_left`` flag.""" + current = start + for s in siblings: + sib = bytes.fromhex(s["hash"]) + if s["is_left"]: + current = hash_combine(sib, current) + else: + current = hash_combine(current, sib) + return current + + +def verify_bundle(bundle: AnswerBundle, *, trust_anchor: str) -> None: + """Verify a bundle against a trusted snapshot_root. Raises on failure. + + Steps per chunk: + 1. hash_leaf(body.encode('utf-8')) must equal the chunk's + declared leaf_hash. (Server cannot lie about chunk bytes.) + 2. Walking the chunk_proof siblings from leaf_hash must reach + the chunk's document_root. (Chunk really belongs to that doc.) + 3. Walking the doc_proof siblings from document_root must reach + the bundle's snapshot_root. (Doc really belongs to that + snapshot.) + 4. The bundle's snapshot_root must equal the trust_anchor the + wallet was configured with. (We're looking at the same + corpus state the wallet trusts.) + + Any mismatch raises VerificationError. A passing run gives the + wallet a cryptographic guarantee that the chunk bodies it just + received are part of the corpus committed to by trust_anchor. + """ + if bundle.snapshot_root != trust_anchor: + raise VerificationError( + f"snapshot_root mismatch: bundle={bundle.snapshot_root[:12]}… " + f"trust_anchor={trust_anchor[:12]}…" + ) + + # Index doc_proofs by document_root for chunk-side lookup. + doc_proof_by_root: dict[str, dict] = { + dp["document_root"]: dp for dp in bundle.doc_proofs + } + + for i, (chunk, cp) in enumerate(zip(bundle.chunks, bundle.chunk_proofs)): + if chunk["leaf_hash"] != cp["leaf_hash"]: + raise VerificationError( + f"chunk[{i}]: chunk.leaf_hash != chunk_proof.leaf_hash" + ) + body_bytes = chunk["body"].encode("utf-8") + computed_leaf = hash_leaf(body_bytes).hex() + if computed_leaf != chunk["leaf_hash"]: + raise VerificationError( + f"chunk[{i}]: body hash {computed_leaf[:12]}… does not " + f"match declared leaf_hash {chunk['leaf_hash'][:12]}…" + ) + + # Chunk → document_root + recomputed_doc = _walk_siblings( + bytes.fromhex(chunk["leaf_hash"]), cp["siblings"] + ).hex() + if recomputed_doc != chunk["document_root"]: + raise VerificationError( + f"chunk[{i}]: chunk_proof recomputed doc_root " + f"{recomputed_doc[:12]}… != declared {chunk['document_root'][:12]}…" + ) + + # Document_root → snapshot_root + dp = doc_proof_by_root.get(chunk["document_root"]) + if dp is None: + raise VerificationError( + f"chunk[{i}]: missing doc_proof for {chunk['document_root'][:12]}…" + ) + if dp.get("degenerate_single_doc"): + recomputed_snap = chunk["document_root"] + else: + recomputed_snap = _walk_siblings( + bytes.fromhex(chunk["document_root"]), dp["siblings"] + ).hex() + if recomputed_snap != bundle.snapshot_root: + raise VerificationError( + f"chunk[{i}]: doc_proof recomputed snapshot_root " + f"{recomputed_snap[:12]}… != bundle.snapshot_root " + f"{bundle.snapshot_root[:12]}…" + ) diff --git a/arborist/wallet/server.py b/arborist/wallet/server.py new file mode 100644 index 0000000..cccf1a3 --- /dev/null +++ b/arborist/wallet/server.py @@ -0,0 +1,173 @@ +"""HTTP server: wraps query() and serves wallet AnswerBundles. + +Endpoints: + GET /snapshot_root → {snapshot_root, doc_count} + POST /ask → AnswerBundle.to_dict() + GET /healthz → {status: "ok"} + +Pure stdlib — no Flask, no FastAPI. ``http.server.ThreadingHTTPServer`` +is enough for the demo and for low-traffic SPV peers. For higher +throughput, sit nginx in front and run multiple workers; or swap in +gunicorn (out of scope for this module). +""" +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Callable + +from arborist.qa.client import StubClient +from arborist.snapshot import compute_snapshot_root +from arborist.store import connect, connect_query +from arborist.wallet.proof import build_answer_bundle + + +class WalletServer: + """Holds the corpus connection factory + ask handler. Stateless across + requests (each request opens its own DB connection so we sidestep + SQLite's single-writer per connection model).""" + + def __init__( + self, + *, + qa_db: Path, + shards_dir: Path | None = None, + single_db: Path | None = None, + chat_client_factory: Callable[[], object] | None = None, + model_id: str = "stub", + default_top_k: int = 4, + ): + self.qa_db = Path(qa_db) + self.shards_dir = Path(shards_dir) if shards_dir else None + self.single_db = Path(single_db) if single_db else None + if shards_dir is None and single_db is None: + raise ValueError("either shards_dir or single_db is required") + self.chat_client_factory = chat_client_factory or (lambda: StubClient()) + self.model_id = model_id + self.default_top_k = default_top_k + + def _open_corpus(self): + """Open a fresh corpus connection for this request.""" + if self.shards_dir: + return connect_query(None, shards_dir=self.shards_dir) + return connect(self.single_db) + + def snapshot_root(self) -> dict: + conn = self._open_corpus() + try: + root, count = compute_snapshot_root(conn) + finally: + conn.close() + return {"snapshot_root": root, "doc_count": count} + + def ask(self, *, question: str, top_k: int | None = None) -> dict: + """Run query, bundle answer + Merkle proofs, return dict.""" + from arborist.qa.query import query + + chat = self.chat_client_factory() + result = query( + question=question, + qa_db=self.qa_db, + chat_client=chat, + model_id=self.model_id, + shards_dir=self.shards_dir, + single_db=self.single_db, + top_k=top_k or self.default_top_k, + ) + conn = self._open_corpus() + try: + snap, _ = compute_snapshot_root(conn) + bundle = build_answer_bundle(conn, result=result, snapshot_root=snap) + finally: + conn.close() + return bundle.to_dict() + + +def make_handler(server: WalletServer) -> type: + """Wrap WalletServer in a BaseHTTPRequestHandler subclass. + + The handler class needs access to the WalletServer instance; the + canonical http.server pattern is a closure-captured factory. + """ + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _write_json(self, status: int, body: dict) -> None: + payload = json.dumps(body).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, fmt, *args): + # Default log writes to stderr; keep it quiet in tests. + # Operators wanting access logs can set ARBORIST_WALLET_VERBOSE=1. + import os + if os.environ.get("ARBORIST_WALLET_VERBOSE", "0") == "1": + super().log_message(fmt, *args) + + def do_GET(self): + try: + if self.path == "/healthz": + self._write_json(200, {"status": "ok"}) + elif self.path == "/snapshot_root": + self._write_json(200, server.snapshot_root()) + else: + self._write_json(404, {"error": "not found"}) + except Exception as e: + self._write_json(500, {"error": str(e)}) + + def do_POST(self): + try: + if self.path != "/ask": + self._write_json(404, {"error": "not found"}) + return + length = int(self.headers.get("Content-Length") or 0) + body = json.loads(self.rfile.read(length).decode("utf-8")) if length else {} + question = body.get("question") + if not question: + self._write_json(400, {"error": "missing 'question'"}) + return + top_k = body.get("top_k") + bundle = server.ask(question=question, top_k=top_k) + self._write_json(200, bundle) + except Exception as e: + self._write_json(500, {"error": str(e), "type": type(e).__name__}) + + return Handler + + +def serve( + *, + server: WalletServer, + host: str = "127.0.0.1", + port: int = 8765, +) -> ThreadingHTTPServer: + """Start the HTTP server. Returns the server instance; caller + drives ``serve_forever()`` or runs it in a background thread.""" + handler_cls = make_handler(server) + httpd = ThreadingHTTPServer((host, port), handler_cls) + return httpd + + +def serve_in_thread( + *, + server: WalletServer, + host: str = "127.0.0.1", + port: int = 0, +) -> tuple[ThreadingHTTPServer, threading.Thread, str]: + """Start the server in a daemon thread. Returns (httpd, thread, url). + + Use ``port=0`` to bind any free port. The returned URL is the + base URL clients should hit (``http://host:port``). + """ + httpd = serve(server=server, host=host, port=port) + actual_port = httpd.server_address[1] + base_url = f"http://{host}:{actual_port}" + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + return httpd, thread, base_url diff --git a/tests/test_wallet_spv.py b/tests/test_wallet_spv.py new file mode 100644 index 0000000..08409fe --- /dev/null +++ b/tests/test_wallet_spv.py @@ -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()