engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

wallet: SPV-style cloud-only consumer (server + thin client + Merkle bundle)

Browse source 
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:
russell@unturf.com 2026-05-30 07:45:06 -04:00
parent d43714a503
commit 7e160584b6
No known key found for this signature in database
6 changed files with 1048 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

142
arborist/cli.py
View file

@ -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)

40
arborist/wallet/__init__.py Normal file
View file

@ -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",
]

124
arborist/wallet/client.py Normal file
View file

@ -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),
)

301
arborist/wallet/proof.py Normal file
View file

@ -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]}…"
)

173
arborist/wallet/server.py Normal file
View file

@ -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

268
tests/test_wallet_spv.py Normal file
View 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()