diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7c575d87bf..cb02be1035 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -30,13 +30,18 @@ jobs: python-version: "3.13" cache: 'pip' - run: pip install -e ".[dev]" - # ChromaDB's rust HNSW core intermittently fails compaction on Windows - # ("Failed to apply logs to the hnsw segment writer") regardless of our - # code — a long-standing, non-reproducible-on-Linux/macOS flake. Retry - # ONLY that specific transient error (via --only-rerun) so real, + # ChromaDB's rust HNSW core intermittently fails compaction or reader + # initialization on Windows ("Failed to apply logs to the hnsw segment + # writer" / "Error creating hnsw segment reader: Nothing found on disk") + # regardless of our code — long-standing, non-reproducible-on-Linux/macOS + # flakes. Retry ONLY those specific transient errors (via --only-rerun) so real, # deterministic failures still fail on the first run. Linux/macOS jobs # deliberately run with no reruns so genuine regressions surface there. - - run: python -m pytest tests/ -v --ignore=tests/benchmarks --cov=mempalace --cov-report=term-missing --cov-fail-under=80 --durations=10 --reruns 2 --reruns-delay 5 --only-rerun "Failed to apply logs to the hnsw segment writer" + - run: >- + python -m pytest tests/ -v --ignore=tests/benchmarks --cov=mempalace + --cov-report=term-missing --cov-fail-under=80 --durations=10 --reruns 2 + --reruns-delay 5 --only-rerun + "Failed to apply logs to the hnsw segment writer|Error creating hnsw segment reader: Nothing found on disk" test-macos: runs-on: macos-latest diff --git a/mempalace/encoding_repair.py b/mempalace/encoding_repair.py new file mode 100644 index 0000000000..999e7991d3 --- /dev/null +++ b/mempalace/encoding_repair.py @@ -0,0 +1,508 @@ +"""Safely repair high-confidence UTF-8 mojibake in MemPalace drawers.""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path +from typing import Callable, Iterator, Optional, TextIO, Union + +_BACKUP_FORMAT = "mempalace-encoding-repair" +_BACKUP_VERSION = 1 +_UNDEFINED_CP1252_BYTES = frozenset( + { + 0x81, + 0x8D, + 0x8F, + 0x90, + 0x9D, + } +) + + +def _cp1252_character( + byte_value: int, +) -> str: + """Map a legacy byte to the character stored in old palaces.""" + if byte_value in _UNDEFINED_CP1252_BYTES: + # Pre-3.1 Windows reads could preserve these five undefined + # Windows-1252 byte values as their corresponding invisible + # C1 control code points. + return chr(byte_value) + + return bytes([byte_value]).decode("cp1252") + + +def _encode_mojibake_candidate( + text: str, +) -> bytes: + """Recover original bytes, including undefined CP1252 values.""" + raw = bytearray() + + for character in text: + codepoint = ord(character) + + if codepoint in _UNDEFINED_CP1252_BYTES: + raw.append(codepoint) + else: + raw.extend(character.encode("cp1252")) + + return bytes(raw) + + +_CONTINUATION_CHARS = "".join( + _cp1252_character(byte_value) + for byte_value in range( + 0x80, + 0xC0, + ) +) +_CONTINUATION_CLASS = re.escape(_CONTINUATION_CHARS) + +# These are the characteristic visible lead characters produced by +# common UTF-8-as-Windows-1252 corruption: +# +# C2/C3 -> Â/à +# E2 -> â +# F0 -> ð +# EF -> ï +# +# C4/C5 -> Ä/Å are deliberately excluded because strings such as Ų +# can be legitimate scientific text. +_HIGH_CONFIDENCE_RUN = re.compile( + rf"(?:" + rf"[ÂÃ][{_CONTINUATION_CLASS}]" + rf"|â[{_CONTINUATION_CLASS}]{{2}}" + rf"|ð[{_CONTINUATION_CLASS}]{{3}}" + rf"|ï[{_CONTINUATION_CLASS}]{{2}}" + rf")+" +) + + +def _decode_high_confidence_run( + match: re.Match, +) -> str: + candidate = match.group(0) + + try: + return _encode_mojibake_candidate(candidate).decode("utf-8") + except ( + UnicodeEncodeError, + UnicodeDecodeError, + ): + return candidate + + +def repair_mojibake_once(text: str) -> str: + """Repair one layer of high-confidence UTF-8-as-CP1252 mojibake.""" + return _HIGH_CONFIDENCE_RUN.sub( + _decode_high_confidence_run, + text, + ) + + +def repair_mojibake( + text: str, + *, + max_passes: int = 3, +) -> str: + """Repair repeated high-confidence mojibake layers until stable.""" + if max_passes < 1: + raise ValueError("max_passes must be at least 1") + + current = text + + for _ in range(max_passes): + repaired = repair_mojibake_once(current) + + if repaired == current: + break + + current = repaired + + return current + + +def _result_field(result, name: str): + if isinstance(result, dict): + return result.get(name) + + return getattr(result, name, None) + + +def _collection_name( + collection, +) -> Optional[str]: + """Resolve a collection name through MemPalace backend wrappers.""" + + def resolve_name(candidate) -> Optional[str]: + name = getattr( + candidate, + "name", + None, + ) + + if callable(name): + try: + name = name() + except TypeError: + name = None + + return str(name) if name else None + + direct_name = resolve_name(collection) + if direct_name: + return direct_name + + # ChromaCollection already provides this resolver for its wrapped + # chromadb collection. + resolver = getattr( + collection, + "_collection_name", + None, + ) + + if callable(resolver): + try: + resolved = resolver() + except ( + AttributeError, + TypeError, + ): + resolved = None + + if resolved: + return str(resolved) + + # Defensive fallback for thin wrappers that expose only their inner + # collection object. + inner = getattr( + collection, + "_collection", + None, + ) + + if inner is not None and inner is not collection: + return resolve_name(inner) + + return None + + +def _open_private_backup( + path: Path, +) -> TextIO: + path.parent.mkdir( + parents=True, + exist_ok=True, + ) + + descriptor = os.open( + str(path), + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + + try: + try: + os.chmod(path, 0o600) + except ( + OSError, + NotImplementedError, + ): + pass + + return os.fdopen( + descriptor, + "w", + encoding="utf-8", + newline="\n", + ) + except Exception: + os.close(descriptor) + raise + + +def _write_json_line( + handle: TextIO, + value: dict, +) -> None: + handle.write( + json.dumps( + value, + ensure_ascii=True, + sort_keys=True, + ) + + "\n" + ) + + +def _write_backup_header( + handle: TextIO, + collection, +) -> None: + _write_json_line( + handle, + { + "format": _BACKUP_FORMAT, + "version": _BACKUP_VERSION, + "collection": _collection_name(collection), + }, + ) + + +def _read_backup_header( + path: Path, +) -> dict: + try: + with path.open( + "r", + encoding="utf-8", + ) as handle: + header = json.loads(handle.readline()) + except OSError as exc: + raise ValueError(f"could not read repair backup: {path}") from exc + except ( + json.JSONDecodeError, + TypeError, + ) as exc: + raise ValueError("repair backup has an invalid header") from exc + + if not isinstance( + header, + dict, + ) or (header.get("format") != _BACKUP_FORMAT or header.get("version") != _BACKUP_VERSION): + raise ValueError("unsupported repair backup format") + + return header + + +def _iter_backup_records( + path: Path, +) -> Iterator[tuple[str, str]]: + _read_backup_header(path) + + with path.open( + "r", + encoding="utf-8", + ) as handle: + # Skip the validated header. + handle.readline() + + for line_number, line in enumerate( + handle, + start=2, + ): + if not line.strip(): + continue + + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"invalid backup JSON at line {line_number}") from exc + + drawer_id = record.get("id") if isinstance(record, dict) else None + document = record.get("original_document") if isinstance(record, dict) else None + + if not isinstance( + drawer_id, + str, + ) or not isinstance( + document, + str, + ): + raise ValueError(f"invalid repair backup record at line {line_number}") + + yield drawer_id, document + + +def repair_collection( + collection, + *, + apply: bool = False, + page_size: int = 500, + backup_path: Optional[Union[str, Path]] = None, + on_change: Optional[Callable[[str, str, str], None]] = None, +) -> dict: + """Scan a collection and optionally repair high-confidence mojibake.""" + if page_size < 1: + raise ValueError("page_size must be at least 1") + + if apply and backup_path is None: + raise ValueError("backup_path is required when apply=True") + + backup = Path(backup_path) if backup_path is not None else None + backup_handle: Optional[TextIO] = None + backup_used: Optional[str] = None + + scanned = 0 + changed = 0 + updated = 0 + offset = 0 + + try: + while True: + page = collection.get( + limit=page_size, + offset=offset, + include=["documents"], + ) + ids = list( + _result_field( + page, + "ids", + ) + or [] + ) + documents = list( + _result_field( + page, + "documents", + ) + or [] + ) + + if len(ids) != len(documents): + raise RuntimeError("collection returned misaligned ids and documents") + + if not ids: + break + + page_changes = [] + + for drawer_id, document in zip( + ids, + documents, + ): + scanned += 1 + + if not isinstance( + document, + str, + ): + continue + + repaired = repair_mojibake(document) + + if repaired == document: + continue + + item = ( + str(drawer_id), + document, + repaired, + ) + page_changes.append(item) + changed += 1 + + if on_change is not None: + on_change(*item) + + if apply and page_changes: + if backup_handle is None: + assert backup is not None + + backup_handle = _open_private_backup(backup) + _write_backup_header( + backup_handle, + collection, + ) + backup_used = str(backup) + + for ( + drawer_id, + original, + _repaired, + ) in page_changes: + _write_json_line( + backup_handle, + { + "id": drawer_id, + "original_document": (original), + }, + ) + + # Originals must reach durable storage before their + # live collection rows are overwritten. + backup_handle.flush() + os.fsync(backup_handle.fileno()) + + collection.update( + ids=[item[0] for item in page_changes], + documents=[item[2] for item in page_changes], + ) + updated += len(page_changes) + + offset += len(ids) + + if len(ids) < page_size: + break + finally: + if backup_handle is not None: + backup_handle.close() + + return { + "scanned": scanned, + "changed": changed, + "updated": updated, + "backup_path": backup_used, + } + + +def restore_collection( + collection, + backup_path: Union[str, Path], + *, + batch_size: int = 500, +) -> dict: + """Restore original documents from an encoding-repair backup.""" + if batch_size < 1: + raise ValueError("batch_size must be at least 1") + + path = Path(backup_path) + header = _read_backup_header(path) + + backup_collection = header.get("collection") + target_collection = _collection_name(collection) + + if backup_collection and target_collection and backup_collection != target_collection: + raise ValueError( + f"backup belongs to collection {backup_collection!r}, not {target_collection!r}" + ) + + # Validate the complete file before performing the first restore write. + validated = sum(1 for _record in _iter_backup_records(path)) + + restored = 0 + batch_ids = [] + batch_documents = [] + + for ( + drawer_id, + document, + ) in _iter_backup_records(path): + batch_ids.append(drawer_id) + batch_documents.append(document) + + if len(batch_ids) < batch_size: + continue + + collection.update( + ids=batch_ids, + documents=batch_documents, + ) + restored += len(batch_ids) + batch_ids = [] + batch_documents = [] + + if batch_ids: + collection.update( + ids=batch_ids, + documents=batch_documents, + ) + restored += len(batch_ids) + + return { + "validated": validated, + "restored": restored, + } diff --git a/scripts/mempalace_repair_encoding.py b/scripts/mempalace_repair_encoding.py new file mode 100755 index 0000000000..b5c5aabadf --- /dev/null +++ b/scripts/mempalace_repair_encoding.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +"""Repair legacy Windows mojibake in a MemPalace collection.""" + +from __future__ import annotations + +import argparse +import sys +from datetime import datetime, timezone +from pathlib import Path + +from mempalace.config import MempalaceConfig +from mempalace.encoding_repair import ( + repair_collection, + restore_collection, +) +from mempalace.palace import ( + get_collection, + mine_palace_lock, +) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Conservatively repair high-confidence UTF-8 mojibake " + "in legacy MemPalace drawers. Dry-run is the default." + ) + ) + parser.add_argument( + "--palace", + help=("Palace path; defaults to the configured palace."), + ) + parser.add_argument( + "--collection", + help=("Collection name; defaults to the configured collection."), + ) + parser.add_argument( + "--page-size", + type=int, + default=500, + help=("Rows scanned or restored per page (default: 500)."), + ) + parser.add_argument( + "--preview-chars", + type=int, + default=180, + help=("Maximum characters shown in each before/after preview."), + ) + + action = parser.add_mutually_exclusive_group() + action.add_argument( + "--apply", + action="store_true", + help=( + "Write repairs after creating a private JSONL backup. " + "Without this flag, the command is read-only." + ), + ) + action.add_argument( + "--restore-backup", + metavar="PATH", + help=("Restore original documents from a prior repair backup."), + ) + + parser.add_argument( + "--backup", + metavar="PATH", + help=( + "Backup destination used with --apply. Defaults to a " + "timestamped file beside the palace. Existing files " + "are never overwritten." + ), + ) + return parser + + +def _default_backup_path( + palace_path: str, +) -> Path: + palace = Path(palace_path).expanduser().resolve() + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + + return palace.parent / (f"{palace.name}.encoding-repair-{timestamp}.jsonl") + + +def _preview( + text: str, + limit: int, +) -> str: + compact = text.replace( + "\r", + "\\r", + ).replace( + "\n", + "\\n", + ) + + if len(compact) <= limit: + return compact + + return ( + compact[ + : max( + 0, + limit - 1, + ) + ] + + "…" + ) + + +def _print_change( + drawer_id: str, + before: str, + after: str, + *, + preview_chars: int, +) -> None: + print() + print(f"Drawer: {drawer_id}") + print( + " before: " + + _preview( + before, + preview_chars, + ) + ) + print( + " after: " + + _preview( + after, + preview_chars, + ) + ) + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + + if args.page_size < 1: + parser.error("--page-size must be at least 1") + + if args.preview_chars < 40: + parser.error("--preview-chars must be at least 40") + + if args.backup and not args.apply: + parser.error("--backup requires --apply") + + config = MempalaceConfig() + palace_path = args.palace or config.palace_path + collection_name = args.collection or getattr( + config, + "collection_name", + "mempalace_drawers", + ) + + if args.restore_backup: + with mine_palace_lock(palace_path): + collection = get_collection( + palace_path, + collection_name=(collection_name), + create=False, + ) + report = restore_collection( + collection, + args.restore_backup, + batch_size=(args.page_size), + ) + + print("Mode: RESTORE") + print(f"Backup records validated: {report['validated']}") + print(f"Documents restored: {report['restored']}") + return 0 + + backup_path = None + + if args.apply: + backup_path = Path(args.backup) if args.backup else _default_backup_path(palace_path) + + def show_change( + drawer_id: str, + before: str, + after: str, + ) -> None: + _print_change( + drawer_id, + before, + after, + preview_chars=(args.preview_chars), + ) + + try: + if args.apply: + with mine_palace_lock(palace_path): + collection = get_collection( + palace_path, + collection_name=(collection_name), + create=False, + ) + report = repair_collection( + collection, + apply=True, + page_size=(args.page_size), + backup_path=(backup_path), + on_change=show_change, + ) + else: + collection = get_collection( + palace_path, + collection_name=(collection_name), + create=False, + ) + report = repair_collection( + collection, + apply=False, + page_size=(args.page_size), + on_change=show_change, + ) + except FileExistsError as exc: + print( + f"ERROR: backup file already exists; refusing to overwrite it: {exc.filename}", + file=sys.stderr, + ) + return 2 + + print() + print("Mode: " + ("APPLY" if args.apply else "DRY RUN")) + print(f"Rows scanned: {report['scanned']}") + print(f"Documents needing repair: {report['changed']}") + print(f"Documents updated: {report['updated']}") + + if report["backup_path"]: + print(f"Original-document backup: {report['backup_path']}") + + if not args.apply and report["changed"]: + print() + print( + "Review every change above, then run again with " + "--apply. A private, non-overwriting backup will " + "be written before any document is updated." + ) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_encoding_repair.py b/tests/test_encoding_repair.py new file mode 100644 index 0000000000..df82818fec --- /dev/null +++ b/tests/test_encoding_repair.py @@ -0,0 +1,790 @@ +import json +import os +import stat + +import pytest + +from mempalace.encoding_repair import ( + repair_collection, + repair_mojibake, + restore_collection, +) + + +CLEAN_MULTILINGUAL = [ + "La canción «PERÚ» abre el disco.", + "Buried surface area 1250 Ų.", + "Volume measured as 42 ų.", + "CAFÉ® is a registered mark.", + "RÉSUMÉ\u00a0: présentation générale.", + "Already clean: café → ✅", + "Été à Noël — déjà vu.", + "L'œuvre d'André coûte 20 €.", + "Größe, Fußgänger und Straße.", + "Übermäßig süß — Öl und Äpfel.", + "A ação começa em São João.", + "Às vezes, o avô lê o jornal.", + "CORAÇÃO, PERÚ e CAFÉ®.", + "Zażółć gęślą jaźń.", + "Łódź — źródło wiedzy.", + "Średnica wynosi 25 µm.", + "L'Àngels diu: «això és català».", + "Per què l'aviació és útil?", + "Temperatura: −5 °C ± 0,2 °C.", + "Trademark™ and registered® symbols.", + "Crème brûlée — déjà vu.", + "São Tomé e Príncipe.", + "François parle à Élise.", + "Smörgåsbord, Ångström and Øresund.", + "naïve façade coöperate.", + "România, când și până.", + "Guðrún lives in Reykjavík.", + "Clean emoji: → ✅ 🚀.", +] + + +DAMAGED_CASES = [ + ( + "café", + "café", + ), + ( + "naïve", + "naïve", + ), + ( + "España", + "España", + ), + ( + "ação", + "ação", + ), + ( + "München", + "München", + ), + ( + "Français", + "Français", + ), + ( + "Plan → result — done.", + "Plan → result — done.", + ), + ( + "Copyright © 2026", + "Copyright © 2026", + ), + ( + "BOM removed", + "BOM \ufeffremoved", + ), + ( + "Emoji 😀", + "Emoji 😀", + ), + ( + "café", + "café", + ), + ( + "Clean prefix, café, clean suffix.", + "Clean prefix, café, clean suffix.", + ), +] + + +AMBIGUOUS_CASES = [ + "Å‚", + "ź", + "Ä™", +] + + +@pytest.mark.parametrize( + "text", + CLEAN_MULTILINGUAL, +) +def test_preserves_clean_multilingual_text( + text, +): + assert repair_mojibake(text) == text + + +@pytest.mark.parametrize( + ( + "damaged", + "expected", + ), + DAMAGED_CASES, +) +def test_repairs_high_confidence_mojibake( + damaged, + expected, +): + assert repair_mojibake(damaged) == expected + + +@pytest.mark.parametrize( + "text", + AMBIGUOUS_CASES, +) +def test_leaves_ambiguous_sequences_for_manual_review( + text, +): + assert repair_mojibake(text) == text + + +def test_repair_is_idempotent(): + repaired = repair_mojibake("café → done") + + assert repair_mojibake(repaired) == repaired + + +def test_rejects_invalid_max_passes(): + with pytest.raises( + ValueError, + match="max_passes", + ): + repair_mojibake( + "café", + max_passes=0, + ) + + +class FakeCollection: + name = "mempalace_drawers" + + def __init__( + self, + documents, + ): + self.ids = [f"drawer-{index}" for index in range(len(documents))] + self.documents = list(documents) + self.updates = [] + + def get( + self, + *, + limit, + offset, + include, + ): + del include + + end = offset + limit + + return { + "ids": self.ids[offset:end], + "documents": self.documents[offset:end], + } + + def update( + self, + *, + ids, + documents, + ): + self.updates.append( + { + "ids": list(ids), + "documents": list(documents), + } + ) + + positions = {drawer_id: index for index, drawer_id in enumerate(self.ids)} + + for drawer_id, document in zip( + ids, + documents, + ): + self.documents[positions[drawer_id]] = document + + +def test_dry_run_flags_only_damaged_documents(): + collection = FakeCollection( + [ + CLEAN_MULTILINGUAL[0], + "café", + CLEAN_MULTILINGUAL[1], + "arrow →", + ] + ) + changes = [] + + report = repair_collection( + collection, + apply=False, + page_size=2, + on_change=( + lambda drawer_id, before, after: changes.append( + ( + drawer_id, + before, + after, + ) + ) + ), + ) + + assert report == { + "scanned": 4, + "changed": 2, + "updated": 0, + "backup_path": None, + } + assert [change[0] for change in changes] == [ + "drawer-1", + "drawer-3", + ] + assert collection.updates == [] + + +def test_apply_requires_backup_path(): + with pytest.raises( + ValueError, + match="backup_path", + ): + repair_collection( + FakeCollection(["café"]), + apply=True, + ) + + +def test_apply_writes_backup_before_update( + tmp_path, +): + backup = tmp_path / "backup.jsonl" + + class BackupCheckingCollection(FakeCollection): + def update( + self, + *, + ids, + documents, + ): + lines = backup.read_text(encoding="utf-8").splitlines() + + assert len(lines) == 2 + assert json.loads(lines[1]) == { + "id": "drawer-0", + "original_document": ("café"), + } + + super().update( + ids=ids, + documents=documents, + ) + + collection = BackupCheckingCollection(["café"]) + + report = repair_collection( + collection, + apply=True, + backup_path=backup, + ) + + assert report["updated"] == 1 + assert report["backup_path"] == str(backup) + assert collection.documents == ["café"] + + if os.name != "nt": + mode = stat.S_IMODE(backup.stat().st_mode) + assert mode & 0o077 == 0 + + +def test_apply_refuses_to_overwrite_existing_backup( + tmp_path, +): + backup = tmp_path / "backup.jsonl" + backup.write_text( + "do not overwrite", + encoding="utf-8", + ) + collection = FakeCollection(["café"]) + + with pytest.raises(FileExistsError): + repair_collection( + collection, + apply=True, + backup_path=backup, + ) + + assert backup.read_text(encoding="utf-8") == "do not overwrite" + assert collection.updates == [] + + +def test_apply_with_no_changes_does_not_create_empty_backup( + tmp_path, +): + backup = tmp_path / "backup.jsonl" + + report = repair_collection( + FakeCollection(CLEAN_MULTILINGUAL[:3]), + apply=True, + backup_path=backup, + ) + + assert report["changed"] == 0 + assert report["updated"] == 0 + assert report["backup_path"] is None + assert not backup.exists() + + +def test_backup_restores_original_documents( + tmp_path, +): + backup = tmp_path / "backup.jsonl" + collection = FakeCollection( + [ + "café", + CLEAN_MULTILINGUAL[1], + "arrow →", + ] + ) + + repair_collection( + collection, + apply=True, + page_size=2, + backup_path=backup, + ) + + assert collection.documents == [ + "café", + CLEAN_MULTILINGUAL[1], + "arrow →", + ] + + report = restore_collection( + collection, + backup, + batch_size=1, + ) + + assert report == { + "validated": 2, + "restored": 2, + } + assert collection.documents == [ + "café", + CLEAN_MULTILINGUAL[1], + "arrow →", + ] + + +def test_restore_validates_whole_backup_before_writing( + tmp_path, +): + backup = tmp_path / "backup.jsonl" + backup.write_text( + ( + '{"format":' + '"mempalace-encoding-repair",' + '"version":1}\n' + '{"id":"drawer-0",' + '"original_document":"café"}\n' + "not-json\n" + ), + encoding="utf-8", + ) + collection = FakeCollection(["café"]) + + with pytest.raises( + ValueError, + match="line 3", + ): + restore_collection( + collection, + backup, + ) + + assert collection.updates == [] + + +def test_collection_rejects_misaligned_results(): + class MisalignedCollection(FakeCollection): + def get( + self, + *, + limit, + offset, + include, + ): + del ( + limit, + offset, + include, + ) + + return { + "ids": ["drawer-0"], + "documents": [], + } + + with pytest.raises( + RuntimeError, + match="misaligned", + ): + repair_collection(MisalignedCollection([])) + + +def test_real_chromadb_repair_path_preserves_review_cases( + tmp_path, +): + from mempalace.palace import ( + get_collection, + ) + + palace_path = str(tmp_path / "palace") + collection = get_collection(palace_path) + + originals = { + "clean-spanish": ("La canción «PERÚ» abre el disco."), + "clean-scientific": ("Buried surface area 1250 Ų."), + "clean-trademark": ("CAFÉ® is a registered mark."), + "clean-french": ("RÉSUMÉ\u00a0: présentation générale."), + "damaged-accent": ("España y café."), + "damaged-punctuation": ("Plan → result — done."), + } + + expected = dict(originals) + expected["damaged-accent"] = "España y café." + expected["damaged-punctuation"] = "Plan → result — done." + + collection.upsert( + ids=list(originals), + documents=list(originals.values()), + ) + + changed_ids = [] + + dry_run = repair_collection( + collection, + apply=False, + page_size=2, + on_change=(lambda drawer_id, _before, _after: changed_ids.append(drawer_id)), + ) + + assert dry_run["changed"] == 2 + assert set(changed_ids) == { + "damaged-accent", + "damaged-punctuation", + } + + backup = tmp_path / "originals.jsonl" + + applied = repair_collection( + collection, + apply=True, + page_size=2, + backup_path=backup, + ) + + assert applied["updated"] == 2 + + result = collection.get( + ids=list(originals), + include=["documents"], + ) + by_id = dict( + zip( + result["ids"], + result["documents"], + ) + ) + + assert by_id == expected + + restored = restore_collection( + collection, + backup, + batch_size=1, + ) + + assert restored == { + "validated": 2, + "restored": 2, + } + + result = collection.get( + ids=list(originals), + include=["documents"], + ) + by_id = dict( + zip( + result["ids"], + result["documents"], + ) + ) + + assert by_id == originals + + +def test_backup_header_resolves_wrapped_chroma_collection_name( + tmp_path, +): + from mempalace.backends.chroma import ( + ChromaCollection, + ) + + raw = FakeCollection(["café"]) + raw.name = "mempalace_drawers" + + wrapped = ChromaCollection(raw) + backup = tmp_path / "wrapped-backup.jsonl" + + report = repair_collection( + wrapped, + apply=True, + backup_path=backup, + ) + + lines = [ + json.loads(line) for line in backup.read_text(encoding="utf-8").splitlines() if line.strip() + ] + + assert report["updated"] == 1 + assert lines[0] == { + "collection": "mempalace_drawers", + "format": "mempalace-encoding-repair", + "version": 1, + } + assert lines[1] == { + "id": "drawer-0", + "original_document": "café", + } + + +def test_restore_rejects_backup_for_another_collection( + tmp_path, +): + backup = tmp_path / "wrong-collection.jsonl" + backup.write_text( + ( + '{"collection":"source_collection",' + '"format":"mempalace-encoding-repair",' + '"version":1}\n' + '{"id":"drawer-0",' + '"original_document":"café"}\n' + ), + encoding="utf-8", + ) + + collection = FakeCollection(["café"]) + collection.name = "different_collection" + + with pytest.raises( + ValueError, + match="source_collection", + ): + restore_collection( + collection, + backup, + ) + + assert collection.updates == [] + + +UNDEFINED_CP1252_CONTINUATION_CASES = [ + ( + "Ã\x81", + "Á", + ), + ( + "Ã\x8d", + "Í", + ), + ( + "Ã\x8f", + "Ï", + ), + ( + "Ã\x90", + "Ð", + ), + ( + "Ã\x9d", + "Ý", + ), + ( + ("Ã\x81LVARO vive en PARÃ\x8dS. Ã\x8dNDICE: página 12."), + ("ÁLVARO vive en PARÍS. ÍNDICE: página 12."), + ), + ( + "Dijo “holaâ€\x9d y se fue.", + "Dijo “hola” y se fue.", + ), +] + + +def _undefined_cp1252_review_rows(): + originals = { + "spanish-controls": ("Ã\x81LVARO vive en PARÃ\x8dS. Ã\x8dNDICE: página 12."), + "all-five-controls": ("Valores: Ã\x81 Ã\x8d Ã\x8f Ã\x90 Ã\x9d."), + "curly-quotes": ("Dijo “holaâ€\x9d y se fue."), + "mixed-damage": ("Texto mixto: café, flecha → y PARÃ\x8dS."), + } + + expected = { + "spanish-controls": ("ÁLVARO vive en PARÍS. ÍNDICE: página 12."), + "all-five-controls": ("Valores: Á Í Ï Ð Ý."), + "curly-quotes": ("Dijo “hola” y se fue."), + "mixed-damage": ("Texto mixto: café, flecha → y PARÍS."), + } + + return originals, expected + + +@pytest.mark.parametrize( + ( + "damaged", + "expected", + ), + UNDEFINED_CP1252_CONTINUATION_CASES, +) +def test_repairs_undefined_cp1252_continuation_bytes( + damaged, + expected, +): + assert repair_mojibake(damaged) == expected + + +@pytest.mark.parametrize( + "text", + [ + ("ÁLVARO vive en PARÍS. ÍNDICE: página 12."), + "Dijo “hola” y se fue.", + ], +) +def test_clean_undefined_cp1252_outputs_remain_unchanged( + text, +): + assert repair_mojibake(text) == text + + +def test_apply_completes_undefined_cp1252_rows_in_one_pass( + tmp_path, +): + originals, expected = _undefined_cp1252_review_rows() + collection = FakeCollection(list(originals.values())) + backup = tmp_path / "undefined-controls.jsonl" + + applied = repair_collection( + collection, + apply=True, + page_size=2, + backup_path=backup, + ) + + assert applied["scanned"] == 4 + assert applied["changed"] == 4 + assert applied["updated"] == 4 + assert collection.documents == list(expected.values()) + + second_run = repair_collection( + collection, + apply=False, + page_size=2, + ) + + assert second_run["scanned"] == 4 + assert second_run["changed"] == 0 + assert second_run["updated"] == 0 + + undefined_controls = { + 0x81, + 0x8D, + 0x8F, + 0x90, + 0x9D, + } + + assert all( + not any(ord(character) in undefined_controls for character in document) + for document in collection.documents + ) + + +def test_real_chromadb_completes_undefined_cp1252_rows_in_one_pass( + tmp_path, +): + from mempalace.palace import ( + get_backend_for_palace, + get_collection, + ) + + palace_path = str(tmp_path / "palace") + originals, expected = _undefined_cp1252_review_rows() + + try: + collection = get_collection(palace_path) + + collection.upsert( + ids=list(originals), + documents=list(originals.values()), + ) + + changed_ids = [] + + dry_run = repair_collection( + collection, + apply=False, + page_size=2, + on_change=(lambda drawer_id, _before, _after: changed_ids.append(drawer_id)), + ) + + assert dry_run["scanned"] == 4 + assert dry_run["changed"] == 4 + assert dry_run["updated"] == 0 + assert set(changed_ids) == set(originals) + + backup = tmp_path / "real-undefined-controls.jsonl" + + applied = repair_collection( + collection, + apply=True, + page_size=2, + backup_path=backup, + ) + + assert applied["scanned"] == 4 + assert applied["changed"] == 4 + assert applied["updated"] == 4 + + result = collection.get( + ids=list(originals), + include=["documents"], + ) + by_id = dict( + zip( + result["ids"], + result["documents"], + ) + ) + + assert by_id == expected + + second_run = repair_collection( + collection, + apply=False, + page_size=2, + ) + + assert second_run["scanned"] == 4 + assert second_run["changed"] == 0 + assert second_run["updated"] == 0 + finally: + try: + backend = get_backend_for_palace(palace_path) + close_palace = getattr( + backend, + "close_palace", + None, + ) + + if callable(close_palace): + close_palace(palace_path) + except Exception: + pass