From d80a1f43d1388c4fc7581b206b58839a48b514c3 Mon Sep 17 00:00:00 2001 From: Evgen197310 Date: Wed, 5 Aug 2026 16:03:39 +0300 Subject: [PATCH 1/4] fix(mcp): refuse vector writes into a diverged HNSW index instead of hanging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The #1222 capacity probe guards reads only: search and check_duplicate fall back to BM25 when the flushed HNSW segment lags sqlite. Writes went straight to chromadb — and an upsert into that index does not fail, it blocks inside chromadb's Rust upsert with no timeout of its own, for the life of the process, while this server holds the dispatch lock, the palace mine lock and the writer lease. Peer sessions drop to read-only behind a server whose handshake still answers, and nothing on the server side says why: the only trace of a 25-minute outage was the client's own "tool still running" ticks, and mempalace_status could not answer because the stuck call held the dispatch lock. This is the #1845 / #1908 symptom. Measured on a 3.7.0 palace whose flushed segment held 803 of 820 embeddings (py-spy: MainThread parked in chromadb/api/rust.py:_upsert, 64 idle tokio-rt-workers, no fcntl lock on chroma.sqlite3 — an internal wait, not lock contention): on copies of that palace, three of five freshly generated vectors blocked forever, the other two committed in 0.03s, and each vector reproduced its own verdict on retry. A write's fate depended on where its embedding landed in the damaged graph. After `mempalace repair rebuild-index`, all five committed in ~0.2s. Two guardrails, both Tier 1 of #1963: - _mcp_diverged_index_refusal: dispatch refuses the vector-write tools while the probe reports divergence, and puts the rebuild command in the error payload. Narrower than _MUTATING_TOOLS on purpose — the knowledge-graph and tunnel/hallway tools keep their own state and never reach the segment, so they must not even pay for the probe. Refusing is also the honest answer for the write that does *not* hang: chromadb acknowledges it into sqlite and the metadata segment, then leaves it out of HNSW, so the drawer reports success while being invisible to vector search. In the same incident 16 such drawers came out of one checkpoint that returned "completed successfully in 2s". The probe is pure sqlite + pickle, so the gate costs no chromadb interaction on the path it protects, and it re-runs per call so a repair in another process un-gates this server without a restart. - a write-stall watchdog: MEMPALACE_MCP_WRITE_STALL_WARN_SECS (default 60, 0 disables) warns on stderr — where the MCP host records it — from a thread the stuck call is not blocking. MEMPALACE_MCP_WRITE_STALL_EXIT_SECS (default 0, opt-in) lets an operator turn the wedge into a restartable failure: a server stuck inside chromadb never recovers, and exiting is what releases the palace locks its peers are queued behind. Exit code 75 (EX_TEMPFAIL) so it does not read as an orderly shutdown. Refs #1845, #1908, #1963 Co-Authored-By: Claude Opus 5 (1M context) --- mempalace/mcp_server.py | 215 ++++++++++++++++++++++++++- tests/test_vector_write_gate.py | 252 ++++++++++++++++++++++++++++++++ 2 files changed, 466 insertions(+), 1 deletion(-) create mode 100644 tests/test_vector_write_gate.py diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 7755c33..6388377 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -44,6 +44,7 @@ except (OSError, AttributeError): sys.stdout = sys.stderr import argparse # noqa: E402 (deferred until after stdio protection above) +import contextlib # noqa: E402 import json # noqa: E402 import logging # noqa: E402 import re # noqa: E402 @@ -405,6 +406,31 @@ _MUTATING_TOOLS = frozenset( } ) +# The subset of _MUTATING_TOOLS whose write path reaches the chroma vector +# segment. Deliberately narrower: the knowledge-graph and tunnel/hallway tools +# keep their own sqlite/JSON state and never touch HNSW, so an unusable vector +# index has no say over them. +# +# The distinction earns its keep because a write into a diverged HNSW segment +# does not fail — it blocks inside chromadb's Rust upsert with no timeout of its +# own, for the life of the process, while this server holds the palace mine lock +# and the writer lease. One stuck call becomes a palace-wide outage that a still +# healthy handshake hides. +_VECTOR_WRITE_TOOLS = frozenset( + { + "mempalace_add_drawer", + "mempalace_update_drawer", + "mempalace_delete_drawer", + "mempalace_delete_by_source", + "mempalace_diary_write", + "mempalace_checkpoint", + "mempalace_mine", + "mempalace_sync", + } +) + +_DIVERGED_INDEX_ERROR_CODE = -32004 + # Read-only mode (#1877) refuses a wider set than the peer-writer guard above. # # _MUTATING_TOOLS is the *palace-write* set: _mcp_peer_writer_refusal consults it @@ -4880,6 +4906,57 @@ def _mcp_read_only_refusal(req_id, tool_name: str): } +def _mcp_diverged_index_refusal(req_id, tool_name: str): + """Refuse vector writes while the HNSW segment is known to be diverged. + + The capacity probe (#1222) already routes *reads* around a diverged index: + ``search`` and ``check_duplicate`` fall back to BM25-only sqlite. Writes had + no such gate — they went straight to chromadb, where an upsert into that same + index can never come back. Observed on a 3.7.0 palace whose flushed segment + held 803 of 820 embeddings: three of five freshly generated vectors blocked + forever (25+ minutes, then killed), the other two committed in 0.03 s, and + the same vector reproduced the same verdict on every retry — so a write's + fate depended on where its embedding landed in the damaged graph. After + ``mempalace repair rebuild-index`` all five committed. + + Refusing is also the honest answer for the write that does not hang: chromadb + acknowledges it into sqlite and the metadata segment, then leaves it out of + the HNSW segment. That drawer is filed and reported as success while being + invisible to vector search — 16 such drawers came out of a single checkpoint + that returned "completed successfully in 2s". + + The probe behind this is pure sqlite + pickle, so the gate costs no chromadb + interaction on the path it is protecting. + """ + if tool_name not in _VECTOR_WRITE_TOOLS: + return None + + _refresh_vector_disabled_flag() + + if not _vector_disabled: + return None + + return { + "jsonrpc": "2.0", + "id": req_id, + "error": { + "code": _DIVERGED_INDEX_ERROR_CODE, + "message": ("Palace vector index is diverged; refusing the write until it is rebuilt"), + "data": { + "tool": tool_name, + "palace": _config.palace_path or "", + "vector_disabled_reason": _vector_disabled_reason, + "hint": ( + "Stop the MemPalace MCP servers, run `mempalace repair rebuild-index`, " + "then mempalace_reconnect. Recall keeps working meanwhile through the " + "BM25 fallback; writes stay refused so they can neither hang inside " + "chromadb nor land outside the index." + ), + }, + }, + } + + def _mcp_tool_preflight_refusal(req_id, tool_name: str): """Run MCP request preflight gates outside handle_request complexity.""" @@ -4891,6 +4968,10 @@ def _mcp_tool_preflight_refusal(req_id, tool_name: str): if sqlite_integrity_error is not None: return sqlite_integrity_error + diverged_index_error = _mcp_diverged_index_refusal(req_id, tool_name) + if diverged_index_error is not None: + return diverged_index_error + return _mcp_peer_writer_refusal(req_id, tool_name) @@ -5041,7 +5122,10 @@ def handle_request(request): if "entry" not in tool_args or tool_args["entry"] is None: tool_args["entry"] = content_val try: - result = _decorate_mcp_tool_result(tool_name, TOOLS[tool_name]["handler"](**tool_args)) + with _write_stall_watch(tool_name): + result = _decorate_mcp_tool_result( + tool_name, TOOLS[tool_name]["handler"](**tool_args) + ) return { "jsonrpc": "2.0", @@ -5267,6 +5351,130 @@ def _maybe_eager_warmup_embedder() -> None: ) +_WRITE_STALL_WARN_ENV = "MEMPALACE_MCP_WRITE_STALL_WARN_SECS" +_WRITE_STALL_WARN_DEFAULT = 60.0 +_WRITE_STALL_EXIT_ENV = "MEMPALACE_MCP_WRITE_STALL_EXIT_SECS" +_WRITE_STALL_EXIT_DEFAULT = 0.0 +# EX_TEMPFAIL: the palace is fine, this process is not. A client that restarts +# the server gets a working one; a zero exit would read as an orderly shutdown. +_WRITE_STALL_EXIT_CODE = 75 + +_write_stall_lock = threading.Lock() +# Optional[dict]: {"tool": str, "since": float(monotonic), "warned": bool} +_write_stall_inflight: Optional[dict] = None + + +def _write_stall_secs(env_name: str, default: float) -> float: + raw = os.environ.get(env_name, "") + if not raw.strip(): + return default + try: + return max(0.0, float(raw)) + except ValueError: + logger.warning("%s=%r is not a number; using %.0fs", env_name, raw, default) + return default + + +def _write_stall_action(elapsed: float, warn_secs: float, exit_secs: float, warned: bool): + """Decide what an in-flight vector write has earned: ``None``/warn/exit. + + Pure so the thresholds can be tested without a stalled write; ``exit`` is + checked first so a single tick can escalate straight past an unsent warning. + """ + if exit_secs > 0 and elapsed >= exit_secs: + return "exit" + if warn_secs > 0 and elapsed >= warn_secs and not warned: + return "warn" + return None + + +@contextlib.contextmanager +def _write_stall_watch(tool_name: str): + """Register a vector write as in flight for the stall watchdog.""" + global _write_stall_inflight + + if tool_name not in _VECTOR_WRITE_TOOLS: + yield + return + + with _write_stall_lock: + _write_stall_inflight = { + "tool": tool_name, + "since": time.monotonic(), + "warned": False, + } + try: + yield + finally: + with _write_stall_lock: + _write_stall_inflight = None + + +def _start_write_stall_watchdog() -> None: + """Start a daemon thread that reports a vector write that stopped returning. + + A chromadb write has no timeout of its own. When one blocks, this process + holds the dispatch lock, the palace mine lock and the writer lease, so it + answers nothing else and every peer session drops to read-only — and no log + on the server side says why. The only trace of a 25-minute outage was the + client's own "tool still running" ticks; the handshake stayed healthy, and + ``mempalace_status`` could not answer because the dispatch lock was held by + the stuck call. So the report has to come from a thread that is not waiting + on that lock, and it has to reach stderr, where the MCP host records it. + + ``MEMPALACE_MCP_WRITE_STALL_WARN_SECS`` (default 60, 0 disables) sets when to + warn. ``MEMPALACE_MCP_WRITE_STALL_EXIT_SECS`` (default 0 = never) lets an + operator turn the wedge into a restartable failure: a server stuck inside + chromadb will not recover, and exiting is what releases the locks its peers + are queued behind. + """ + warn_secs = _write_stall_secs(_WRITE_STALL_WARN_ENV, _WRITE_STALL_WARN_DEFAULT) + exit_secs = _write_stall_secs(_WRITE_STALL_EXIT_ENV, _WRITE_STALL_EXIT_DEFAULT) + if warn_secs <= 0 and exit_secs <= 0: + return + + thresholds = [t for t in (warn_secs, exit_secs) if t > 0] + interval = max(1.0, min(15.0, min(thresholds) / 4)) + + def _watchdog() -> None: + while True: + time.sleep(interval) + with _write_stall_lock: + inflight = _write_stall_inflight + if inflight is None: + continue + elapsed = time.monotonic() - inflight["since"] + action = _write_stall_action(elapsed, warn_secs, exit_secs, inflight["warned"]) + tool = inflight["tool"] + if action == "warn": + inflight["warned"] = True + if action == "warn": + logger.warning( + "%s has been inside the palace write path for %.0fs and has not " + "returned. chromadb writes have no timeout: this server now answers " + "nothing else and holds the writer lease, so peer sessions are " + "read-only. Check `mempalace repair --dry-run` for HNSW divergence; " + "restarting this MCP server releases the locks.", + tool, + elapsed, + ) + elif action == "exit": + logger.error( + "%s stalled in the palace write path for %.0fs (limit %s=%.0fs); " + "exiting so the palace locks are released and the client can " + "reconnect. The stalled write is lost — rebuild the index before " + "retrying it.", + tool, + elapsed, + _WRITE_STALL_EXIT_ENV, + exit_secs, + ) + os._exit(_WRITE_STALL_EXIT_CODE) + + t = threading.Thread(target=_watchdog, name="mcp-write-stall-watchdog", daemon=True) + t.start() + + def _start_idle_exit_watchdog() -> None: """Start a daemon thread that exits the process after an idle period. @@ -5686,6 +5894,10 @@ def _run_stdio_loop() -> None: # that outlived their Claude Code session (#1552). _start_idle_exit_watchdog() + # Say so when a chromadb write stops coming back, from a thread the stuck + # call is not blocking. + _start_write_stall_watchdog() + while True: try: line = sys.stdin.readline() @@ -5759,6 +5971,7 @@ def _run_http_loop() -> None: # soon as the process is alive. _refresh_vector_disabled_flag() _start_idle_exit_watchdog() + _start_write_stall_watchdog() raw_warmup = os.environ.get("MEMPALACE_EAGER_WARMUP", "").strip().lower() if raw_warmup in _WARMUP_TRUTHY: diff --git a/tests/test_vector_write_gate.py b/tests/test_vector_write_gate.py new file mode 100644 index 0000000..22542a8 --- /dev/null +++ b/tests/test_vector_write_gate.py @@ -0,0 +1,252 @@ +"""Tests for the write-side half of the #1222 HNSW divergence guard. + +The capacity probe already routed *reads* to the BM25 fallback when the flushed +HNSW segment lags sqlite. These cover the two gaps that left writes exposed: + +* ``_mcp_diverged_index_refusal`` — a vector write into a diverged index is + refused at dispatch instead of reaching chromadb, where an upsert can block + for the life of the process. +* the write-stall watchdog — when a write does stop coming back anyway, a + thread that the stuck call is not blocking says so on stderr, and an operator + can opt into turning the wedge into a restartable exit. +""" + +from __future__ import annotations + +import threading + +import pytest + + +REASON = "HNSW index holds 803 elements but sqlite has 820 embeddings" + + +@pytest.fixture +def diverged(monkeypatch): + """A palace whose vector index is known-diverged, probe counted.""" + from mempalace import mcp_server + + probes = {"n": 0} + + def _probe(): + probes["n"] += 1 + + monkeypatch.setattr(mcp_server, "_refresh_vector_disabled_flag", _probe) + monkeypatch.setattr(mcp_server, "_vector_disabled", True) + monkeypatch.setattr(mcp_server, "_vector_disabled_reason", REASON) + return probes + + +# ── Dispatch gate ───────────────────────────────────────────────────────── + + +def test_vector_write_refused_while_index_diverged(diverged): + from mempalace import mcp_server + + err = mcp_server._mcp_diverged_index_refusal(req_id=7, tool_name="mempalace_add_drawer") + + assert err is not None + assert err["id"] == 7 + assert err["error"]["code"] == mcp_server._DIVERGED_INDEX_ERROR_CODE + data = err["error"]["data"] + assert data["tool"] == "mempalace_add_drawer" + assert data["vector_disabled_reason"] == REASON + # The refusal has to carry the way out, not just the verdict. + assert "rebuild-index" in data["hint"] + + +@pytest.mark.parametrize( + "tool", + sorted( + { + "mempalace_add_drawer", + "mempalace_update_drawer", + "mempalace_delete_drawer", + "mempalace_delete_by_source", + "mempalace_diary_write", + "mempalace_checkpoint", + "mempalace_mine", + "mempalace_sync", + } + ), +) +def test_every_vector_write_tool_is_gated(diverged, tool): + from mempalace import mcp_server + + assert tool in mcp_server._MUTATING_TOOLS, "the vector set must stay a subset" + assert mcp_server._mcp_diverged_index_refusal(req_id=1, tool_name=tool) is not None + + +@pytest.mark.parametrize( + "tool", + [ + "mempalace_kg_add", + "mempalace_kg_invalidate", + "mempalace_kg_supersede", + "mempalace_create_tunnel", + "mempalace_delete_tunnel", + "mempalace_delete_hallway", + ], +) +def test_non_vector_writes_are_not_gated(diverged, tool): + """The knowledge graph and hallways keep their own state — a broken HNSW + segment has no say over them, and must not even cost them a probe.""" + from mempalace import mcp_server + + assert mcp_server._mcp_diverged_index_refusal(req_id=1, tool_name=tool) is None + assert diverged["n"] == 0 + + +def test_read_tools_are_not_gated(diverged): + """Reads have their own fallback (BM25); refusing them here would break it.""" + from mempalace import mcp_server + + assert mcp_server._mcp_diverged_index_refusal(req_id=1, tool_name="mempalace_search") is None + + +def test_healthy_index_allows_the_write(monkeypatch): + from mempalace import mcp_server + + monkeypatch.setattr(mcp_server, "_refresh_vector_disabled_flag", lambda: None) + monkeypatch.setattr(mcp_server, "_vector_disabled", False) + + assert ( + mcp_server._mcp_diverged_index_refusal(req_id=1, tool_name="mempalace_add_drawer") is None + ) + + +def test_gate_re_probes_so_a_repair_un_gates_without_restart(diverged): + """The gate must consult the probe on every call: a long-lived stdio server + has to notice `mempalace repair` finishing in another process.""" + from mempalace import mcp_server + + mcp_server._mcp_diverged_index_refusal(req_id=1, tool_name="mempalace_add_drawer") + mcp_server._mcp_diverged_index_refusal(req_id=2, tool_name="mempalace_add_drawer") + + assert diverged["n"] == 2 + + +def test_preflight_reports_divergence_ahead_of_the_peer_writer_lock(diverged, monkeypatch): + """Both gates can be up at once — a peer holds the lease *because* this + palace is wedged. The diverged verdict is the actionable one.""" + from mempalace import mcp_server + + monkeypatch.setattr(mcp_server, "_mcp_read_only_refusal", lambda req_id, tool_name: None) + monkeypatch.setattr(mcp_server, "_mcp_sqlite_integrity_refusal", lambda req_id, tool_name: None) + monkeypatch.setattr( + mcp_server, + "_mcp_peer_writer_refusal", + lambda req_id, tool_name: {"error": {"code": -32001}}, + ) + + err = mcp_server._mcp_tool_preflight_refusal(req_id=3, tool_name="mempalace_add_drawer") + + assert err["error"]["code"] == mcp_server._DIVERGED_INDEX_ERROR_CODE + + +def test_dispatch_refuses_before_the_handler_runs(diverged, monkeypatch): + from mempalace import mcp_server + + monkeypatch.setattr(mcp_server, "_mcp_sqlite_integrity_refusal", lambda req_id, tool_name: None) + monkeypatch.setattr(mcp_server, "_mcp_peer_writer_refusal", lambda req_id, tool_name: None) + + def _must_not_run(**kwargs): # pragma: no cover - the point is that it never runs + raise AssertionError("handler reached chromadb despite the diverged index") + + monkeypatch.setitem(mcp_server.TOOLS["mempalace_add_drawer"], "handler", _must_not_run) + + resp = mcp_server.handle_request( + { + "method": "tools/call", + "id": 11, + "params": { + "name": "mempalace_add_drawer", + "arguments": {"wing": "w", "room": "r", "content": "c"}, + }, + } + ) + + assert resp["error"]["code"] == mcp_server._DIVERGED_INDEX_ERROR_CODE + + +# ── Write-stall watchdog ────────────────────────────────────────────────── + + +def test_stall_action_warns_once_then_stays_quiet(): + from mempalace import mcp_server + + assert mcp_server._write_stall_action(59.0, 60.0, 0.0, False) is None + assert mcp_server._write_stall_action(60.0, 60.0, 0.0, False) == "warn" + assert mcp_server._write_stall_action(600.0, 60.0, 0.0, True) is None + + +def test_stall_action_escalates_to_exit_when_opted_in(): + from mempalace import mcp_server + + assert mcp_server._write_stall_action(299.0, 60.0, 300.0, True) is None + assert mcp_server._write_stall_action(300.0, 60.0, 300.0, True) == "exit" + # A single tick may cross both lines; exit wins over an unsent warning. + assert mcp_server._write_stall_action(300.0, 60.0, 300.0, False) == "exit" + + +def test_stall_action_respects_disabled_thresholds(): + from mempalace import mcp_server + + assert mcp_server._write_stall_action(10_000.0, 0.0, 0.0, False) is None + + +def test_stall_secs_falls_back_on_garbage(monkeypatch): + from mempalace import mcp_server + + monkeypatch.setenv(mcp_server._WRITE_STALL_WARN_ENV, "soon") + assert mcp_server._write_stall_secs(mcp_server._WRITE_STALL_WARN_ENV, 60.0) == 60.0 + + monkeypatch.setenv(mcp_server._WRITE_STALL_WARN_ENV, "-5") + assert mcp_server._write_stall_secs(mcp_server._WRITE_STALL_WARN_ENV, 60.0) == 0.0 + + monkeypatch.delenv(mcp_server._WRITE_STALL_WARN_ENV) + assert mcp_server._write_stall_secs(mcp_server._WRITE_STALL_WARN_ENV, 60.0) == 60.0 + + +def test_stall_watch_registers_the_write_and_clears_it(): + from mempalace import mcp_server + + with mcp_server._write_stall_watch("mempalace_add_drawer"): + inflight = mcp_server._write_stall_inflight + assert inflight is not None + assert inflight["tool"] == "mempalace_add_drawer" + assert inflight["warned"] is False + + assert mcp_server._write_stall_inflight is None + + +def test_stall_watch_clears_on_failure(): + """A raising handler must not leave a phantom write in flight — the next + write would inherit its clock and trip the watchdog.""" + from mempalace import mcp_server + + with pytest.raises(RuntimeError): + with mcp_server._write_stall_watch("mempalace_checkpoint"): + raise RuntimeError("boom") + + assert mcp_server._write_stall_inflight is None + + +def test_stall_watch_ignores_tools_that_never_reach_chromadb(): + from mempalace import mcp_server + + with mcp_server._write_stall_watch("mempalace_search"): + assert mcp_server._write_stall_inflight is None + + +def test_watchdog_thread_not_started_when_both_thresholds_are_zero(monkeypatch): + from mempalace import mcp_server + + monkeypatch.setenv(mcp_server._WRITE_STALL_WARN_ENV, "0") + monkeypatch.setenv(mcp_server._WRITE_STALL_EXIT_ENV, "0") + before = [t.name for t in threading.enumerate()] + + mcp_server._start_write_stall_watchdog() + + after = [t.name for t in threading.enumerate()] + assert after == before From 121a4552ace3574dd1a8577a436ab82446bbee77 Mon Sep 17 00:00:00 2001 From: mvalentsev Date: Thu, 6 Aug 2026 16:33:10 +0500 Subject: [PATCH 2/4] perf(embedding): group documents by size before sub-batching (#2104) The tokenizer pads every row of a sub-batch to the longest sequence in it, and attention costs batch x heads x length^2 per layer, so arrival order decided what an embeddinggemma run cost: one long verbatim message dragged its whole sub-batch up to its own length. Sort by UTF-8 byte length before the split and scatter each row back to its input index. An input that fits one sub-batch is left alone, since every row pads to the same width either way. Over 43,157 sweep drawers from 160 transcripts, padded token slots drop 39.7% and the quadratic attention term 45.0%. Wall clock on 52 messages with one model instance: 1169.70 s to 490.83 s, control baseline 1286.23 s. Row values move by at most one float32 ULP (1.2e-07, cosine 0.99999992), which is reduction-order rounding, not a change of meaning. Only palaces configured for embeddinggemma are affected; the default MiniLM embedder pads to a fixed width. --- CHANGELOG.md | 1 + mempalace/embedding.py | 65 +++++++++-- tests/test_embeddinggemma.py | 206 ++++++++++++++++++++++++++++++++++- 3 files changed, 258 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5b1b31..5714327 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Performance +- **EmbeddingGemma groups documents by size before sub-batching.** The tokenizer pads every row of a sub-batch to the longest sequence in it, so arrival order decided the bill: one long verbatim message dragged a whole sub-batch up to its own length. Measured over 43,157 `sweep` drawers from 160 Claude Code transcripts, padded token slots drop 39.7% and the quadratic attention term 45.0%. Vectors move by at most one float32 ULP (1.2e-07 absolute, cosine 0.99999992), which is reduction-order rounding and not a change of meaning. Applies to `embedding_model: embeddinggemma` only; the default MiniLM embedder pads to a fixed width and was never affected. (#2104) - **HNSW capacity probes are cached** and invalidated by palace file signature, so repeated MCP status/taxonomy paths no longer re-scan native segment files on every call. (#2051, #1471) - **`chunk_text` line numbering is O(N)** via incremental tallies, fixing multi-second hangs on large sources. (#2054, #2055) diff --git a/mempalace/embedding.py b/mempalace/embedding.py index 9c701bc..0e4ae74 100644 --- a/mempalace/embedding.py +++ b/mempalace/embedding.py @@ -216,7 +216,9 @@ _EMBEDDINGGEMMA_MAX_LEN = 2048 # matches the internal batch size of chromadb's ONNXMiniLM_L6_V2, whose # chunked _forward survives the same call sites. embeddinggemma's # sentence_embedding output is attention-masked, so sub-batch padding -# does not change any row's vector. +# does not change any row's vector. __call__ decides which documents share +# a sub-batch by size rather than by arrival order (#2104), because the run +# is priced on that padded length and not on the document count. _EMBEDDINGGEMMA_BATCH_SIZE = 32 @@ -356,6 +358,33 @@ class EmbeddinggemmaONNX: self._session = session def __call__(self, input: str | list[str] | None) -> list[list[float]]: # noqa: A002 — ChromaDB EF protocol + """Embed ``input``, returning one vector per document in input order. + + Documents are grouped by size before the sub-batch split. The + tokenizer pads every row of a sub-batch to the longest sequence in + it, and attention cost per layer is batch x heads x length^2, so one + long document drags a whole sub-batch up to its own length. Without + grouping the bill is set by arrival order: a verbatim transcript + whose long tool results sit between one-line replies pays the long + length for nearly every row (#2104). + + An input that fits a single sub-batch is left in arrival order: every + row pads to the same width either way, so the keys would buy nothing + on the one-document search path. + + Regrouping does not change what a row means. The model's + ``sentence_embedding`` output is attention-masked, so padding never + enters a row's values; what does move is float32 rounding, because a + different padded width changes the reduction order inside the GEMMs. + Measured against the same documents embedded in arrival order, that + residual peaks at one float32 ULP (1.2e-07 absolute, cosine + 0.99999992). + + The key is UTF-8 byte length rather than character count: this model + is multilingual, and bytes per token vary far less across scripts + than characters per token do. The sort is stable, so equal-size + documents keep arrival order and the split stays reproducible. + """ if isinstance(input, str): # A bare string would be iterated character by character below, # silently producing one garbage vector per character. @@ -367,14 +396,21 @@ class EmbeddinggemmaONNX: return [] self._lazy_load() np = self._np - embeddings: list[list[float]] = [] - # Tokenize and run per sub-batch, not over the whole input: padding - # is to the longest sequence in the sub-batch, and the ONNX runtime - # only ever holds batch_size rows of attention buffers at a time - # (#1770). - for start in range(0, len(input), self._batch_size): - chunk = input[start : start + self._batch_size] - texts = [_EMBEDDINGGEMMA_PREFIX + t for t in chunk] + # One sub-batch pads identically whatever the order, so the sort is + # only worth its keys once the input splits into several. + order: range | list[int] = range(len(input)) + if len(input) > self._batch_size: + order = sorted(range(len(input)), key=lambda i: len(input[i].encode("utf-8"))) + # Row i is filled by the sub-batch that carries document i. ``order`` + # is a permutation of every index, so no placeholder survives; callers + # (ChromaDB included) zip the result against their ids positionally. + embeddings: list[list[float] | None] = [None] * len(input) + # Tokenize and run per sub-batch, not over the whole input: the ONNX + # runtime only ever holds batch_size rows of attention buffers at a + # time (#1770). + for start in range(0, len(order), self._batch_size): + idxs = order[start : start + self._batch_size] + texts = [_EMBEDDINGGEMMA_PREFIX + input[i] for i in idxs] encs = self._tokenizer.encode_batch(texts) input_ids = np.asarray([e.ids for e in encs], dtype=np.int64) input_ids = _sanitize_embeddinggemma_input_ids( @@ -390,7 +426,16 @@ class EmbeddinggemmaONNX: # L2-normalize so cosine similarity == dot product (matches what the # MTEB methodology assumes; ChromaDB's distance is configured for it). norms = np.linalg.norm(sent_emb, axis=1, keepdims=True) + 1e-12 - embeddings.extend((sent_emb / norms).tolist()) + rows = (sent_emb / norms).tolist() + if len(rows) != len(idxs): + # zip would truncate silently and leave a None in the result, + # which only surfaces far downstream in the caller's array + # conversion. Fail on the sub-batch that came back short. + raise RuntimeError( + f"embeddinggemma returned {len(rows)} rows for a {len(idxs)}-document sub-batch" + ) + for row_index, row in zip(idxs, rows): + embeddings[row_index] = row return embeddings def embed_query(self, input: list[str]) -> list[list[float]]: # noqa: A002 — ChromaDB EF protocol diff --git a/tests/test_embeddinggemma.py b/tests/test_embeddinggemma.py index eb512ed..35aa0c1 100644 --- a/tests/test_embeddinggemma.py +++ b/tests/test_embeddinggemma.py @@ -190,7 +190,9 @@ def test_call_chunks_large_batches(patched_lazy_load, monkeypatch): monkeypatch.setattr(_FakeTokenizer, "encode_batch", recording_encode_batch) ef = embedding.EmbeddinggemmaONNX() n = embedding._EMBEDDINGGEMMA_BATCH_SIZE * 2 + 6 - docs = [f"doc {i}" for i in range(n)] + # Descending sizes, so size-sorted order is the reverse of arrival order + # and the assertion below cannot pass on both. + docs = [f"{'x' * (n - i)} doc {i}" for i in range(n)] out = ef(docs) assert batch_sizes == [ @@ -198,9 +200,11 @@ def test_call_chunks_large_batches(patched_lazy_load, monkeypatch): embedding._EMBEDDINGGEMMA_BATCH_SIZE, 6, ], f"expected bounded sub-batches, got {batch_sizes}" - # Sub-batches must cover the input in order; combined with the per-chunk - # extend in __call__ this pins output row order to input order. - assert captured_texts == [embedding._EMBEDDINGGEMMA_PREFIX + d for d in docs] + # Sub-batches cover the input in size-sorted order; the scatter in + # __call__ puts every row back at its own input index afterwards. + ordered = sorted(docs, key=lambda d: len(d.encode("utf-8"))) + assert ordered != docs, "fixture must not already be in size order" + assert captured_texts == [embedding._EMBEDDINGGEMMA_PREFIX + d for d in ordered] arr = np.asarray(out) assert arr.shape == (n, 384), f"chunked outputs must concatenate to (n, 384), got {arr.shape}" assert np.allclose(np.linalg.norm(arr, axis=1), 1.0, atol=1e-5) @@ -251,6 +255,200 @@ def test_custom_batch_size_is_honored(patched_lazy_load, monkeypatch): assert len(out) == 24 +# Bound before any test can monkeypatch the method, so the width helper +# below measures with the real fake and never re-enters a recorder. +_UNPATCHED_ENCODE_BATCH = _FakeTokenizer.encode_batch + + +def _record_batches(monkeypatch, sink): + """Capture the texts handed to each encode_batch call, in order.""" + + def recording_encode_batch(self, texts): + sink.append(list(texts)) + return _UNPATCHED_ENCODE_BATCH(self, texts) + + monkeypatch.setattr(_FakeTokenizer, "encode_batch", recording_encode_batch) + + +def _fake_padded_width(batches): + """Total padded token slots the fake tokenizer produces for `batches`. + + Measured by encoding, not by re-deriving the fake's padding rule, so the + two cannot drift apart and quietly turn the assertion into a tautology. + """ + tokenizer = _FakeTokenizer() + return sum(sum(len(e.ids) for e in _UNPATCHED_ENCODE_BATCH(tokenizer, b)) for b in batches) + + +def test_call_groups_documents_by_size(patched_lazy_load, monkeypatch): + """Similar-size documents must share a sub-batch. + + encode_batch pads every row to the longest sequence in the sub-batch and + attention cost per layer is batch x heads x length^2, so interleaving one + long document with short ones makes the short ones pay the long length + (#2104). Grouping by size is what keeps that bill proportional to the + text actually being embedded. + """ + batches = [] + _record_batches(monkeypatch, batches) + # One long document per sub-batch's worth of short ones: the pathological + # arrival order a verbatim transcript sweep produces. + long_doc = " ".join(["word"] * 200) + docs = [long_doc if i % _B == 0 else f"short {i}" for i in range(4 * _B)] + + ef = embedding.EmbeddinggemmaONNX() + out = ef(docs) + assert len(out) == len(docs) + + for texts in batches: + sizes = [len(t.encode("utf-8")) for t in texts] + assert sizes == sorted(sizes), f"sub-batch is not size-grouped: {sizes}" + + prefixed = [embedding._EMBEDDINGGEMMA_PREFIX + d for d in docs] + arrival_order = [prefixed[s : s + _B] for s in range(0, len(prefixed), _B)] + assert _fake_padded_width(batches) < _fake_padded_width(arrival_order), ( + "size grouping must lower the total padded width" + ) + + +def test_call_groups_by_utf8_size_not_character_count(patched_lazy_load, monkeypatch): + """The key is UTF-8 bytes, because this model is multilingual. + + A CJK document is ~3 bytes per character and roughly a token per + character, so ordering by character count would file it next to Latin + documents several times cheaper to embed. + """ + batches = [] + _record_batches(monkeypatch, batches) + # Same character count, very different byte count (and token count). + docs = ["a" * 90] * _B + ["中" * 90] * _B + ef = embedding.EmbeddinggemmaONNX() + ef(list(reversed(docs))) + + assert len(batches) == 2, f"expected two sub-batches, got {len(batches)}" + sizes = [[len(t.encode("utf-8")) for t in b] for b in batches] + assert max(sizes[0]) < min(sizes[1]), ( + f"CJK documents must not share a sub-batch with Latin ones: {sizes}" + ) + + +def test_call_size_grouping_is_stable(patched_lazy_load, monkeypatch): + """Equal-size documents keep their arrival order. + + An unstable sort would make the sub-batch split depend on nothing the + caller can see, so two identical inputs could take different code paths. + """ + batches = [] + _record_batches(monkeypatch, batches) + docs = [f"doc{i:03d}" for i in range(_B + 8)] # identical size, distinct text + ef = embedding.EmbeddinggemmaONNX() + ef(docs) + captured = [t for b in batches for t in b] + assert captured == [embedding._EMBEDDINGGEMMA_PREFIX + d for d in docs] + + +def test_call_keeps_arrival_order_within_a_single_sub_batch(patched_lazy_load, monkeypatch): + """An input that fits one sub-batch is not reordered. + + Every row pads to the same width either way, so the sort would buy + nothing and only add keys to compute on the search hot path. + """ + batches = [] + _record_batches(monkeypatch, batches) + docs = [f"{'x' * (_B - i)} doc {i}" for i in range(_B)] # descending size + ef = embedding.EmbeddinggemmaONNX() + ef(docs) + assert batches == [[embedding._EMBEDDINGGEMMA_PREFIX + d for d in docs]] + + +class _MarkerTokenizer(_FakeTokenizer): + """Tokenizer whose first token id carries that document's own length. + + ``_FakeTokenizer`` emits all-zero ids padded to one width, so a fake + session cannot tell its rows apart, which is exactly what an + order-restoration test has to observe. + """ + + def encode_batch(self, texts): + widths = [len(t) for t in texts] + padded = max(widths) + + class _Enc: + def __init__(self, marker): + self.ids = [marker] + [0] * (padded - 1) + self.attention_mask = [1] * marker + [0] * (padded - marker) + + return [_Enc(w) for w in widths] + + +class _MarkerSession: + """Emit a vector whose first two dims encode the row's marker id. + + Both dims scale by the same L2 norm, so ``row[0] / row[1]`` survives + normalization and identifies which document produced the row. + """ + + _WIDTH = 2 * embedding._EMBEDDINGGEMMA_DIM + + def run(self, _output_names, feed): + ids = feed["input_ids"] + batch, length = ids.shape + sent = np.zeros((batch, self._WIDTH), dtype=np.float64) + sent[:, 0] = ids[:, 0] + sent[:, 1] = 1.0 + return [np.zeros((batch, length, self._WIDTH), dtype=np.float64), sent] + + +def _marker_ef(patched_lazy_load, session=None): + """An EF wired to the marker fakes, with the real lazy load short-circuited.""" + # patched_lazy_load is taken so a future tightening of _lazy_load's + # early-return cannot turn these tests into a 300 MB model download. + ef = embedding.EmbeddinggemmaONNX() + ef._tokenizer = _MarkerTokenizer() + ef._session = session if session is not None else _MarkerSession() + ef._output_idx = 1 + ef._np = np + return ef + + +def test_call_returns_rows_at_their_input_index(patched_lazy_load): + """Row i of the result must be the embedding of document i. + + Sub-batching by size reorders the work; ChromaDB zips the returned + vectors against the ids positionally, so grouping without the matching + scatter would file every drawer under another drawer's vector. That + half-applied state is what this pins: arrival order trivially satisfies + it, so it is the grouping tests that cover the other direction. + """ + ef = _marker_ef(patched_lazy_load) + # Strictly descending sizes, so grouping reverses arrival order and an + # unscattered result would be visibly wrong. + docs = ["x" * n for n in range(200, 200 - 3 * _B, -1)] + out = ef(docs) + + assert len(out) == len(docs) + assert all(row is not None for row in out), "every index must be filled" + markers = [round(row[0] / row[1]) for row in out] + assert markers == [len(embedding._EMBEDDINGGEMMA_PREFIX + d) for d in docs] + + +def test_call_rejects_a_short_row_count_from_the_session(patched_lazy_load): + """A session returning fewer rows than documents must fail loudly. + + Scattering by index would otherwise leave a None in the result and the + caller would only trip over it much later, converting to an array. + """ + + class _ShortSession(_MarkerSession): + def run(self, output_names, feed): + last_hidden, sent = super().run(output_names, feed) + return [last_hidden, sent[:-1]] + + ef = _marker_ef(patched_lazy_load, session=_ShortSession()) + with pytest.raises(RuntimeError, match="rows for a"): + ef(["x" * n for n in range(200, 200 - 2 * _B, -1)]) + + def test_batch_size_below_one_is_rejected(): """A zero or negative batch size would loop forever or embed nothing.""" with pytest.raises(ValueError, match="batch_size"): From 6093340754688ed1c84932a885e4b2cef4599cc7 Mon Sep 17 00:00:00 2001 From: mvalentsev Date: Thu, 6 Aug 2026 16:48:56 +0500 Subject: [PATCH 3/4] fix(normalize): parse ChatGPT export arrays, not just single conversations (#2160) A real ChatGPT conversations.json is a top-level array of conversation objects. _try_chatgpt_json requires a single dict with a "mapping" key, and no other parser in the chain claims the array either, so the export fell through to the plain-text path: mine --mode convos reported success while filing serialized JSON sliced at arbitrary offsets. Unwrap the array and normalize each conversation separately, mirroring _try_claude_ai_json_split, so a freshly downloaded export sitting next to the previous one files only the conversations that are actually new. The ChatGPT parser is now reached for every element of every top-level JSON array under the mined tree, so type-check its nested shapes: ten of the twelve malformed payloads now covered by tests raised AttributeError, KeyError or TypeError, which would abort a whole mine run since convo_miner catches only OSError and ValueError. --- CHANGELOG.md | 1 + mempalace/normalize.py | 82 +++++++++++++--- tests/test_normalize.py | 212 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 283 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5b1b31..5cf0ac2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Bug Fixes +- **ChatGPT data exports are parsed instead of stored as raw JSON.** A real `conversations.json` is a top-level array of conversations, which no parser claimed, so `mine --mode convos` chunked the raw JSON and lost every speaker turn while reporting success. Each conversation now normalizes to its own transcript, as Claude.ai privacy exports already do, so per-conversation dedup survives a re-export. The ChatGPT parser also type-checks its nested shapes, so an unrelated array carrying a `mapping` key is declined instead of raising. (#2160) - **Local backends enforce process-lifetime single-writer ownership.** File-backed and unknown backends require one writer owner for the full process lifetime (daemon holds the lease until workers exit; writable MCP HTTP acquires ownership before bind and refuses startup when blocked). Read-only MCP may coexist; `sqlite_exact` opens genuine query-only/immutable readers; remote Milvus/Zilliz remain multi-process. Addresses multi-writer SQLite/WAL corruption from MCP HTTP + daemon + mine topologies. (#2079, #2045) - **Chroma HNSW write defaults match chromadb** (`batch_size=100` / `sync_threshold=1000`) instead of the old 2/2 bloat guard that rewrote segments thousands of times on large mines. (#2107, #2106) - **Repair and recovery are safer under contention.** `repair --mode from-sqlite` takes the mine-lock before archiving; rebuilds preserve a verified temp collection when the live swap fails; sparse drawers with zero `embedding_metadata` rows are no longer dropped; truncated ID pagination fails loud instead of pretending success. (#2109, #2086, #2087) diff --git a/mempalace/normalize.py b/mempalace/normalize.py index 7ac4f6b..9503aeb 100644 --- a/mempalace/normalize.py +++ b/mempalace/normalize.py @@ -5,7 +5,8 @@ normalize.py — Convert any chat export format to MemPalace transcript format. Supported: - Plain text with > markers (pass through) - Claude.ai JSON export - - ChatGPT conversations.json + - ChatGPT conversations.json (a single conversation, or the top-level + array of them that a real data export ships) - Claude Code JSONL (with tool_use/tool_result block capture) - OpenAI Codex CLI JSONL - Gemini CLI JSONL (~/.gemini/tmp//chats/session-*.jsonl) @@ -182,9 +183,13 @@ def normalize_conversations(filepath: str) -> list: the existing conversations did. This returns the pieces un-joined so callers can hash and dedup per conversation instead. - Non-bundle formats (a single Claude Code session, a ChatGPT export, - plain text, ...) always normalize to one conversation, so this returns - a one-element list for those — identical dedup granularity to before. + A ChatGPT data export is a bundle for the same reason: its + ``conversations.json`` is an array of conversations, so it splits per + conversation too. + + Non-bundle formats (a single Claude Code session, plain text, ...) + always normalize to one conversation, so this returns a one-element + list for those — identical dedup granularity to before. """ content = _read_transcript_file(filepath) @@ -248,6 +253,10 @@ def _try_normalize_json_split(content: str) -> Optional[list]: if split: return split + split = _try_chatgpt_export_json_split(data) + if split: + return split + for parser in (_try_chatgpt_json, _try_continue_json, _try_slack_json): normalized = parser(data) if normalized: @@ -623,8 +632,14 @@ def _collect_claude_messages(items) -> list: def _try_chatgpt_json(data) -> Optional[str]: - """ChatGPT conversations.json with mapping tree.""" - if not isinstance(data, dict) or "mapping" not in data: + """ChatGPT conversations.json with mapping tree. + + Every nested shape is type-checked rather than assumed: this parser is + reached from ``_try_chatgpt_export_json_split`` for each element of any + top-level JSON array, so it must return None on unrelated payloads that + merely carry a ``mapping`` key instead of raising. + """ + if not isinstance(data, dict) or not isinstance(data.get("mapping"), dict): return None mapping = data["mapping"] messages = [] @@ -632,6 +647,8 @@ def _try_chatgpt_json(data) -> Optional[str]: root_id = None fallback_root = None for node_id, node in mapping.items(): + if not isinstance(node, dict): + continue if node.get("parent") is None: if node.get("message") is None: root_id = node_id @@ -645,24 +662,65 @@ def _try_chatgpt_json(data) -> Optional[str]: visited = set() while current_id and current_id not in visited: visited.add(current_id) - node = mapping.get(current_id, {}) + node = mapping.get(current_id) + if not isinstance(node, dict): + break msg = node.get("message") - if msg: - role = msg.get("author", {}).get("role", "") + if isinstance(msg, dict): + author = msg.get("author") + role = author.get("role", "") if isinstance(author, dict) else "" content = msg.get("content", {}) - parts = content.get("parts", []) if isinstance(content, dict) else [] + parts = content.get("parts") if isinstance(content, dict) else None + if not isinstance(parts, list): + parts = [] text = " ".join(str(p) for p in parts if isinstance(p, str) and p).strip() if role == "user" and text: messages.append(("user", text)) elif role == "assistant" and text: messages.append(("assistant", text)) - children = node.get("children", []) - current_id = children[0] if children else None + children = node.get("children") + next_id = children[0] if isinstance(children, list) and children else None + # Node ids index a dict and a visited set, so anything unhashable + # (a nested child object rather than an id) ends the walk. + current_id = next_id if isinstance(next_id, str) else None if len(messages) >= 2: return _messages_to_transcript(messages) return None +def _try_chatgpt_export_json_split(data) -> Optional[list]: + """ChatGPT data export: top-level array of conversation objects. + + The ``conversations.json`` OpenAI ships is an *array*, while + ``_try_chatgpt_json`` handles the single conversation object inside it. + Without this the whole export falls through to the plain-text path and is + chunked as raw JSON: the drawers hold serialized structure sliced at + arbitrary offsets, and every speaker turn is gone. + + Each conversation is kept as its own segment rather than concatenated, so + per-conversation dedup survives a re-export (see ``normalize_conversations``); + the joined form is reached through ``_try_normalize_json``. + + Runs after ``_try_gemini_json`` and ``_try_claude_ai_json_split`` and before + the ``_try_chatgpt_json``/``_try_continue_json``/``_try_slack_json`` loop. + That position is safe in both directions: Gemini requires a ``role="model"`` + entry and Claude.ai requires ``chat_messages``/``messages`` on the first + element, neither of which a ChatGPT conversation object has, while Slack + entries carry no ``mapping`` and Continue.dev sessions are not arrays at + all, so this parser declines them and they fall through unchanged. + """ + if not isinstance(data, list): + return None + + transcripts = [] + for convo in data: + transcript = _try_chatgpt_json(convo) + if transcript: + transcripts.append(transcript) + # None, not [], so an array of other JSON still reaches the later parsers. + return transcripts or None + + def _try_slack_json(data) -> Optional[str]: """ Slack channel export: [{"type": "message", "user": "...", "text": "..."}] diff --git a/tests/test_normalize.py b/tests/test_normalize.py index 8e2d17b..3336e0d 100644 --- a/tests/test_normalize.py +++ b/tests/test_normalize.py @@ -8,6 +8,7 @@ from mempalace.normalize import ( _format_tool_result, _format_tool_use, _messages_to_transcript, + _try_chatgpt_export_json_split, _try_chatgpt_json, _try_claude_ai_json, _try_claude_code_jsonl, @@ -16,9 +17,11 @@ from mempalace.normalize import ( _try_gemini_jsonl, _try_continue_json, _try_normalize_json, + _try_normalize_json_split, _try_pi_jsonl, _try_slack_json, normalize, + normalize_conversations, strip_noise, ) @@ -1067,6 +1070,215 @@ def test_chatgpt_json_too_few_messages(): assert result is None +# ── _try_chatgpt_export_json_split ──────────────────────────────────── + + +def _chatgpt_convo(title, question, answer): + """One conversation object, shaped like a real ChatGPT export entry.""" + return { + "title": title, + "create_time": 1740000000.0, + "mapping": { + "root": {"parent": None, "message": None, "children": ["msg1"]}, + "msg1": { + "parent": "root", + "message": { + "author": {"role": "user"}, + "content": {"parts": [question]}, + }, + "children": ["msg2"], + }, + "msg2": { + "parent": "msg1", + "message": { + "author": {"role": "assistant"}, + "content": {"parts": [answer]}, + }, + "children": [], + }, + }, + } + + +def test_chatgpt_export_array_is_normalized(): + """A real export is a top-level array, not a single conversation object.""" + # Deliberately not in alphabetical order, so a sort would be caught. + data = [ + _chatgpt_convo("Debug", "Why the segfault?", "You free the node twice."), + _chatgpt_convo("Trip", "Best month for Lisbon?", "May, for the light."), + ] + result = _try_chatgpt_export_json_split(data) + assert result is not None + assert len(result) == 2 + # Export order is preserved, and neither segment bleeds into the other. + assert result[0] == "> Why the segfault?\nYou free the node twice.\n" + assert result[1] == "> Best month for Lisbon?\nMay, for the light.\n" + + +def test_chatgpt_export_array_keeps_conversations_separate(): + """Each conversation stays its own segment, as for Claude.ai bundles.""" + data = [ + _chatgpt_convo("One", "Q1", "A1"), + _chatgpt_convo("Two", "Q2", "A2"), + ] + split = _try_normalize_json_split(json.dumps(data)) + assert split is not None + assert len(split) == 2 + assert sum("> Q1" in segment for segment in split) == 1 + assert sum("> Q2" in segment for segment in split) == 1 + + +def test_chatgpt_export_array_skips_entries_without_mapping(): + """Export metadata entries alongside conversations must not break parsing.""" + data = [ + {"title": "no mapping here"}, + _chatgpt_convo("Real", "Q1", "A1"), + ] + result = _try_chatgpt_export_json_split(data) + assert result is not None + assert len(result) == 1 + assert "> Q1" in result[0] + + +def test_chatgpt_export_array_ignores_unrelated_payloads(): + """Anything that is not an array of ChatGPT conversations is declined.""" + assert _try_chatgpt_export_json_split([]) is None + assert _try_chatgpt_export_json_split([1, 2, 3]) is None + assert _try_chatgpt_export_json_split({"mapping": {}}) is None + assert _try_chatgpt_export_json_split([{"role": "user", "content": "hi"}]) is None + # A .json file holding a bare scalar reaches the parsers too; it must not raise. + assert _try_chatgpt_export_json_split(42) is None + assert _try_chatgpt_export_json_split(None) is None + + +def test_chatgpt_export_array_yields_none_when_every_conversation_is_empty(): + """An array of unusable mappings must return None, never []. + + Returning [] would be falsy at the call site today, but the contract is + Optional[list]; None is what makes the dispatcher fall through cleanly. + """ + data = [ + { + "mapping": { + "root": {"parent": None, "message": None, "children": ["m1"]}, + "m1": { + "parent": "root", + "message": { + "author": {"role": "user"}, + "content": {"parts": ["only one turn"]}, + }, + "children": [], + }, + } + } + ] + assert _try_chatgpt_export_json_split(data) is None + + +def test_chatgpt_export_array_survives_malformed_mappings(): + """Arrays are now fed element-wise to the parser, so it must not raise. + + Any JSON array under the mined tree reaches this parser (an ES mapping + dump, a DB field map), and a raised exception would abort the whole + `mine --mode convos` run, since convo_miner only catches OSError/ValueError. + """ + malformed = [ + [{"mapping": [1, 2, 3]}], + [{"mapping": "oops"}], + [{"mapping": {"a": "not-a-dict"}}], + [{"mapping": {"r": {"parent": None, "message": {"author": None}, "children": []}}}], + [{"mapping": {"r": {"parent": None, "message": "text", "children": []}}}], + [{"mapping": {"r": {"parent": None, "message": None, "children": "nope"}}}], + [{"name": "users", "mapping": {"uid": "int", "email": "str"}}], + # Root is well formed but the node it points at is not a dict. + [{"mapping": {"r": {"parent": None, "message": None, "children": ["m1"]}, "m1": "oops"}}], + # `children` as a mapping would raise KeyError on children[0]. + [{"mapping": {"r": {"parent": None, "message": None, "children": {"a": "m1"}}}}], + # A tree serialised with child objects instead of ids: unhashable ids. + [{"mapping": {"r": {"parent": None, "message": None, "children": [["x"]]}}}], + [ + { + "name": "org-chart", + "mapping": { + "r": { + "parent": None, + "message": None, + "children": [{"id": "n1", "label": "Engineering"}], + } + }, + } + ], + # A cycle must terminate rather than spin. + [ + { + "mapping": { + "r": {"parent": None, "message": None, "children": ["a"]}, + "a": {"parent": "r", "message": None, "children": ["r"]}, + } + } + ], + ] + for payload in malformed: + assert _try_chatgpt_export_json_split(payload) is None + assert _try_normalize_json_split(json.dumps(payload)) is None + + +def test_chatgpt_json_string_parts_are_not_split_into_characters(): + """`parts` must be a list; a bare string would be iterated character-wise.""" + data = { + "mapping": { + "root": {"parent": None, "message": None, "children": ["m1"]}, + "m1": { + "parent": "root", + "message": {"author": {"role": "user"}, "content": {"parts": "secret"}}, + "children": ["m2"], + }, + "m2": { + "parent": "m1", + "message": {"author": {"role": "assistant"}, "content": {"parts": ["ok"]}}, + "children": [], + }, + } + } + result = _try_chatgpt_json(data) + assert result is None or "s e c r e t" not in result + + +def test_chatgpt_export_parser_does_not_shadow_slack_exports(): + """A Slack export is also a top-level array; it must still reach its parser.""" + slack = [ + {"type": "message", "user": "U1", "text": "deploy is green"}, + {"type": "message", "user": "U2", "text": "shipping it"}, + ] + split = _try_normalize_json_split(json.dumps(slack)) + assert split is not None + assert _SLACK_PROVENANCE_FOOTER.strip() in split[0] + + +def test_chatgpt_export_array_end_to_end(tmp_path): + """The miner path: an export file yields one entry per conversation.""" + f = tmp_path / "conversations.json" + f.write_text( + json.dumps( + [ + _chatgpt_convo("One", "Q1", "A1"), + _chatgpt_convo("Two", "Q2", "A2"), + ] + ), + encoding="utf-8", + ) + conversations = normalize_conversations(str(f)) + assert len(conversations) == 2 + assert conversations[0].startswith("> Q1") + assert conversations[1].startswith("> Q2") + + # normalize() joins the bundle, but must still produce a transcript + joined = normalize(str(f)) + assert joined.startswith(">") + assert "> Q1" in joined + assert "> Q2" in joined + + # ── _try_slack_json ──────────────────────────────────────────────────── From 495a12bff55b0099f67384a1f9dec0c7eb3c9060 Mon Sep 17 00:00:00 2001 From: mvalentsev Date: Thu, 6 Aug 2026 16:48:59 +0500 Subject: [PATCH 4/4] fix(repair): honor --dry-run in the default mode (#2144) The default (legacy) repair path ignored --dry-run and ran the real rebuild: it deleted any existing .backup, copied the live palace over it, re-filed the drawers collection through a staged temp copy, then rebuilt FTS5 and VACUUMed. #2095 and #2133 fixed this for --mode from-sqlite only. The preview now returns before ChromaBackend() is constructed. Opening a chromadb client is itself a write to chroma.sqlite3, so a preview that reached one could not be inert; the row count comes from sqlite_drawer_count instead, the read-only SQLite ground truth check_extraction_safety already trusts. Staying off the chromadb layer also keeps a dry run clear of the layer repair is separately reported to segfault in on a large palace (#2113). resolve_repair_preflight_errors() decides what a dry run does about the FTS5 autoheal. The autoheal is a write, so a preview must not run it, but skipping it routed an isolated inverted-index error into the abort banner and exit 1 - telling the operator to run offline sqlite3 .recover on a palace a real run heals by itself (#1596). The dry run now classifies the errors with the same _errors_are_isolated_fts5 predicate the real path gates on and continues; broader corruption still aborts in both modes. It is worded as an attempt rather than a promise, because the real heal still gives up when another process holds the mine lock, when the rebuild raises, or when quick_check is still dirty afterwards. The plan describes the real run in execution order. It names the live-collection delete the rebuild performs, since "re-file via a staged temp copy" alone reads as additive; it warns when an existing backup would be deleted; and it reports a no-op instead of a rebuild when the collection holds no rows. It states the #1208 truncation guard that can abort the run, and reports that guard as disabled when --confirm-truncation-ok is set, because check_extraction_safety returns immediately then and the promised abort would not happen. An unreadable count fails closed with a non-zero exit, for parity with the from-sqlite preview. --- CHANGELOG.md | 1 + mempalace/cli.py | 35 +++++- mempalace/repair.py | 135 +++++++++++++++++++++-- tests/test_cli.py | 189 +++++++++++++++++++++++++++++++ tests/test_repair.py | 233 +++++++++++++++++++++++++++++++++++++++ website/reference/cli.md | 15 ++- 6 files changed, 592 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5b1b31..e73cfba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), - **Chroma HNSW write defaults match chromadb** (`batch_size=100` / `sync_threshold=1000`) instead of the old 2/2 bloat guard that rewrote segments thousands of times on large mines. (#2107, #2106) - **Repair and recovery are safer under contention.** `repair --mode from-sqlite` takes the mine-lock before archiving; rebuilds preserve a verified temp collection when the live swap fails; sparse drawers with zero `embedding_metadata` rows are no longer dropped; truncated ID pagination fails loud instead of pretending success. (#2109, #2086, #2087) - **`repair --mode from-sqlite --dry-run` is a true preview.** It no longer archives or re-embeds; it prints per-collection would-be counts from SQLite ground truth and exits without touching the palace. Unreadable counts fail closed instead of inventing zeros. (#2133, #2095, #1654) +- **`repair --dry-run` is a true preview in the default (legacy) mode too.** That path ignored the flag entirely and ran the real rebuild — deleting any existing `.backup`, copying the palace over it, and re-filing the drawers collection. It now prints a read-only plan and exits without opening a chromadb client, which is itself a write to `chroma.sqlite3`. The plan names the live-collection delete the rebuild performs, warns when an existing backup would be destroyed, and reports the truncation guard as disabled when `--confirm-truncation-ok` is set. An isolated FTS5 inverted-index error is reported as auto-healable instead of raising the manual-recovery abort a real run never reaches, unreadable counts fail closed with a non-zero exit, and the `--dry-run` help no longer claims to be `--mode max-seq-id` only. (#2144) - **HNSW divergence is preflighted before remaining `col.count()` crash sites** across mine, dedup, migrate, repair, and palace helpers. (#2093) - **Re-mine and conversation ingest no longer lose or duplicate drawers.** Content-hash dedup prevents duplicate LLM conversation drawers; sweeper drawers are excluded from convo extract-mode purge scope and failed purges abort; search returns round-trippable `drawer_id` values for `get_drawer`. (#2050, #2125, #2089, #2090, #2044, #2080) - **MCP and daemon lifecycle harden multi-agent use.** Read-only mode refuses config and checkpoint-ack tools that rewrite host state; stdio MCP exits on stdin EOF/broken pipe so orphaned sessions release locks; daemon jobs refused the palace lock are deferred instead of failed permanently. (#2126, #2103, #2101, #2072, #2029, #2014) diff --git a/mempalace/cli.py b/mempalace/cli.py index 58ec3f1..99eab46 100644 --- a/mempalace/cli.py +++ b/mempalace/cli.py @@ -1140,13 +1140,14 @@ def cmd_repair(args): _close_chroma_handles, _extract_drawers, _post_rebuild_cleanup, + _preview_legacy_repair, _promote_temp_collection, _rebuild_collection_via_temp, check_extraction_safety, index_read_recovery_guidance, - maybe_autoheal_fts5_index, maybe_repair_poisoned_max_seq_id_before_rebuild, print_sqlite_integrity_abort, + resolve_repair_preflight_errors, sqlite_integrity_errors, ) @@ -1250,9 +1251,13 @@ def cmd_repair(args): # stack trace instead of the friendly abort message. Run quick_check # here so we can surface the clear recovery instructions and exit # cleanly before chromadb's compactor touches the disk. - sqlite_errors = sqlite_integrity_errors(palace_path) - if sqlite_errors: - sqlite_errors = maybe_autoheal_fts5_index(palace_path, sqlite_errors) + dry_run = getattr(args, "dry_run", False) + # The FTS5 autoheal inside this call is a write, so a --dry-run predicts + # its outcome instead of performing it (#1596 is auto-healable and must + # not surface as an abort in a preview). + sqlite_errors = resolve_repair_preflight_errors( + palace_path, sqlite_integrity_errors(palace_path), dry_run=dry_run + ) if sqlite_errors: print_sqlite_integrity_abort(palace_path, sqlite_errors) sys.exit(1) @@ -1260,7 +1265,7 @@ def cmd_repair(args): preflight = maybe_repair_poisoned_max_seq_id_before_rebuild( palace_path, backup=getattr(args, "backup", True), - dry_run=getattr(args, "dry_run", False), + dry_run=dry_run, assume_yes=getattr(args, "yes", False), ) if preflight is not None: @@ -1271,6 +1276,24 @@ def cmd_repair(args): print(f"{'=' * 55}\n") print(f" Palace: {palace_path}") + if dry_run: + # Return before the backend is used at all: the chromadb client this + # path opens is itself a write to chroma.sqlite3 (measured — the file + # hash changes on get_collection alone, before count()), so a preview + # that reached it could not be inert. Staying off the chromadb layer + # also keeps a dry run clear of the layer repair is separately reported + # to segfault in on a large palace (#2113). Exit non-zero on an + # unreadable count for parity with the from-sqlite preview above, so + # `--dry-run && repair --yes` cannot walk into the destructive run + # after a failed preview (#2095, #2133). + if not _preview_legacy_repair( + palace_path=palace_path, + collection_name=collection_name, + confirm_truncation_ok=getattr(args, "confirm_truncation_ok", False), + ): + sys.exit(1) + return + backend = ChromaBackend() # Try to read existing drawers @@ -2135,7 +2158,7 @@ def main(): 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 what the repair would do and exit without modifying the palace", ) # repair-status — read-only HNSW capacity health check (#1222) diff --git a/mempalace/repair.py b/mempalace/repair.py index f5257da..4ab67dc 100644 --- a/mempalace/repair.py +++ b/mempalace/repair.py @@ -1715,6 +1715,21 @@ def rebuild_from_sqlite( ) +def _print_unreadable_count_refusal(*, collection_name: str, palace_path: str) -> None: + """Refuse to preview a collection whose SQLite row count cannot be read. + + Fail closed: inventing 0 would hide an unreadable source and make the + operator believe a real run would upsert nothing (review note on #1654 / + #2095). Shared by both previews so the wording cannot drift apart. + """ + print( + f"\n Cannot preview [{collection_name}]: SQLite row count is unreadable " + f"at {os.path.join(palace_path, 'chroma.sqlite3')}.\n" + " Fix source readability (schema, lock, permissions) and re-run " + "--dry-run; refusing to invent zero counts." + ) + + def _preview_rebuild_from_sqlite( *, source_palace: str, @@ -1739,15 +1754,7 @@ def _preview_rebuild_from_sqlite( for cname in _recoverable_collections(): n = sqlite_drawer_count(source_palace, cname) if n is None: - # Fail closed: inventing 0 would hide an unreadable source and - # make the operator believe a real rebuild would upsert nothing - # (review note on #1654 / #2095). - print( - f"\n Cannot preview [{cname}]: SQLite row count is unreadable " - f"at {os.path.join(source_palace, 'chroma.sqlite3')}.\n" - " Fix source readability (schema, lock, permissions) and re-run " - "--dry-run; refusing to invent zero counts." - ) + _print_unreadable_count_refusal(collection_name=cname, palace_path=source_palace) return {} counts[cname] = n print(f" [{cname}] would re-embed and upsert {n} rows") @@ -1758,6 +1765,116 @@ def _preview_rebuild_from_sqlite( return counts +def _preview_legacy_repair( + *, + palace_path: str, + collection_name: str, + confirm_truncation_ok: bool = False, +) -> dict[str, int]: + """Read-only preview for the default (legacy) ``repair`` path (``dry_run=True``). + + Never opens a chromadb client, takes a lock, or writes. Opening a client is + itself a write to ``chroma.sqlite3``, so the row count comes from the + read-only SQLite ground truth :func:`check_extraction_safety` already + trusts. That is a different source than the real run rebuilds from (it + re-files what the chromadb collection layer returns), so the plan below + states the ``#1208`` contingency rather than promising the number. + + ``confirm_truncation_ok`` mirrors the real run's flag: it switches that + contingency off, so the preview has to say the guard is disabled rather + than promise an abort that would not happen. + + Returns ``{}`` when the count is unreadable so a broken preview cannot look + like a valid plan (#1654, #2095, #2133). + """ + print("\n DRY RUN — no changes will be made.") + n = sqlite_drawer_count(palace_path, collection_name) + if n is None: + _print_unreadable_count_refusal(collection_name=collection_name, palace_path=palace_path) + print(f"{'=' * 55}\n") + return {} + + if n == 0: + # The real run stops at ``total == 0`` with "Nothing to repair.", or — + # when the collection is absent altogether — at the index-read error + # that points to --mode from-sqlite. Neither backs up nor rebuilds, so + # promising a backup and a VACUUM here would describe a run that does + # not happen. + print( + f" [{collection_name}] chroma.sqlite3 holds no rows. A real run would report\n" + " nothing to repair, or an index read error, and change nothing." + ) + print(f"{'=' * 55}\n") + return {collection_name: 0} + + backup_path = os.path.normpath(palace_path) + ".backup" + if confirm_truncation_ok: + print( + f" [{collection_name}] chroma.sqlite3 holds {n} rows, and --confirm-truncation-ok\n" + " is set, so the #1208 truncation guard is DISABLED. A real run would re-file\n" + f" whatever the chromadb collection layer returns, even if that is fewer than {n}\n" + " rows, and the difference would be destroyed. It would, in order:" + ) + else: + print( + f" [{collection_name}] chroma.sqlite3 holds {n} rows. A real run would extract them\n" + " through the chromadb collection layer first and abort without changes if that\n" + f" returns fewer than {n} (#1208 truncation guard). It would then, in order:" + ) + if os.path.exists(backup_path): + print(f" 1. DELETE the existing backup at {backup_path} — or refuse outright") + print(" if it is not a palace — and copy the live palace in its place") + else: + print(f" 1. copy the palace directory to {backup_path}") + print(f" 2. DELETE the live '{collection_name}' collection and re-file the extracted rows") + print(" into a fresh one, staged and verified in a temp collection first") + print(" 3. rebuild the FTS5 index and VACUUM chroma.sqlite3") + print("\n Without --yes it would ask for confirmation before step 1.") + print(" Re-run without --dry-run to execute.") + print(f"{'=' * 55}\n") + return {collection_name: n} + + +def resolve_repair_preflight_errors( + palace_path: str, + errors: list[str], + *, + dry_run: bool, + progress=print, +) -> list[str]: + """Return the quick_check errors that still block a repair. + + A real run heals an isolated malformed FTS5 inverted index in place and + carries on (#1596). ``--dry-run`` must not perform that write, so it + classifies the errors with the same :func:`_errors_are_isolated_fts5` + predicate the real path gates on: an isolated FTS5 error is reported and + cleared, anything broader still aborts. Without this a preview would print + the ABORT banner — offline ``sqlite3 .recover``, recreate the FTS5 table — + for a palace the tool repairs by itself. + + The prediction is deliberately the optimistic branch, and it is stated as + an attempt rather than a promise: the real heal still returns the errors + unchanged when another process holds the mine lock, when the rebuild + raises, or when ``quick_check`` is still dirty afterwards. A dry run cannot + tell those apart without taking the lock and writing, which is exactly what + it must not do, so the wording names them instead. + """ + if not errors: + return errors + if not dry_run: + return maybe_autoheal_fts5_index(palace_path, errors, progress=progress) + if _errors_are_isolated_fts5(errors): + progress( + "\n DRY RUN — quick_check reports an isolated FTS5 inverted-index error.\n" + " A real run would attempt an in-place rebuild of that index from the\n" + " intact content table and continue if it succeeds; it aborts instead if\n" + " another process holds the mine lock or the rebuild leaves quick_check\n" + " dirty. This preview leaves the index untouched." + ) + return [] + return errors + + def _rebuild_from_sqlite_locked( *, source_palace: str, diff --git a/tests/test_cli.py b/tests/test_cli.py index a5c8b69..a190c94 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1214,6 +1214,195 @@ def test_cmd_repair_uses_configured_collection(mock_config_cls, tmp_path, capsys ] +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_default_mode_dry_run_writes_nothing(mock_config_cls, tmp_path, capsys): + """``repair --dry-run`` with no --mode must print a plan and leave the palace alone. + + #2095 / #2133 fixed this for ``--mode from-sqlite``; the default (legacy) + path still ran the backup copy and the full rebuild. + """ + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True, dry_run=True) + mock_col = MagicMock() + mock_col.count.return_value = 2 + # A concrete extract payload, so that if the dry-run guard ever regresses + # this test fails on its assertions instead of spinning in _extract_drawers. + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_backend = _mock_backend_for(col=mock_col) + + with ( + patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend), + patch("mempalace.repair.sqlite_drawer_count", return_value=2) as mock_count, + patch("mempalace.migrate.confirm_destructive_action") as mock_confirm, + ): + cmd_repair(args) + + out = capsys.readouterr().out + assert "DRY RUN — no changes will be made." in out + assert "holds 2 rows" in out + assert "Repair complete" not in out + # The count comes from read-only SQLite, never from a chromadb client: + # opening one is itself a write to chroma.sqlite3. + mock_count.assert_called_once_with(str(palace_dir), "mempalace_drawers") + mock_backend.get_collection.assert_not_called() + mock_col.get.assert_not_called() + mock_backend.create_collection.assert_not_called() + mock_backend.delete_collection.assert_not_called() + mock_confirm.assert_not_called() + assert not (tmp_path / "palace.backup").exists() + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_default_mode_dry_run_honours_custom_collection( + mock_config_cls, tmp_path, capsys +): + """The preview must target the configured collection, not a hardcoded name.""" + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "custom_drawers" + args = argparse.Namespace(palace=None, yes=True, dry_run=True) + + with ( + patch("mempalace.backends.chroma.ChromaBackend", return_value=_mock_backend_for()), + patch("mempalace.repair.sqlite_drawer_count", return_value=7) as mock_count, + ): + cmd_repair(args) + + assert "[custom_drawers]" in capsys.readouterr().out + mock_count.assert_called_once_with(str(palace_dir), "custom_drawers") + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_default_mode_dry_run_reports_healable_fts5_and_continues( + mock_config_cls, tmp_path, capsys +): + """An isolated FTS5 error is auto-healed by a real run, so the preview must not abort. + + Aborting here would hand the operator the manual ``sqlite3 .recover`` + recipe for a palace ``maybe_autoheal_fts5_index`` repairs by itself (#1596). + """ + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True, dry_run=True) + + with ( + patch( + "mempalace.repair.sqlite_integrity_errors", + return_value=["malformed inverted index for FTS5 table x"], + ), + patch("mempalace.repair.maybe_autoheal_fts5_index") as mock_autoheal, + patch("mempalace.repair.sqlite_drawer_count", return_value=4), + patch("mempalace.backends.chroma.ChromaBackend", return_value=_mock_backend_for()), + ): + cmd_repair(args) + + out = capsys.readouterr().out + mock_autoheal.assert_not_called() + assert "isolated FTS5 inverted-index error" in out + assert "ABORT" not in out + assert "holds 4 rows" in out + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_default_mode_dry_run_still_aborts_on_broad_corruption( + mock_config_cls, tmp_path, capsys +): + """Corruption a real run cannot heal must still abort the preview with exit 1.""" + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True, dry_run=True) + + with ( + patch( + "mempalace.repair.sqlite_integrity_errors", + return_value=["*** in database main *** Page 42 is never used"], + ), + patch("mempalace.repair.maybe_autoheal_fts5_index") as mock_autoheal, + patch("mempalace.backends.chroma.ChromaBackend"), + pytest.raises(SystemExit) as excinfo, + ): + cmd_repair(args) + + assert excinfo.value.code == 1 + mock_autoheal.assert_not_called() + assert "ABORT" in capsys.readouterr().out + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_non_dry_run_still_autoheals_fts5(mock_config_cls, tmp_path, capsys): + """The real path must keep calling the autoheal — the patch restructured this branch.""" + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True) + mock_col = MagicMock() + mock_col.count.return_value = 0 + + with ( + patch( + "mempalace.repair.sqlite_integrity_errors", + return_value=["malformed inverted index for FTS5 table x"], + ), + patch("mempalace.repair.maybe_autoheal_fts5_index", return_value=[]) as mock_autoheal, + patch( + "mempalace.backends.chroma.ChromaBackend", + return_value=_mock_backend_for(col=mock_col), + ), + ): + cmd_repair(args) + + mock_autoheal.assert_called_once() + assert "Nothing to repair" in capsys.readouterr().out + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_default_mode_dry_run_refuses_to_invent_zero(mock_config_cls, tmp_path, capsys): + """An unreadable count must refuse AND exit non-zero, like the from-sqlite preview. + + Otherwise ``repair --dry-run && repair --yes`` walks into the destructive + run after a preview that could not be produced. + """ + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + sqlite3.connect(str(palace_dir / "chroma.sqlite3")).close() + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True, dry_run=True) + mock_backend = _mock_backend_for(col=MagicMock()) + + with ( + patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend), + patch("mempalace.repair.sqlite_drawer_count", return_value=None), + pytest.raises(SystemExit) as excinfo, + ): + cmd_repair(args) + + out = capsys.readouterr().out + assert excinfo.value.code == 1 + assert "Cannot preview [mempalace_drawers]" in out + assert "refusing to invent zero counts" in out + assert "would extract" not in out + mock_backend.get_collection.assert_not_called() + + @patch("mempalace.cli.MempalaceConfig") def test_cmd_repair_restores_backup_on_live_rebuild_failure(mock_config_cls, tmp_path, capsys): """When the live swap fails after the delete, recovery must PROMOTE the diff --git a/tests/test_repair.py b/tests/test_repair.py index 082b92a..60fce21 100644 --- a/tests/test_repair.py +++ b/tests/test_repair.py @@ -2420,6 +2420,239 @@ def test_rebuild_from_sqlite_dry_run_fails_closed_when_count_unreadable(tmp_path assert not dest.exists() +# ── _preview_legacy_repair — the default (legacy) repair --dry-run ───── + + +def test_preview_legacy_repair_leaves_the_palace_byte_identical(tmp_path, capsys): + """The preview must not change a single byte of a real palace.""" + import hashlib + + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [(f"d{i}", f"b{i}", {"wing": "w"}) for i in range(6)]) + db = palace / "chroma.sqlite3" + before = hashlib.sha256(db.read_bytes()).hexdigest() + tree_before = sorted((p.name, p.stat().st_size) for p in palace.iterdir()) + + counts = repair._preview_legacy_repair( + palace_path=str(palace), collection_name="mempalace_drawers" + ) + + assert counts == {"mempalace_drawers": 6} + assert hashlib.sha256(db.read_bytes()).hexdigest() == before + assert sorted((p.name, p.stat().st_size) for p in palace.iterdir()) == tree_before + assert [p for p in tmp_path.iterdir() if p.name.endswith(".backup")] == [] + out = capsys.readouterr().out + assert "holds 6 rows" in out + assert "#1208 truncation guard" in out + + +def test_cmd_repair_dry_run_leaves_a_real_palace_byte_identical(tmp_path, capsys): + """End-to-end: everything cmd_repair touches ahead of the preview is read-only. + + The helper test above covers ``_preview_legacy_repair`` on its own. This one + covers the calls ``cmd_repair`` makes before reaching it — the quick_check + preflight and the poisoned-bookmark detector — against a palace that really + exists on disk, which is the only assertion that would catch a write + sneaking into any of them. + """ + import argparse + import hashlib + + from mempalace.cli import cmd_repair + + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [(f"d{i}", f"b{i}", {"wing": "w"}) for i in range(4)]) + + def snapshot(): + return { + str(p.relative_to(palace)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(palace.rglob("*")) + if p.is_file() + } + + before = snapshot() + args = argparse.Namespace(palace=str(palace), yes=True, dry_run=True) + + with patch("mempalace.cli.MempalaceConfig") as mock_config_cls: + mock_config_cls.return_value.palace_path = str(palace) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + cmd_repair(args) + + out = capsys.readouterr().out + assert snapshot() == before + assert not (tmp_path / "palace.backup").exists() + assert "DRY RUN — no changes will be made." in out + assert "holds 4 rows" in out + assert "Repair complete" not in out + + +def test_preview_legacy_repair_zero_rows_promises_no_backup(tmp_path, capsys): + """A real run stops at "Nothing to repair." — the preview must not promise a rebuild. + + ``sqlite_drawer_count`` returns 0 (not None) for an absent collection, so + the fail-closed guard does not cover this case. + """ + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [("d1", "b1", {"wing": "w"})]) + + counts = repair._preview_legacy_repair( + palace_path=str(palace), collection_name="collection_that_does_not_exist" + ) + + out = capsys.readouterr().out + assert counts == {"collection_that_does_not_exist": 0} + assert "holds no rows" in out + assert "copy the palace directory" not in out + assert "VACUUM" not in out + + +def test_preview_legacy_repair_warns_it_would_delete_an_existing_backup(tmp_path, capsys): + """Destroying the operator's previous backup is the step a preview must not hide.""" + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [("d1", "b1", {"wing": "w"})]) + backup = tmp_path / "palace.backup" + backup.mkdir() + + repair._preview_legacy_repair(palace_path=str(palace), collection_name="mempalace_drawers") + + out = capsys.readouterr().out + assert f"DELETE the existing backup at {backup}" in out + assert backup.exists() + + +def test_preview_legacy_repair_warns_when_the_backup_path_is_a_file(tmp_path, capsys): + """The real run branches on exists(), not isdir(). + + A regular file at .backup makes a real run refuse at the backup + validation step, so a preview that promised a plain copy would describe a + run that never happens. + """ + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [("d1", "b1", {"wing": "w"})]) + backup = tmp_path / "palace.backup" + backup.write_text("not a palace") + + repair._preview_legacy_repair(palace_path=str(palace), collection_name="mempalace_drawers") + + out = capsys.readouterr().out + assert f"DELETE the existing backup at {backup}" in out + assert "refuse outright" in out + assert "copy the palace directory" not in out + assert backup.read_text() == "not a palace" + + +def test_preview_legacy_repair_names_the_live_collection_delete(tmp_path, capsys): + """The real run calls delete_collection on the live collection. + + "re-file via a staged temp copy" alone reads as additive, which is the one + thing an operator must not misread about a destructive rebuild. + """ + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [("d1", "b1", {"wing": "w"})]) + + repair._preview_legacy_repair(palace_path=str(palace), collection_name="mempalace_drawers") + + out = capsys.readouterr().out + assert "DELETE the live 'mempalace_drawers' collection" in out + + +def test_preview_legacy_repair_reports_the_truncation_guard_as_disabled(tmp_path, capsys): + """--confirm-truncation-ok switches the #1208 abort off, so the preview must say so. + + The abort message the guard prints tells operators to re-run with this + flag, so the flag plus --dry-run is a combination they are actively + steered into. Promising the guard there would be a false safety claim. + """ + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [(f"d{i}", f"b{i}", {"wing": "w"}) for i in range(3)]) + + counts = repair._preview_legacy_repair( + palace_path=str(palace), + collection_name="mempalace_drawers", + confirm_truncation_ok=True, + ) + + out = capsys.readouterr().out + assert counts == {"mempalace_drawers": 3} + assert "#1208 truncation guard is DISABLED" in out + assert "the difference would be destroyed" in out + assert "abort without changes" not in out + + +def test_preview_legacy_repair_fails_closed_when_count_unreadable(tmp_path, capsys, monkeypatch): + """Unreadable count must return {} rather than render a zero-row plan.""" + palace = tmp_path / "palace" + _seed_palace(palace, "mempalace_drawers", [("d1", "b1", {"wing": "w"})]) + monkeypatch.setattr(repair, "sqlite_drawer_count", lambda *a, **k: None) + + counts = repair._preview_legacy_repair( + palace_path=str(palace), collection_name="mempalace_drawers" + ) + + out = capsys.readouterr().out + assert counts == {} + assert "refusing to invent zero counts" in out + assert "holds" not in out + + +# ── resolve_repair_preflight_errors ─────────────────────────────────── + + +def test_resolve_preflight_dry_run_clears_isolated_fts5_without_healing(tmp_path, capsys): + """A dry run predicts the autoheal instead of performing its write.""" + called = [] + errs = ["malformed inverted index for FTS5 table x"] + + out_errors = repair.resolve_repair_preflight_errors( + str(tmp_path), + errs, + dry_run=True, + progress=lambda *a, **k: called.append(a), + ) + + assert out_errors == [] + assert called and "isolated FTS5 inverted-index error" in called[0][0] + + +def test_resolve_preflight_dry_run_keeps_broad_corruption(tmp_path): + """Errors a real run cannot heal must survive so the caller still aborts.""" + errs = ["*** in database main *** Page 42 is never used"] + + assert repair.resolve_repair_preflight_errors(str(tmp_path), errs, dry_run=True) == errs + + +def test_resolve_preflight_real_run_delegates_to_autoheal(tmp_path, monkeypatch): + """Outside a dry run the behaviour is unchanged: hand off to the autoheal.""" + seen = {} + + def fake_autoheal(path, errors, *, progress=print): + seen["args"] = (path, errors) + return [] + + monkeypatch.setattr(repair, "maybe_autoheal_fts5_index", fake_autoheal) + errs = ["malformed inverted index for FTS5 table x"] + + assert repair.resolve_repair_preflight_errors(str(tmp_path), errs, dry_run=False) == [] + assert seen["args"] == (str(tmp_path), errs) + + +def test_resolve_preflight_passes_empty_errors_through(tmp_path, monkeypatch): + """A clean quick_check must not invoke the autoheal at all. + + Asserting only on the return value would pass with the early-out deleted, + since the autoheal hands empty errors straight back. + """ + called = [] + monkeypatch.setattr( + repair, + "maybe_autoheal_fts5_index", + lambda path, errors, **kw: called.append(path) or errors, + ) + + assert repair.resolve_repair_preflight_errors(str(tmp_path), [], dry_run=False) == [] + assert called == [] + + def test_rebuild_from_sqlite_in_place_validates_source_before_archiving(tmp_path): """In-place + archive_existing_dest=True with a dir that lacks chroma.sqlite3 must NOT rename the dir before bailing. An earlier diff --git a/website/reference/cli.md b/website/reference/cli.md index f1c244b..d15d2e7 100644 --- a/website/reference/cli.md +++ b/website/reference/cli.md @@ -134,7 +134,20 @@ Rebuild palace vector index from stored data. Fixes segfaults after database cor mempalace repair ``` -Creates a backup at `.backup` before rebuilding. +Creates a backup at `.backup` before rebuilding, replacing any backup already there. + +| Flag | Description | +|------|-------------| +| `rebuild-index` | Positional alias for `--mode from-sqlite --archive-existing` | +| `--mode` | `legacy` (default), `max-seq-id`, or `from-sqlite` | +| `--dry-run` | Print what the repair would do and exit without modifying the palace | +| `--yes` | Skip confirmation for destructive changes | +| `--backup` | Back up SQLite before mutation (default: on) | +| `--source` | Source palace for `--mode from-sqlite` (defaults to `--palace`) | +| `--archive-existing` | Rename the existing palace to `.pre-rebuild-` first | +| `--segment` | Segment UUID filter for `--mode max-seq-id` | +| `--from-sidecar` | Pre-corruption `chroma.sqlite3` to copy clean `max_seq_id` values from | +| `--confirm-truncation-ok` | Override the truncation safety guard. Disables the abort that protects you when the collection layer returns fewer drawers than SQLite holds | ## `mempalace mcp`