diff --git a/CHANGELOG.md b/CHANGELOG.md index f5b1b31..1215292 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,15 +18,18 @@ 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) ### 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) - **`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/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/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/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/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_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"): 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 ──────────────────────────────────────────────────── 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/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 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`