From 7ada42823fd95d4b5fb7eb72537e9cbe1e61b73e Mon Sep 17 00:00:00 2001
From: "russell@unturf.com" <russell@unturf.com>
Date: Sat, 30 May 2026 06:51:04 -0400
Subject: [PATCH] qa/query: JIT-hydrate doc chunks from bucket blobs on Tier B
 consumer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds hydrate_doc_jit(conn, doc_root, backend): for every chunk in this
document whose local content IS NULL, fetch blobs/<leaf_hash> from the
configured cold backend, verify hash, cache into the row. _load_doc_text
and _load_doc_chunks call it transparently before reading.

Env-gated: ARBORIST_JIT_CHUNKS=1 opts in (so non-JIT environments stay
a pure no-op); the backend comes from the standard ARBORIST_COLD_*
config. With this wired in, a Tier B consumer can clone metadata-only
shards and the query path transparently fetches answer chunks JIT — no
caller code changes.
---
 arborist/cold_clone.py | 50 ++++++++++++++++++++++++++++++++++++++++++
 arborist/qa/query.py   | 30 +++++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/arborist/cold_clone.py b/arborist/cold_clone.py
index 2cbb80c..7aa018f 100644
--- a/arborist/cold_clone.py
+++ b/arborist/cold_clone.py
@@ -386,6 +386,56 @@ def clone_from_bucket(
 # ---------- JIT chunk fetch (Tier B consumer) -------------------------------
 
 
+def hydrate_doc_jit(
+    conn: sqlite3.Connection,
+    document_root: str,
+    backend: ObjectStoreBackend,
+    *,
+    fetch_workers: int = 16,
+) -> int:
+    """Fetch any chunks for ``document_root`` whose ``content`` is NULL
+    in the local shard from ``blobs/<leaf_hash>`` on the bucket and
+    cache them back into the row. No-op if every chunk already has
+    content (the cheap WHERE count short-circuits). Returns the number
+    of chunks fetched.
+
+    Used by Tier B consumers ('clone --just-enough' produced metadata-
+    only shards): the query path calls this just before reading chunk
+    text for a hit, so the rest of the pipeline ('SELECT content FROM
+    chunks WHERE content IS NOT NULL') sees the same locally-resident
+    shape it would for a full clone."""
+    from arborist.merkle import hash_leaf
+
+    rows = conn.execute(
+        "SELECT chunk_id, leaf_hash FROM chunks "
+        "WHERE document_root = ? AND content IS NULL",
+        (document_root,),
+    ).fetchall()
+    if not rows:
+        return 0
+
+    def _fetch(pair: tuple[int, str]) -> tuple[int, bytes]:
+        chunk_id, leaf_hash = pair
+        key = f"{BLOB_PREFIX}/{leaf_hash[:2]}/{leaf_hash[2:]}"
+        body = backend.get(key)
+        if hash_leaf(body).hex() != leaf_hash:
+            raise ValueError(
+                f"chunk {chunk_id}: bucket blob hash mismatch (leaf_hash {leaf_hash[:12]})"
+            )
+        return chunk_id, body
+
+    pairs = [(r[0], r[1]) for r in rows]
+    n_workers = max(1, min(fetch_workers, len(pairs)))
+    with ThreadPoolExecutor(max_workers=n_workers) as ex:
+        fetched = list(ex.map(_fetch, pairs))
+    with conn:
+        conn.executemany(
+            "UPDATE chunks SET content = ? WHERE chunk_id = ?",
+            [(body, cid) for cid, body in fetched],
+        )
+    return len(fetched)
+
+
 def fetch_chunk_jit(
     conn: sqlite3.Connection,
     chunk_id: int,
diff --git a/arborist/qa/query.py b/arborist/qa/query.py
index 3a53b3b..d8595ff 100644
--- a/arborist/qa/query.py
+++ b/arborist/qa/query.py
@@ -1850,8 +1850,37 @@ def _rerank_by_body_coverage(
     return hits
 
 
+def _maybe_jit_hydrate(shard_path: str, document_root: str) -> None:
+    """If ARBORIST_JIT_CHUNKS=1 and a cold backend is configured
+    (ARBORIST_COLD_BUCKET / _ENDPOINT_URL), fetch any chunks for this
+    document whose ``content`` is NULL from ``blobs/<hash>`` on the
+    bucket and cache them locally. No-op otherwise. Lets a Tier B
+    consumer (metadata-only clone) serve queries by pulling answer
+    chunks from S3 just in time."""
+    import os
+    if os.environ.get("ARBORIST_JIT_CHUNKS", "0") != "1":
+        return
+    try:
+        from arborist.cli import _make_cold_backend
+        from arborist.cold_clone import hydrate_doc_jit
+        backend = _make_cold_backend()
+    except Exception:
+        return
+    conn = connect(shard_path)
+    try:
+        n = hydrate_doc_jit(conn, document_root, backend)
+        if n:
+            print(
+                f"[jit] hydrated {n} chunks for {document_root[:12]} from bucket",
+                file=sys.stderr,
+            )
+    finally:
+        conn.close()
+
+
 def _load_doc_text(shard_path: str, document_root: str) -> str | None:
     """Concatenate all hot chunks of a document. Returns None if cold or missing."""
+    _maybe_jit_hydrate(shard_path, document_root)
     conn = connect(shard_path)
     try:
         rows = conn.execute(
@@ -1880,6 +1909,7 @@ def _load_doc_chunks(
     paragraph that supports a claim instead of lazy-anchoring the
     whole article.
     """
+    _maybe_jit_hydrate(shard_path, document_root)
     conn = connect(shard_path)
     try:
         rows = conn.execute(