diff --git a/mempalace/cli.py b/mempalace/cli.py index c17078f768..7428674225 100644 --- a/mempalace/cli.py +++ b/mempalace/cli.py @@ -774,7 +774,12 @@ def cmd_repair_status(args): def cmd_repair(args): - """Rebuild palace vector index from SQLite metadata.""" + """Rebuild palace vector index from SQLite metadata. + + ``--mode hnsw`` dispatches to the segment-level rebuild path + (:func:`mempalace.repair.rebuild_hnsw_segment`); the legacy default + mode rebuilds the whole collection via re-embed. + """ import shutil from .backends.chroma import ChromaBackend from .migrate import confirm_destructive_action, contains_palace_database @@ -796,6 +801,24 @@ def cmd_repair(args): os.path.expanduser(args.palace) if args.palace else config.palace_path ) + if getattr(args, "mode", "legacy") == "hnsw": + if not getattr(args, "segment", None): + print(" --mode hnsw requires --segment ") + return + from .repair import rebuild_hnsw_segment + + rebuild_hnsw_segment( + palace_path, + segment=args.segment, + max_elements=getattr(args, "max_elements", None), + backup=getattr(args, "backup", True), + purge_queue=getattr(args, "purge_queue", False), + quarantine_orphans=getattr(args, "quarantine_orphans", False), + dry_run=getattr(args, "dry_run", False), + assume_yes=getattr(args, "yes", False), + ) + return + if getattr(args, "mode", "legacy") == "max-seq-id": from .repair import repair_max_seq_id @@ -809,6 +832,23 @@ def cmd_repair(args): ) return + if getattr(args, "mode", "legacy") == "reconcile": + if not getattr(args, "segment", None): + print(" --mode reconcile requires --segment ") + return + from .repair import reconcile_orphan_sql_rows + + reconcile_orphan_sql_rows( + palace_path, + segment=args.segment, + metadata_segment=getattr(args, "metadata_segment", None), + max_elements=getattr(args, "max_elements", None), + backup=getattr(args, "backup", True), + dry_run=getattr(args, "dry_run", False), + assume_yes=getattr(args, "yes", False), + ) + return + if getattr(args, "mode", "legacy") == "from-sqlite": from .migrate import confirm_destructive_action from .repair import RebuildPartialError, rebuild_from_sqlite @@ -1477,8 +1517,8 @@ def main(): p_repair = sub.add_parser( "repair", help=( - "Rebuild palace vector index (legacy mode) or un-poison max_seq_id rows " - "(--mode max-seq-id)" + "Rebuild palace vector index (legacy mode), rebuild a single HNSW segment " + "(--mode hnsw, issue #1046), or un-poison max_seq_id rows (--mode max-seq-id)" ), ) p_repair.add_argument( @@ -1496,14 +1536,16 @@ def main(): ) p_repair.add_argument( "--mode", - choices=["legacy", "max-seq-id", "from-sqlite"], + choices=["legacy", "hnsw", "max-seq-id", "from-sqlite", "reconcile"], default="legacy", help=( "legacy: full-palace rebuild via the chromadb client (default). " + "hnsw: rebuild one segment from data_level0.bin (issue #1046). " "max-seq-id: un-poison max_seq_id rows corrupted by the legacy 0.6.x shim. " "from-sqlite: rebuild by reading rows directly from chroma.sqlite3, " "bypassing the chromadb client. Use when legacy mode bails because the " - "chromadb client cannot open the collection." + "chromadb client cannot open the collection. " + "reconcile: re-embed SQL-only orphan rows into the HNSW segment." ), ) p_repair.add_argument( @@ -1526,7 +1568,10 @@ def main(): p_repair.add_argument( "--segment", default=None, - help="Segment UUID filter for --mode max-seq-id (repairs only that segment).", + help=( + "Segment UUID. For --mode max-seq-id: repair only that segment (optional filter). " + "For --mode hnsw or reconcile: required, points to //." + ), ) p_repair.add_argument( "--from-sidecar", @@ -1536,16 +1581,50 @@ def main(): "clean values are copied from its max_seq_id table verbatim." ), ) + p_repair.add_argument( + "--metadata-segment", + default=None, + help=( + "METADATA segment UUID for --mode reconcile " + "(auto-detected from sibling-segment lookup when omitted)" + ), + ) + p_repair.add_argument( + "--max-elements", + type=int, + default=None, + help="HNSW max_elements for new index (--mode hnsw only; default: max(count*1.3, 200_000))", + ) p_repair.add_argument( "--backup", action=argparse.BooleanOptionalAction, default=True, - help="Back up SQLite before mutation (default: on)", + help=( + "Back up before mutation (default: on). " + "--mode hnsw: SQLite + data_level0.bin + pickle. " + "--mode max-seq-id: SQLite only." + ), + ) + p_repair.add_argument( + "--purge-queue", + action="store_true", + help=( + "(--mode hnsw only) Clear the embeddings_queue rows for this segment's " + "collection after rebuild" + ), + ) + p_repair.add_argument( + "--quarantine-orphans", + action="store_true", + help=( + "(--mode hnsw only) Append dropped UUIDs + orphan HNSW labels to " + "quarantined_orphans.json" + ), ) p_repair.add_argument( "--dry-run", action="store_true", - help="Print detected poisoned rows and exit without mutation (--mode max-seq-id only)", + help="Print rebuild/repair report and exit without mutation (--mode hnsw / max-seq-id)", ) # repair-status — read-only HNSW capacity health check (#1222) diff --git a/mempalace/repair.py b/mempalace/repair.py index 7a4a28cd19..0d37e7d87c 100644 --- a/mempalace/repair.py +++ b/mempalace/repair.py @@ -6,13 +6,19 @@ add() calls with the same ID), link_lists.bin can grow unbounded — terabytes on large palaces — eventually causing segfaults. -This module provides four operations: - - status — compare sqlite vs HNSW element counts (read-only health check) - scan — find every corrupt/unfetchable ID in the palace - prune — delete only the corrupt IDs (surgical) - rebuild — extract all drawers, delete the collection, recreate with - correct HNSW settings, and upsert everything back +This module provides several operations: + + status — compare sqlite vs HNSW element counts (read-only health check) + scan — find every corrupt/unfetchable ID in the palace + prune — delete only the corrupt IDs (surgical) + rebuild — extract all drawers, delete the collection, recreate with + correct HNSW settings, and upsert everything back + hnsw-rebuild — segment-level HNSW rebuild from data_level0.bin + + index_metadata.pickle; avoids re-embedding, bounded memory, + atomic swap-aside with rollback. Productionises the + recovery path from the 2026-04-19 incident. Issue #1046. + max-seq-id — un-poison ``max_seq_id`` rows corrupted by the legacy 0.6.x + BLOB shim misreading chromadb 1.5.x's native format. The rebuild backs up ONLY chroma.sqlite3 (the source of truth), not the full palace directory — so it works even when link_lists.bin is bloated. @@ -22,20 +28,35 @@ python -m mempalace.repair scan [--wing X] python -m mempalace.repair prune --confirm python -m mempalace.repair rebuild + python -m mempalace.repair hnsw --segment [--dry-run] [--purge-queue] ... + python -m mempalace.repair max-seq-id [--segment ] [--from-sidecar ] Usage (from CLI): mempalace repair - mempalace repair-scan [--wing X] - mempalace repair-prune --confirm + mempalace repair --mode hnsw --segment + mempalace repair --mode max-seq-id [--segment ] [--from-sidecar ] + +The hnsw-rebuild path imports numpy and hnswlib lazily — they are not +core mempalace dependencies (per CONTRIBUTING.md). Install only when +running this rescue command. """ +from __future__ import annotations + import argparse +import gc +import json +import logging import os +import pickle import shutil import sqlite3 +import struct +import tempfile import time from collections import defaultdict from contextlib import closing +from dataclasses import dataclass from datetime import datetime import re from typing import Callable, Iterator, Optional @@ -44,6 +65,8 @@ from .backends.chroma import ChromaBackend, hnsw_capacity_status +logger = logging.getLogger(__name__) + COLLECTION_NAME = "mempalace_drawers" REPAIR_TEMP_COLLECTION = f"{COLLECTION_NAME}__repair_tmp" @@ -89,6 +112,8 @@ def _recoverable_collections() -> tuple[str, ...]: # up at call time. RECOVERABLE_COLLECTIONS = (COLLECTION_NAME, CLOSETS_COLLECTION_NAME) +_FLOAT32_SIZE = 4 + def _get_palace_path(): """Resolve palace path from config.""" @@ -1280,6 +1305,8 @@ def status(palace_path=None, collection_name: Optional[str] = None) -> dict: for label, info in (("drawers", drawers), ("closets", closets)): print(f"\n [{label}]") + if info.get("segment_id"): + print(f" segment id: {info['segment_id']}") if info["sqlite_count"] is None: print(" sqlite count: (unreadable)") else: @@ -1295,12 +1322,399 @@ def status(palace_path=None, collection_name: Optional[str] = None) -> dict: if info["message"]: print(f" note: {info['message']}") + diverged_segments = [ + (label, info["segment_id"]) + for label, info in (("drawers", drawers), ("closets", closets)) + if info["diverged"] and info.get("segment_id") + ] if drawers["diverged"] or closets["diverged"]: - print("\n Recommended: run `mempalace repair` to rebuild the index.") + print("\n Recommended next steps:") + if diverged_segments: + print( + " - Targeted segment rebuild (faster, no re-embed):" + " `mempalace repair --mode hnsw --segment `" + ) + for label, seg_id in diverged_segments: + print(f" {label}: {seg_id}") + print(" - Full-palace rebuild (re-embeds, slower): `mempalace repair`") print() return {"drawers": drawers, "closets": closets} +# --------------------------------------------------------------------------- +# hnsw-mode: segment-level rebuild (issue #1046) +# --------------------------------------------------------------------------- + + +class RebuildVerificationError(RuntimeError): + """Raised when the rebuilt index fails its self-query sanity check.""" + + +@dataclass +class _HnswHeader: + """Subset of the chromadb-wrapped hnswlib header we rely on. + + Byte layout (little-endian) of the first 100 bytes: + off 0: u32 format_version + off 4: u64 offset_level0 + off 12: u64 max_elements + off 20: u64 cur_count + off 28: u64 size_per_element + off 36: u64 label_offset + off 44: u64 offset_data + off 52: i32 maxlevel + off 56: u32 enterpoint_node + off 60: u64 maxM + off 68: u64 maxM0 + off 76: u64 M + off 84: f64 mult + off 92: u64 ef_construction + """ + + format_version: int + max_elements: int + cur_count: int + size_per_element: int + label_offset: int + offset_data: int + M: int + ef_construction: int + dim: int + + +def _parse_hnsw_header(data: bytes) -> _HnswHeader: + """Parse the 100-byte hnswlib header. + + Derives ``dim`` from ``size_per_element - offset_data - 8`` (trailing + 8 bytes are the u64 label), divided by 4 (float32). + """ + if len(data) < 100: + raise ValueError(f"HNSW header too short: {len(data)} bytes") + + format_version = struct.unpack_from("= {expected}") + + labels = np.empty(n, dtype=np.uint64) + vectors = np.empty((n, hdr.dim), dtype=np.float32) + vec_end = hdr.offset_data + hdr.dim * _FLOAT32_SIZE + for i in range(n): + slot = i * stride + vectors[i] = np.frombuffer(data[slot + hdr.offset_data : slot + vec_end], dtype=np.float32) + labels[i] = struct.unpack_from(" str: + """Look up ``hnsw:space`` for the collection that owns ``segment``. + + Falls back to ``"l2"`` (hnswlib's default) with a warning if absent — + matches ChromaDB's default when no metadata is recorded. + """ + db_path = os.path.join(palace_path, "chroma.sqlite3") + if not os.path.isfile(db_path): + logger.warning("No chroma.sqlite3 at %s — defaulting space to l2", palace_path) + return "l2" + + row = None + with sqlite3.connect(db_path) as conn: + cols = [r[1] for r in conn.execute("PRAGMA table_info(collection_metadata)").fetchall()] + value_col = "str_value" if "str_value" in cols else "string_value" + try: + row = conn.execute( + f""" + SELECT cm.{value_col} + FROM segments s + JOIN collection_metadata cm ON cm.collection_id = s.collection + WHERE s.id = ? AND cm.key = 'hnsw:space' + """, + (segment,), + ).fetchone() + except sqlite3.OperationalError: + row = None + + if row and row[0]: + return row[0] + logger.warning("No hnsw:space for segment %s — defaulting to l2 (ChromaDB default)", segment) + return "l2" + + +def _meta_get(meta, key): + """Read a field from index_metadata (0.6.x attr-object or 1.5.x dict).""" + return meta[key] if isinstance(meta, dict) else getattr(meta, key) + + +def _meta_set(meta, key, value): + """Write a field to index_metadata (0.6.x attr-object or 1.5.x dict).""" + if isinstance(meta, dict): + meta[key] = value + else: + setattr(meta, key, value) + + +def _reconcile_with_pickle(labels, pickle_path: str): + """Intersect HNSW labels with the pickle's ``label_to_id`` mapping. + + Returns ``(keep_mask, orphan_hnsw_labels, stale_pickle_ids, meta)`` + where ``meta`` has had its three mapping tables pruned to the healthy + set (caller persists it afterwards). + """ + import numpy as np + + with open(pickle_path, "rb") as f: + meta = pickle.load(f) + + label_to_id = _meta_get(meta, "label_to_id") + id_to_label = _meta_get(meta, "id_to_label") + id_to_seq_id = _meta_get(meta, "id_to_seq_id") + + mapped_labels = set(label_to_id.keys()) + hnsw_labels = set(int(x) for x in labels) + healthy = hnsw_labels & mapped_labels + orphan_hnsw = hnsw_labels - mapped_labels + stale_uids = [uid for lbl, uid in label_to_id.items() if lbl not in healthy] + dropped_uid_set = set(stale_uids) + + _meta_set( + meta, + "label_to_id", + {lbl: uid for lbl, uid in label_to_id.items() if lbl in healthy}, + ) + _meta_set( + meta, + "id_to_label", + {uid: lbl for uid, lbl in id_to_label.items() if uid not in dropped_uid_set}, + ) + _meta_set( + meta, + "id_to_seq_id", + {uid: sid for uid, sid in id_to_seq_id.items() if uid not in dropped_uid_set}, + ) + + keep_mask = np.fromiter((int(x) in healthy for x in labels), dtype=bool, count=len(labels)) + return keep_mask, sorted(orphan_hnsw), stale_uids, meta + + +def _compute_max_elements(count: int, override: Optional[int]) -> int: + """Pick the ``max_elements`` value for the new index. + + Default ``max(count * 1.3, 200_000)`` leaves headroom so the next + flush does not auto-resize (the very bug #2594 we are fixing). + """ + if override is not None: + if override < count: + raise ValueError( + f"--max-elements={override} is smaller than healthy vector count {count}" + ) + return int(override) + return max(int(count * 1.3), 200_000) + + +def _build_persistent_index( + vectors, + labels, + *, + space: str, + dim: int, + max_elements: int, + persistence_location: str, + M: int = 16, + ef_construction: int = 100, +): + """Build a persistent hnswlib index and write it to ``persistence_location``.""" + import hnswlib + + idx = hnswlib.Index(space=space, dim=dim) + idx.init_index( + max_elements=max_elements, + ef_construction=ef_construction, + M=M, + is_persistent_index=True, + persistence_location=persistence_location, + ) + idx.set_num_threads(1) + + n = len(labels) + chunk = 10_000 + for i in range(0, n, chunk): + j = min(i + chunk, n) + idx.add_items(vectors[i:j], labels[i:j], num_threads=1) + idx.persist_dirty() + return idx + + +def _self_query_verify(index, sample_vectors, sample_labels, k: int = 10) -> None: + """Verify the rebuilt index returns each sample within its own top-k neighbors. + + Looser than top-1 to tolerate near-duplicates — mined corpora regularly contain + drawers with byte-identical vectors (e.g. the same code snippet appearing in + multiple transcripts), so an exact top-1 check false-positives on corpora that + are actually indexed correctly. + + Bumps ``ef`` before querying: hnswlib's default (≈10) is too tight a search + beam for a ~500k-element index with ``M=16`` and can miss even a byte-identical + self-match because its neighborhood in the HNSW graph is sparse. ChromaDB sets + its own ``ef`` at query time, so this only affects the verify step. + """ + if len(sample_labels) == 0: + return + index.set_ef(max(200, k * 4)) + labels, _dists = index.knn_query(sample_vectors, k=k) + expected = [int(x) for x in sample_labels] + misses = [] + for i, exp in enumerate(expected): + row = [int(x) for x in labels[i]] + if exp not in row: + misses.append((exp, row)) + if misses: + raise RebuildVerificationError( + f"Self-query mismatch (top-{k}): {len(misses)}/{len(expected)} samples " + f"did not include their own label — first miss: expected {misses[0][0]}, " + f"got {misses[0][1]}" + ) + + +def _atomic_swap_segment(tmpdir: str, segment_dir: str) -> None: + """Rename-aside swap: move live out of the way, drop new in, rollback on failure.""" + stamp = datetime.now().strftime("%Y%m%d-%H%M%S") + stale = f"{segment_dir}.old-{stamp}" + os.rename(segment_dir, stale) + try: + os.replace(tmpdir, segment_dir) + except OSError: + try: + os.rename(stale, segment_dir) + except OSError: + logger.exception("Swap failed AND rollback failed. Live segment left at %s", stale) + raise + shutil.rmtree(stale, ignore_errors=True) + + +def _backup_segment(palace_path: str, segment: str, timestamp: str) -> str: + """Copy chroma.sqlite3 plus the small HNSW files (skip bloated link_lists.bin).""" + seg_dir = os.path.join(palace_path, segment) + backup_dir = os.path.join(palace_path, f"{segment}.hnsw-backup-{timestamp}") + os.makedirs(backup_dir, exist_ok=True) + + sqlite_path = os.path.join(palace_path, "chroma.sqlite3") + if os.path.isfile(sqlite_path): + shutil.copy2(sqlite_path, os.path.join(backup_dir, "chroma.sqlite3")) + + for fname in ("header.bin", "data_level0.bin", "index_metadata.pickle", "length.bin"): + src = os.path.join(seg_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(backup_dir, fname)) + return backup_dir + + +def _purge_segment_queue(palace_path: str, segment: str) -> int: + """Delete ``embeddings_queue`` rows for the collection that owns ``segment``. + + ``topic`` is ``persistent://default/default/`` — we look + up the collection UUID via ``segments.collection`` and match by pattern. + """ + db_path = os.path.join(palace_path, "chroma.sqlite3") + if not os.path.isfile(db_path): + return 0 + with sqlite3.connect(db_path) as conn: + row = conn.execute("SELECT collection FROM segments WHERE id = ?", (segment,)).fetchone() + if not row: + return 0 + collection_uuid = row[0] + cur = conn.execute( + "DELETE FROM embeddings_queue WHERE topic LIKE ?", + (f"%{collection_uuid}%",), + ) + deleted = cur.rowcount + conn.commit() + return int(deleted or 0) + + +def _quarantine_orphans(palace_path: str, stale_ids, orphan_labels) -> str: + """Append dropped UUIDs + orphan HNSW labels to a sidecar JSON file.""" + sidecar = os.path.join(palace_path, "quarantined_orphans.json") + entry = { + "timestamp": datetime.now().isoformat(timespec="seconds"), + "stale_pickle_ids": list(stale_ids), + "orphan_hnsw_labels": [int(x) for x in orphan_labels], + } + history: list = [] + if os.path.isfile(sidecar): + try: + with open(sidecar) as f: + history = json.load(f) + if not isinstance(history, list): + history = [history] + except Exception: + history = [] + history.append(entry) + with open(sidecar, "w") as f: + json.dump(history, f, indent=2) + return sidecar + + # --------------------------------------------------------------------------- # max-seq-id mode: un-poison max_seq_id rows corrupted by the old shim # --------------------------------------------------------------------------- @@ -1313,7 +1727,6 @@ def _close_chroma_handles(palace_path: str, backend: "ChromaBackend | None" = No releases the handles it was already using. Otherwise fall back to a transient backend instance for the max-seq-id repair path. """ - import gc try: closer = backend if backend is not None else ChromaBackend() @@ -1563,14 +1976,575 @@ def repair_max_seq_id( return result +# --------------------------------------------------------------------------- +# hnsw-mode driver: rebuild a single segment from data_level0.bin (issue #1046) +# --------------------------------------------------------------------------- + + +def _peak_memory_mb() -> float: + """Return the process peak RSS in MB (mac returns bytes, linux kilobytes).""" + try: + import resource + import sys + + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform == "darwin": + return peak / (1024 * 1024) + return peak / 1024 + except Exception: + return 0.0 + + +def rebuild_hnsw_segment( + palace_path: str, + *, + segment: str, + max_elements: Optional[int] = None, + backup: bool = True, + purge_queue: bool = False, + quarantine_orphans: bool = False, + dry_run: bool = False, + assume_yes: bool = False, +) -> dict: + """Rebuild a single HNSW segment from on-disk ``data_level0.bin`` + pickle. + + Avoids re-embedding by reading vectors straight out of the persistent + HNSW data file. Atomic swap-aside with rollback keeps the live palace + untouched on any failure. Issue #1046. + + On successful completion the palace is healthy on disk; if the running + MCP server still has ``_vector_disabled`` set from a prior #1222 capacity + probe, calling ``mempalace_reconnect`` will refresh the probe and clear + the flag — the runtime check is the authoritative source of truth and + re-runs at every reconnect. + """ + from .migrate import confirm_destructive_action, contains_palace_database + + palace_path = os.path.abspath(os.path.expanduser(palace_path)) + seg_dir = os.path.join(palace_path, segment) + header_path = os.path.join(seg_dir, "header.bin") + data_path = os.path.join(seg_dir, "data_level0.bin") + pickle_path = os.path.join(seg_dir, "index_metadata.pickle") + + result: dict = { + "palace_path": palace_path, + "segment": segment, + "dry_run": dry_run, + "aborted": False, + } + + print(f"\n{'=' * 55}") + print(" MemPalace Repair — HNSW Segment Rebuild") + print(f"{'=' * 55}\n") + print(f" Palace: {palace_path}") + print(f" Segment: {segment}") + + if not os.path.isdir(palace_path): + print(f" No palace found at {palace_path}") + result["aborted"] = True + result["reason"] = "palace-missing" + return result + if not contains_palace_database(palace_path): + print(f" No palace database at {palace_path}") + result["aborted"] = True + result["reason"] = "db-missing" + return result + if not os.path.isdir(seg_dir): + print(f" Segment directory not found: {seg_dir}") + result["aborted"] = True + result["reason"] = "segment-missing" + return result + if not os.path.isfile(data_path): + print(f" data_level0.bin not found in {seg_dir}") + result["aborted"] = True + result["reason"] = "data-missing" + return result + + try: + import numpy # noqa: F401 + import hnswlib # noqa: F401 + except ImportError as e: + print(f" Required dependency missing: {e}") + print(" Install with: pip install numpy chroma-hnswlib") + result["aborted"] = True + result["reason"] = "deps-missing" + return result + + header_src = header_path if os.path.isfile(header_path) else data_path + with open(header_src, "rb") as f: + header_bytes = f.read(100) + hdr = _parse_hnsw_header(header_bytes) + print( + f" Header: dim={hdr.dim}, cur_count={hdr.cur_count:,}, " + f"max_elements={hdr.max_elements:,}, size_per_element={hdr.size_per_element}" + ) + + space = _detect_space(palace_path, segment) + print(f" Space: {space}") + + with open(data_path, "rb") as f: + data_bytes = f.read() + labels, vectors = _extract_vectors(data_bytes, hdr) + del data_bytes + raw_n = len(labels) + labels, vectors = _sanitize_vectors(labels, vectors) + sanitized_n = len(labels) + + orphan_hnsw: list = [] + stale_uids: list = [] + meta = None + if os.path.isfile(pickle_path): + keep_mask, orphan_hnsw, stale_uids, meta = _reconcile_with_pickle(labels, pickle_path) + labels = labels[keep_mask] + vectors = vectors[keep_mask] + else: + logger.warning("No index_metadata.pickle for segment %s — skipping reconcile", segment) + + healthy_n = len(labels) + new_max = _compute_max_elements(healthy_n, max_elements) + + data_bytes_size = os.path.getsize(data_path) + link_lists_path = os.path.join(seg_dir, "link_lists.bin") + link_lists_size = os.path.getsize(link_lists_path) if os.path.isfile(link_lists_path) else 0 + + print() + print(" Report") + print(f" raw labels {raw_n:>10,}") + print(f" after dedup/zeros {sanitized_n:>10,}") + print(f" healthy (in pickle) {healthy_n:>10,}") + print(f" orphan HNSW labels {len(orphan_hnsw):>10,}") + print(f" stale pickle ids {len(stale_uids):>10,}") + print(f" new max_elements {new_max:>10,}") + print(f" data_level0.bin {data_bytes_size:>10,} bytes") + print(f" link_lists.bin {link_lists_size:>10,} bytes (will be rebuilt)") + + result.update( + { + "raw_labels": raw_n, + "sanitized_labels": sanitized_n, + "healthy_labels": healthy_n, + "orphan_hnsw_labels": len(orphan_hnsw), + "stale_pickle_ids": len(stale_uids), + "max_elements": new_max, + "data_bytes": data_bytes_size, + "link_lists_bytes": link_lists_size, + "space": space, + "dim": hdr.dim, + } + ) + + if dry_run: + print("\n DRY RUN — no files modified.\n" + "=" * 55 + "\n") + return result + + if healthy_n == 0: + print(" No healthy labels to rebuild — aborting.") + result["aborted"] = True + result["reason"] = "no-healthy-labels" + return result + + if not confirm_destructive_action("Rebuild HNSW segment", palace_path, assume_yes=assume_yes): + result["aborted"] = True + result["reason"] = "user-aborted" + return result + + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + + backup_dir: Optional[str] = None + if backup: + backup_dir = _backup_segment(palace_path, segment, timestamp) + print(f" Backup: {backup_dir}") + + _close_chroma_handles(palace_path) + + tmpdir = tempfile.mkdtemp(prefix="mempalace_hnsw_", dir=palace_path) + t0 = time.time() + try: + idx = _build_persistent_index( + vectors, + labels, + space=space, + dim=hdr.dim, + max_elements=new_max, + persistence_location=tmpdir, + M=hdr.M or 16, + ef_construction=hdr.ef_construction or 100, + ) + except Exception: + shutil.rmtree(tmpdir, ignore_errors=True) + raise + build_seconds = time.time() - t0 + print(f" Built: {healthy_n:,} vectors in {build_seconds:.1f}s") + + try: + sample_n = min(3, healthy_n) + _self_query_verify(idx, vectors[:sample_n], labels[:sample_n], k=min(10, healthy_n)) + except RebuildVerificationError: + shutil.rmtree(tmpdir, ignore_errors=True) + raise + finally: + del idx + gc.collect() + + if meta is not None: + _meta_set(meta, "total_elements_added", healthy_n) + with open(os.path.join(tmpdir, "index_metadata.pickle"), "wb") as f: + pickle.dump(meta, f, protocol=pickle.HIGHEST_PROTOCOL) + + _atomic_swap_segment(tmpdir, seg_dir) + + if purge_queue: + deleted = _purge_segment_queue(palace_path, segment) + print(f" Queue: purged {deleted:,} embeddings_queue rows") + result["queue_rows_purged"] = deleted + if quarantine_orphans and (stale_uids or orphan_hnsw): + sidecar = _quarantine_orphans(palace_path, stale_uids, orphan_hnsw) + print(f" Orphans: appended to {sidecar}") + result["orphan_sidecar"] = sidecar + + peak_mb = _peak_memory_mb() + print(f"\n Rebuild complete in {build_seconds:.1f}s (peak RSS ≈ {peak_mb:.0f} MB)") + print(f" Backup: {backup_dir or '(skipped)'}") + print("\n If the MCP server is currently running with vector_disabled set") + print(" (e.g. after a #1222 capacity-divergence detection), call the") + print(" `mempalace_reconnect` tool to refresh the capacity probe and") + print(" restore vector search.") + print(f"\n{'=' * 55}\n") + + result.update({"build_seconds": build_seconds, "peak_rss_mb": peak_mb, "backup": backup_dir}) + return result + + +# --------------------------------------------------------------------------- +# reconcile mode: re-embed SQL-only rows that never landed an HNSW label +# --------------------------------------------------------------------------- + + +def _resolve_metadata_segment(db_path: str, vector_segment: str) -> Optional[str]: + """Find the METADATA segment that shares a collection with ``vector_segment``.""" + if not os.path.isfile(db_path): + return None + with sqlite3.connect(db_path) as conn: + row = conn.execute( + "SELECT collection FROM segments WHERE id = ?", (vector_segment,) + ).fetchone() + if not row: + return None + rows = conn.execute( + "SELECT id FROM segments WHERE collection = ? AND scope = 'METADATA'", + (row[0],), + ).fetchall() + if len(rows) != 1: + return None + return str(rows[0][0]) + + +def _fetch_sql_only_docs(db_path: str, metadata_segment: str, hnsw_uuids: set) -> list: + """Return ``[(embedding_id, document), ...]`` for SQL rows missing from HNSW. + + Stable order by ``embedding_id`` so reconciles are reproducible. + """ + with sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) as conn: + rows = conn.execute( + """ + SELECT e.embedding_id, em.string_value + FROM embeddings e + JOIN embedding_metadata em ON e.id = em.id + WHERE e.segment_id = ? AND em.key = 'chroma:document' + ORDER BY e.embedding_id + """, + (metadata_segment,), + ).fetchall() + return [(uid, doc) for uid, doc in rows if uid not in hnsw_uuids] + + +def _embed_in_batches(ef, docs, *, dim: int, batch: int = 64): + """Encode ``docs`` with ``ef`` in fixed-size batches; return ``np.ndarray``.""" + import numpy as np + + out = np.empty((len(docs), dim), dtype=np.float32) + for i in range(0, len(docs), batch): + j = min(i + batch, len(docs)) + out[i:j] = np.asarray(ef(docs[i:j]), dtype=np.float32) + return out + + +def reconcile_orphan_sql_rows( + palace_path: str, + *, + segment: str, + metadata_segment: Optional[str] = None, + max_elements: Optional[int] = None, + backup: bool = True, + dry_run: bool = False, + assume_yes: bool = False, + embedding_function=None, +) -> dict: + """Re-embed SQL-only embeddings into the HNSW vector ``segment``. + + Some chromadb crash modes (e.g. issue #6979) commit ``embeddings``/ + ``embedding_metadata`` rows transactionally to SQL but lose the + corresponding HNSW additions, leaving a subset of drawers visible to + metadata queries but unreachable to vector search. This mode finds + those orphans, embeds their ``chroma:document`` payloads with the + palace's configured embedding function, and writes a fresh persistent + index containing both the existing labels (extracted from + ``data_level0.bin``) and freshly-allocated labels for the SQL-only + rows. Atomic swap with rollback on failure — same safety profile as + ``--mode hnsw``. + + ``metadata_segment`` is auto-detected from the sibling METADATA + segment in the ``segments`` table when omitted. ``embedding_function`` + is injectable for tests; production callers should leave it as + ``None`` so the palace's configured EF is resolved. + """ + from .migrate import confirm_destructive_action, contains_palace_database + + palace_path = os.path.abspath(os.path.expanduser(palace_path)) + seg_dir = os.path.join(palace_path, segment) + db_path = os.path.join(palace_path, "chroma.sqlite3") + pickle_path = os.path.join(seg_dir, "index_metadata.pickle") + data_path = os.path.join(seg_dir, "data_level0.bin") + header_path = os.path.join(seg_dir, "header.bin") + + result: dict = { + "palace_path": palace_path, + "segment": segment, + "dry_run": dry_run, + "aborted": False, + } + + print(f"\n{'=' * 55}") + print(" MemPalace Repair — SQL/HNSW Reconcile") + print(f"{'=' * 55}\n") + print(f" Palace: {palace_path}") + print(f" Segment: {segment}") + + if not os.path.isdir(palace_path): + print(f" No palace found at {palace_path}") + result["aborted"] = True + result["reason"] = "palace-missing" + return result + if not contains_palace_database(palace_path): + print(f" No palace database at {palace_path}") + result["aborted"] = True + result["reason"] = "db-missing" + return result + if not os.path.isdir(seg_dir): + print(f" Segment directory not found: {seg_dir}") + result["aborted"] = True + result["reason"] = "segment-missing" + return result + if not os.path.isfile(pickle_path): + print(f" index_metadata.pickle not found in {seg_dir}") + result["aborted"] = True + result["reason"] = "pickle-missing" + return result + + try: + import hnswlib # noqa: F401 + import numpy # noqa: F401 + except ImportError as e: + print(f" Required dependency missing: {e}") + result["aborted"] = True + result["reason"] = "deps-missing" + return result + + if metadata_segment is None: + metadata_segment = _resolve_metadata_segment(db_path, segment) + if not metadata_segment: + print(" Could not resolve sibling METADATA segment — pass --metadata-segment") + result["aborted"] = True + result["reason"] = "metadata-segment-unresolved" + return result + print(f" Metadata segment: {metadata_segment}") + + with open(pickle_path, "rb") as f: + meta = pickle.load(f) + id_to_label = dict(_meta_get(meta, "id_to_label") or {}) + label_to_id = dict(_meta_get(meta, "label_to_id") or {}) + pickle_total = int(_meta_get(meta, "total_elements_added") or 0) + if len(id_to_label) != pickle_total: + print( + f" Pickle inconsistent: id_to_label={len(id_to_label)} " + f"vs total_elements_added={pickle_total}. Run --mode hnsw first." + ) + result["aborted"] = True + result["reason"] = "pickle-inconsistent" + return result + + src = header_path if os.path.isfile(header_path) else data_path + with open(src, "rb") as f: + hdr = _parse_hnsw_header(f.read(100)) + space = _detect_space(palace_path, segment) + + missing = _fetch_sql_only_docs(db_path, metadata_segment, set(id_to_label.keys())) + print() + print(" Report") + print(f" existing labels {len(id_to_label):>10,}") + print(f" sql-only orphans {len(missing):>10,}") + print(f" space {space}") + print(f" dim {hdr.dim}") + + result.update( + { + "existing_labels": len(id_to_label), + "sql_only_orphans": len(missing), + "space": space, + "dim": hdr.dim, + } + ) + + if not missing: + print(" Nothing to reconcile.") + print(f"\n{'=' * 55}\n") + return result + + if dry_run: + print("\n DRY RUN — no embedding, no HNSW write, no swap.\n" + "=" * 55 + "\n") + return result + + if not confirm_destructive_action("Reconcile HNSW segment", palace_path, assume_yes=assume_yes): + result["aborted"] = True + result["reason"] = "user-aborted" + return result + + if embedding_function is None: + from .embedding import get_embedding_function + + embedding_function = get_embedding_function() + + import numpy as np + + with open(data_path, "rb") as f: + data_bytes = f.read() + raw_labels, raw_vectors = _extract_vectors(data_bytes, hdr) + del data_bytes + keep_set = set(id_to_label.values()) + keep_mask = np.array([int(x) in keep_set for x in raw_labels], dtype=bool) + existing_labels = raw_labels[keep_mask] + existing_vectors = raw_vectors[keep_mask] + if len(existing_labels) != len(id_to_label): + print(f" WARN: extracted healthy ({len(existing_labels)}) != pickle ({len(id_to_label)}).") + + docs = [d for _uid, d in missing] + uuids = [u for u, _d in missing] + new_vectors = _embed_in_batches(embedding_function, docs, dim=hdr.dim) + max_label = max(id_to_label.values()) if id_to_label else 0 + new_labels = np.arange(max_label + 1, max_label + 1 + len(missing), dtype=np.int64) + new_total = len(id_to_label) + len(missing) + new_max = _compute_max_elements(new_total, max_elements) + if int(new_labels[-1]) >= new_max: + new_max = max(new_max, int(new_labels[-1]) + 1) + + all_vectors = np.concatenate([existing_vectors, new_vectors]) + all_labels = np.concatenate([existing_labels.astype(np.int64), new_labels]) + + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + backup_dir: Optional[str] = None + if backup: + backup_dir = _backup_segment(palace_path, segment, timestamp) + print(f" Backup: {backup_dir}") + + _close_chroma_handles(palace_path) + + tmpdir = tempfile.mkdtemp(prefix="mempalace_reconcile_", dir=palace_path) + t0 = time.time() + try: + idx = _build_persistent_index( + all_vectors, + all_labels, + space=space, + dim=hdr.dim, + max_elements=new_max, + persistence_location=tmpdir, + M=hdr.M or 16, + ef_construction=hdr.ef_construction or 100, + ) + except Exception: + shutil.rmtree(tmpdir, ignore_errors=True) + raise + build_seconds = time.time() - t0 + + try: + sample_n = min(3, len(missing)) + existing_n = min(3, len(existing_vectors)) + sv = np.concatenate([existing_vectors[:existing_n], new_vectors[:sample_n]]) + sl = np.concatenate( + [ + existing_labels[:existing_n].astype(np.int64), + new_labels[:sample_n], + ] + ) + _self_query_verify(idx, sv, sl, k=min(10, len(all_labels))) + except Exception: + shutil.rmtree(tmpdir, ignore_errors=True) + raise + finally: + del idx + gc.collect() + + new_id_to_label = dict(id_to_label) + new_label_to_id = dict(label_to_id) + for uid, lbl in zip(uuids, new_labels): + new_id_to_label[uid] = int(lbl) + new_label_to_id[int(lbl)] = uid + _meta_set(meta, "id_to_label", new_id_to_label) + _meta_set(meta, "label_to_id", new_label_to_id) + _meta_set(meta, "total_elements_added", new_total) + with open(os.path.join(tmpdir, "index_metadata.pickle"), "wb") as f: + pickle.dump(meta, f, protocol=pickle.HIGHEST_PROTOCOL) + + _atomic_swap_segment(tmpdir, seg_dir) + + peak_mb = _peak_memory_mb() + print( + f"\n Reconcile complete: {len(missing):,} new labels appended in " + f"{build_seconds:.1f}s (peak RSS ≈ {peak_mb:.0f} MB)" + ) + print(f" Backup: {backup_dir or '(skipped)'}") + print(f"\n{'=' * 55}\n") + + result.update( + { + "new_labels": len(missing), + "total_elements_added": new_total, + "build_seconds": build_seconds, + "peak_rss_mb": peak_mb, + "backup": backup_dir, + } + ) + return result + + if __name__ == "__main__": p = argparse.ArgumentParser(description="MemPalace repair tools") - p.add_argument("command", choices=["status", "scan", "prune", "rebuild"]) p.add_argument("--palace", default=None, help="Palace directory path") - p.add_argument("--wing", default=None, help="Scan only this wing") - p.add_argument("--confirm", action="store_true", help="Actually delete corrupt IDs") - args = p.parse_args() + sub = p.add_subparsers(dest="command", required=True) + sub.add_parser("status", help="Read-only sqlite-vs-HNSW capacity health check") + p_scan = sub.add_parser("scan") + p_scan.add_argument("--wing", default=None) + p_prune = sub.add_parser("prune") + p_prune.add_argument("--confirm", action="store_true") + sub.add_parser("rebuild") + p_hnsw = sub.add_parser("hnsw", help="Single-segment HNSW rebuild (issue #1046)") + p_hnsw.add_argument("--segment", required=True) + p_hnsw.add_argument("--max-elements", type=int, default=None) + p_hnsw.add_argument("--backup", action=argparse.BooleanOptionalAction, default=True) + p_hnsw.add_argument("--purge-queue", action="store_true") + p_hnsw.add_argument("--quarantine-orphans", action="store_true") + p_hnsw.add_argument("--dry-run", action="store_true") + p_hnsw.add_argument("--yes", action="store_true") + p_msi = sub.add_parser( + "max-seq-id", help="Un-poison max_seq_id rows (legacy 0.6.x shim damage)" + ) + p_msi.add_argument("--segment", default=None) + p_msi.add_argument("--from-sidecar", default=None) + p_msi.add_argument("--backup", action=argparse.BooleanOptionalAction, default=True) + p_msi.add_argument("--dry-run", action="store_true") + p_msi.add_argument("--yes", action="store_true") + args = p.parse_args() path = os.path.expanduser(args.palace) if args.palace else None if args.command == "status": @@ -1581,3 +2555,23 @@ def repair_max_seq_id( prune_corrupt(palace_path=path, confirm=args.confirm) elif args.command == "rebuild": rebuild_index(palace_path=path) + elif args.command == "hnsw": + rebuild_hnsw_segment( + path or _get_palace_path(), + segment=args.segment, + max_elements=args.max_elements, + backup=args.backup, + purge_queue=args.purge_queue, + quarantine_orphans=args.quarantine_orphans, + dry_run=args.dry_run, + assume_yes=args.yes, + ) + elif args.command == "max-seq-id": + repair_max_seq_id( + path or _get_palace_path(), + segment=args.segment, + from_sidecar=args.from_sidecar, + backup=args.backup, + dry_run=args.dry_run, + assume_yes=args.yes, + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 0caf75c3ed..8e00e7b450 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1041,6 +1041,123 @@ def test_cmd_repair_aborts_without_confirmation(mock_config_cls, tmp_path, capsy mock_backend.create_collection.assert_not_called() +# ── cmd_repair --mode hnsw (issue #1046) ─────────────────────────────── + + +def _hnsw_args(tmp_path, **overrides): + defaults = dict( + palace=str(tmp_path), + mode="hnsw", + segment=None, + max_elements=None, + backup=True, + purge_queue=False, + quarantine_orphans=False, + dry_run=False, + yes=True, + ) + defaults.update(overrides) + return argparse.Namespace(**defaults) + + +def test_cmd_repair_hnsw_requires_segment(tmp_path, capsys): + args = _hnsw_args(tmp_path, segment=None) + with patch("mempalace.repair.rebuild_hnsw_segment") as mock_rebuild: + cmd_repair(args) + mock_rebuild.assert_not_called() + assert "--mode hnsw requires --segment" in capsys.readouterr().out + + +def test_cmd_repair_hnsw_dispatches_with_defaults(tmp_path): + args = _hnsw_args(tmp_path, segment="seg-abc") + with patch("mempalace.repair.rebuild_hnsw_segment") as mock_rebuild: + cmd_repair(args) + mock_rebuild.assert_called_once() + call = mock_rebuild.call_args + assert call.kwargs["segment"] == "seg-abc" + assert call.kwargs["backup"] is True + assert call.kwargs["dry_run"] is False + assert call.kwargs["purge_queue"] is False + assert call.kwargs["quarantine_orphans"] is False + assert call.kwargs["max_elements"] is None + assert call.kwargs["assume_yes"] is True + + +def test_cmd_repair_hnsw_forwards_all_flags(tmp_path): + args = _hnsw_args( + tmp_path, + segment="seg-xyz", + max_elements=1234, + backup=False, + purge_queue=True, + quarantine_orphans=True, + dry_run=True, + yes=False, + ) + with patch("mempalace.repair.rebuild_hnsw_segment") as mock_rebuild: + cmd_repair(args) + mock_rebuild.assert_called_once() + call = mock_rebuild.call_args + assert call.kwargs["segment"] == "seg-xyz" + assert call.kwargs["max_elements"] == 1234 + assert call.kwargs["backup"] is False + assert call.kwargs["purge_queue"] is True + assert call.kwargs["quarantine_orphans"] is True + assert call.kwargs["dry_run"] is True + assert call.kwargs["assume_yes"] is False + + +def test_cmd_repair_legacy_mode_does_not_invoke_hnsw(tmp_path, capsys): + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + (palace_dir / "chroma.sqlite3").write_text("db") + args = argparse.Namespace(palace=str(palace_dir), mode="legacy", yes=True) + mock_col = MagicMock() + mock_col.count.return_value = 0 + mock_backend = _mock_backend_for(col=mock_col) + # The sqlite integrity preflight added in the #1362/#1364 ordering fix + # aborts with sys.exit(1) on the bogus "db" content seeded above before + # any rebuild path runs. That still proves the legacy path didn't fan + # out to the hnsw rebuild helper, which is what this test guards. + with ( + patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend), + patch("mempalace.repair.rebuild_hnsw_segment") as mock_hnsw, + pytest.raises(SystemExit), + ): + cmd_repair(args) + mock_hnsw.assert_not_called() + + +def test_main_repair_hnsw_parses_flags(): + argv = [ + "mempalace", + "repair", + "--mode", + "hnsw", + "--segment", + "seg-abc", + "--max-elements", + "777", + "--no-backup", + "--purge-queue", + "--quarantine-orphans", + "--dry-run", + "--yes", + ] + with patch("sys.argv", argv), patch("mempalace.cli.cmd_repair") as mock_cmd: + main() + mock_cmd.assert_called_once() + ns = mock_cmd.call_args.args[0] + assert ns.mode == "hnsw" + assert ns.segment == "seg-abc" + assert ns.max_elements == 777 + assert ns.backup is False + assert ns.purge_queue is True + assert ns.quarantine_orphans is True + assert ns.dry_run is True + assert ns.yes is True + + # ── cmd_compress ─────────────────────────────────────────────────────── diff --git a/tests/test_repair.py b/tests/test_repair.py index 981351ef8b..bfef85d7d7 100644 --- a/tests/test_repair.py +++ b/tests/test_repair.py @@ -1,7 +1,10 @@ """Tests for mempalace.repair — scan, prune, and rebuild HNSW index.""" +import json import os +import pickle import sqlite3 +import struct from contextlib import closing from unittest.mock import MagicMock, call, patch @@ -10,6 +13,133 @@ from mempalace import repair +# ── helpers: synthesize a legacy-format HNSW segment on disk ────────── + +_DIM = 8 +_SIZE_PER_ELEMENT = 132 + _DIM * 4 + 8 +_LABEL_OFFSET = 132 + _DIM * 4 +_OFFSET_DATA = 132 + + +def _pack_header(max_elements: int, cur_count: int) -> bytes: + hdr = bytearray(100) + struct.pack_into(" 0 + + result = repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True, purge_queue=True) + assert result["queue_rows_purged"] == before + + after = sqlite3.connect(db_path).execute("SELECT COUNT(*) FROM embeddings_queue").fetchone()[0] + assert after == 0 + + +def test_rebuild_hnsw_quarantine_orphans_writes_sidecar(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + sidecar = os.path.join(palace, "quarantined_orphans.json") + assert not os.path.exists(sidecar) + + repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True, quarantine_orphans=True) + + assert os.path.isfile(sidecar) + with open(sidecar) as f: + data = json.load(f) + assert isinstance(data, list) and len(data) == 1 + assert "uid-stale" in data[0]["stale_pickle_ids"] + + +def test_rebuild_hnsw_max_elements_override(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + seg_dir = os.path.join(palace, segment) + + result = repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True, max_elements=500) + assert result["max_elements"] == 500 + + with open(os.path.join(seg_dir, "header.bin"), "rb") as f: + hdr = repair._parse_hnsw_header(f.read(100)) + assert hdr.max_elements == 500 + + +def test_rebuild_hnsw_max_elements_override_below_count(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + with pytest.raises(ValueError, match="smaller than healthy"): + repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True, max_elements=2) + + +def test_rebuild_hnsw_rollback_on_build_failure(synthetic_segment, monkeypatch): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + seg_dir = os.path.join(palace, segment) + pre_contents = sorted(os.listdir(seg_dir)) + pre_sizes = {name: os.path.getsize(os.path.join(seg_dir, name)) for name in pre_contents} + + def _boom(*args, **kwargs): + raise RuntimeError("synthetic build failure") + + monkeypatch.setattr(repair, "_build_persistent_index", _boom) + + with pytest.raises(RuntimeError, match="synthetic build failure"): + repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True) + + assert os.path.isdir(seg_dir), "live segment dir must survive a failed build" + assert sorted(os.listdir(seg_dir)) == pre_contents + for name in pre_contents: + assert os.path.getsize(os.path.join(seg_dir, name)) == pre_sizes[name], ( + f"{name} was modified despite rollback" + ) + # No stray .old-* dirs left around + assert not any(n.startswith(segment + ".old-") for n in os.listdir(palace)) + + +def test_rebuild_hnsw_no_backup_flag(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + + result = repair.rebuild_hnsw_segment(palace, segment=segment, assume_yes=True, backup=False) + assert result["backup"] is None + assert not any(n.startswith(segment + ".hnsw-backup-") for n in os.listdir(palace)) + + +def test_detect_space_fallback_when_missing(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + db_path = palace / "chroma.sqlite3" + conn = sqlite3.connect(str(db_path)) + conn.executescript( + """ + CREATE TABLE segments(id TEXT PRIMARY KEY, type TEXT, scope TEXT, collection TEXT); + CREATE TABLE collection_metadata(collection_id TEXT, key TEXT, str_value TEXT); + INSERT INTO segments VALUES ('seg-x', 'VECTOR', 'VECTOR', 'coll-x'); + """ + ) + conn.commit() + conn.close() + + assert repair._detect_space(str(palace), "seg-x") == "l2" + + +def test_detect_space_returns_configured_value(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + db_path = palace / "chroma.sqlite3" + conn = sqlite3.connect(str(db_path)) + conn.executescript( + """ + CREATE TABLE segments(id TEXT PRIMARY KEY, type TEXT, scope TEXT, collection TEXT); + CREATE TABLE collection_metadata(collection_id TEXT, key TEXT, str_value TEXT); + INSERT INTO segments VALUES ('seg-x', 'VECTOR', 'VECTOR', 'coll-x'); + INSERT INTO collection_metadata VALUES ('coll-x', 'hnsw:space', 'ip'); + """ + ) + conn.commit() + conn.close() + + assert repair._detect_space(str(palace), "seg-x") == "ip" + + +def test_parse_hnsw_header_round_trip(): + header = _pack_header(max_elements=1000, cur_count=42) + hdr = repair._parse_hnsw_header(header) + assert hdr.max_elements == 1000 + assert hdr.cur_count == 42 + assert hdr.dim == _DIM + assert hdr.size_per_element == _SIZE_PER_ELEMENT + + +def test_parse_hnsw_header_too_short(): + with pytest.raises(ValueError, match="too short"): + repair._parse_hnsw_header(b"\x00" * 20) + + +def test_extract_vectors_accepts_cur_count_sized_file(): + np = pytest.importorskip("numpy") + cur_count, max_elements = 5, 100 + header = _pack_header(max_elements=max_elements, cur_count=cur_count) + hdr = repair._parse_hnsw_header(header) + + data = bytearray(cur_count * _SIZE_PER_ELEMENT) + data[:100] = header + vectors = np.arange(cur_count * _DIM, dtype=np.float32).reshape(cur_count, _DIM) + for i in range(cur_count): + slot = i * _SIZE_PER_ELEMENT + data[slot + _OFFSET_DATA : slot + _OFFSET_DATA + _DIM * 4] = vectors[i].tobytes() + struct.pack_into("` recovery command alongside + the legacy full-rebuild option, with the segment UUID inline. + """ + palace = tmp_path / "palace" + palace.mkdir() + (palace / "chroma.sqlite3").write_text("db") + seg_uuid = "deadbeef-1111-2222-3333-444455556666" + + def _fake_status(_palace, collection): + if collection == "mempalace_drawers": + return { + "segment_id": seg_uuid, + "sqlite_count": 200_000, + "hnsw_count": 16_384, + "divergence": 183_616, + "diverged": True, + "status": "diverged", + "message": "HNSW frozen at stale max_elements", + } + return { + "segment_id": None, + "sqlite_count": 0, + "hnsw_count": None, + "divergence": None, + "diverged": False, + "status": "ok", + "message": "", + } + + monkeypatch.setattr(repair, "hnsw_capacity_status", _fake_status) + monkeypatch.setattr(repair, "sqlite_drawer_count", lambda *a, **k: 200_000) + result = repair.status(palace_path=str(palace)) + out = capsys.readouterr().out + + assert result["drawers"]["diverged"] is True + assert "--mode hnsw --segment" in out + assert seg_uuid in out + assert "mempalace repair" in out # legacy full-rebuild path also surfaced + + +# ── reconcile_orphan_sql_rows ───────────────────────────────────────── + + +@pytest.fixture +def reconcile_palace(tmp_path): + """Build a synthetic palace with consistent pickle (no stale entries).""" + pytest.importorskip("numpy") + pytest.importorskip("hnswlib") + palace = tmp_path / "palace" + segment, coll, healthy_uids, vectors = _seed_hnsw_segment(str(palace), extra_pickle_ids=()) + return { + "palace": str(palace), + "segment": segment, + "collection": coll, + "uids": healthy_uids, + "vectors": vectors, + } + + +def _add_metadata_segment(palace_path, coll_uuid, healthy_uids, orphan_uids): + """Augment a ``_seed_hnsw_segment`` palace with a sibling METADATA segment. + + Adds the ``embeddings`` and ``embedding_metadata`` tables chromadb would + normally create, with one row per UID — both healthy (already in HNSW) + and orphan (SQL-only). + """ + metadata_segment = "11111111-2222-3333-4444-555566667777" + db_path = os.path.join(palace_path, "chroma.sqlite3") + conn = sqlite3.connect(db_path) + conn.executescript( + """ + CREATE TABLE embeddings( + id INTEGER PRIMARY KEY, + segment_id TEXT NOT NULL, + embedding_id TEXT NOT NULL, + seq_id BLOB NOT NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE (segment_id, embedding_id) + ); + CREATE TABLE embedding_metadata( + id INTEGER REFERENCES embeddings(id), + key TEXT NOT NULL, + string_value TEXT, + int_value INTEGER, + float_value REAL, + bool_value INTEGER, + PRIMARY KEY (id, key) + ); + """ + ) + conn.execute( + "INSERT INTO segments VALUES (?, 'urn:chroma:segment/metadata/sqlite', 'METADATA', ?)", + (metadata_segment, coll_uuid), + ) + rowid = 0 + for uid in list(healthy_uids) + list(orphan_uids): + rowid += 1 + conn.execute( + "INSERT INTO embeddings(id, segment_id, embedding_id, seq_id) VALUES (?, ?, ?, ?)", + (rowid, metadata_segment, uid, b"\x00" * 6), + ) + conn.execute( + "INSERT INTO embedding_metadata(id, key, string_value) VALUES (?, ?, ?)", + (rowid, "chroma:document", f"document text for {uid}"), + ) + conn.commit() + conn.close() + return metadata_segment + + +def _make_recording_ef(dim): + """Return a deterministic embedding function that records calls.""" + import numpy as np + + captured = [] + + def ef(docs): + captured.append(list(docs)) + out = np.zeros((len(docs), dim), dtype=np.float32) + for i, d in enumerate(docs): + out[i, 0] = float(abs(hash(d)) % 1000) / 1000.0 + 0.001 + out[i, 1] = float(len(d) % 100) / 100.0 + 0.001 + norms = np.linalg.norm(out, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + return out / norms + + return ef, captured + + +def test_reconcile_resolves_metadata_segment(tmp_path): + db_path = str(tmp_path / "chroma.sqlite3") + conn = sqlite3.connect(db_path) + conn.executescript( + """ + CREATE TABLE segments(id TEXT PRIMARY KEY, type TEXT, scope TEXT, collection TEXT); + INSERT INTO segments VALUES ('vec-1', 't', 'VECTOR', 'coll-1'); + INSERT INTO segments VALUES ('meta-1', 't', 'METADATA', 'coll-1'); + INSERT INTO segments VALUES ('vec-2', 't', 'VECTOR', 'coll-2'); + """ + ) + conn.commit() + conn.close() + assert repair._resolve_metadata_segment(db_path, "vec-1") == "meta-1" + # vec-2 has no METADATA sibling — returns None + assert repair._resolve_metadata_segment(db_path, "vec-2") is None + + +def test_reconcile_dry_run_reports_orphans(reconcile_palace): + palace = reconcile_palace["palace"] + segment = reconcile_palace["segment"] + coll = reconcile_palace["collection"] + healthy = reconcile_palace["uids"] + orphans = ["uid-orphan-a", "uid-orphan-b"] + metadata_segment = _add_metadata_segment(palace, coll, healthy, orphans) + seg_dir = os.path.join(palace, segment) + before = { + name: os.stat(os.path.join(seg_dir, name)).st_mtime_ns for name in os.listdir(seg_dir) + } + + result = repair.reconcile_orphan_sql_rows( + palace, + segment=segment, + metadata_segment=metadata_segment, + dry_run=True, + assume_yes=True, + ) + assert result["aborted"] is False + assert result["dry_run"] is True + assert result["existing_labels"] == 4 + assert result["sql_only_orphans"] == 2 + + after = {name: os.stat(os.path.join(seg_dir, name)).st_mtime_ns for name in os.listdir(seg_dir)} + assert before == after + + +def test_reconcile_appends_orphans_and_updates_pickle(reconcile_palace): + palace = reconcile_palace["palace"] + segment = reconcile_palace["segment"] + coll = reconcile_palace["collection"] + healthy = reconcile_palace["uids"] + orphans = ["uid-orphan-a", "uid-orphan-b", "uid-orphan-c"] + metadata_segment = _add_metadata_segment(palace, coll, healthy, orphans) + ef, captured = _make_recording_ef(_DIM) + + result = repair.reconcile_orphan_sql_rows( + palace, + segment=segment, + metadata_segment=metadata_segment, + max_elements=500, + assume_yes=True, + embedding_function=ef, + ) + assert result["aborted"] is False + assert result["new_labels"] == 3 + assert result["total_elements_added"] == 7 + + embedded = [d for batch in captured for d in batch] + for uid in orphans: + assert any(uid in d for d in embedded), f"orphan {uid} was not embedded" + + seg_dir = os.path.join(palace, segment) + with open(os.path.join(seg_dir, "index_metadata.pickle"), "rb") as f: + meta = pickle.load(f) + id_to_label = repair._meta_get(meta, "id_to_label") + label_to_id = repair._meta_get(meta, "label_to_id") + assert len(id_to_label) == 7 + assert len(label_to_id) == 7 + assert repair._meta_get(meta, "total_elements_added") == 7 + for uid in orphans: + assert uid in id_to_label + assert id_to_label[uid] in label_to_id + + +def test_reconcile_no_orphans_short_circuits(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + coll = synthetic_segment["collection"] + healthy = synthetic_segment["uids"] + metadata_segment = _add_metadata_segment(palace, coll, healthy, []) + + result = repair.reconcile_orphan_sql_rows( + palace, + segment=segment, + metadata_segment=metadata_segment, + assume_yes=True, + ) + assert result["aborted"] is False + assert result["sql_only_orphans"] == 0 + assert "new_labels" not in result + + +def test_reconcile_aborts_on_inconsistent_pickle(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + coll = synthetic_segment["collection"] + metadata_segment = _add_metadata_segment(palace, coll, synthetic_segment["uids"], ["x"]) + + pickle_path = os.path.join(palace, segment, "index_metadata.pickle") + with open(pickle_path, "rb") as f: + meta = pickle.load(f) + repair._meta_set(meta, "total_elements_added", 999) + with open(pickle_path, "wb") as f: + pickle.dump(meta, f) + + result = repair.reconcile_orphan_sql_rows( + palace, + segment=segment, + metadata_segment=metadata_segment, + assume_yes=True, + ) + assert result["aborted"] is True + assert result["reason"] == "pickle-inconsistent" + + +def test_reconcile_aborts_when_metadata_segment_unresolved(synthetic_segment): + palace = synthetic_segment["palace"] + segment = synthetic_segment["segment"] + # No METADATA segment added → auto-detect fails. + result = repair.reconcile_orphan_sql_rows( + palace, + segment=segment, + assume_yes=True, + ) + assert result["aborted"] is True + assert result["reason"] == "metadata-segment-unresolved"