Coverage for src/local_deep_research/vector_stores/legacy

1"""One-time migration of legacy RAG index docstore files off disk.

3The previous RAG index format persisted a companion ``<hash>.pkl`` docstore next

4to each ``<hash>.faiss``. That docstore held the chunk **text** (plaintext at

5rest) plus the ``index_to_docstore_id`` map (FAISS position -> uuid) that the

6index needs. The current format stores no text in the index — text lives only in

7the encrypted DB and is rehydrated by id.

9This runs at startup (pure filesystem, no login / DB access). For each legacy

10``.pkl`` it:

12 1. extracts ONLY the ``index_to_docstore_id`` map (position -> uuid) into a

13 small, **text-free** ``<hash>.idmap.json`` sidecar (plain JSON — no pickle,

14 no text, no metadata), then

15 2. deletes the ``.pkl`` — removing all plaintext from disk.

17The per-user, login-time re-key (elsewhere) later consumes the sidecar to

18rebuild the index keyed by ``DocumentChunk.id`` (no re-embedding), then deletes

19the sidecar. If a ``.pkl`` can't be read (corrupt), it is deleted anyway

20(plaintext removed) and that collection falls back to a normal reindex.

22IMPORTANT: this MUST ship in the same release as the vector-store cutover.

23Before the cutover the old code still reads the ``.pkl``'s text for snippets, so

24removing it early would break search — do NOT wire this into startup until the

25old FAISS path is gone. It is exposed as a plain function so it can be called

26from tests and from the cutover's startup hook.

27"""

29import json

30import os

31import re

32import stat

33import tempfile

34import time

35from pathlib import Path

36from typing import Dict, List, Mapping, Tuple

38import numpy as np

39from faiss import (

40 IndexFlatIP,

41 IndexFlatL2,

42 IndexHNSWFlat,

43 IndexIDMap,

44 IndexIDMap2,

45 read_index,

46 vector_to_array,

47)

48from loguru import logger

50from ..config.paths import get_cache_directory

51from ..research_library.services.faiss_safe_load import (

52 load_index_to_docstore_id,

53)

55_RAG_CACHE_SUBDIR = "rag_indices"

56_SIDECAR_SUFFIX = ".idmap.json"

57# Only sweep sidecar ``.tmp`` files older than this — long enough that no single

58# write is still in flight, so a concurrent (multi-process) phase-1 run's temp is

59# never mistaken for an orphan. Mirrors faiss_store._STALE_TMP_AGE_SECONDS.

60_STALE_TMP_AGE_SECONDS = 3600

62# A legitimate docstore key is a ``uuid.uuid4().hex`` (32 hex chars, what

63# _store_chunks_to_db writes); accept the dashed uuid form too. Anything else in

64# a (tampered) ``.pkl`` — arbitrary text, whitespace — is rejected so it can

65# never be smuggled into the "text-free" sidecar.

66_ID_RE = re.compile(

67 r"[0-9a-f]{32}"

68 r"|[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}",

69 re.IGNORECASE,

70)

71# Anchored so a crafted name like ``x.pkl.corrupt-1.idmap.json`` is NOT matched.

72_QUARANTINE_RE = re.compile(r"\.pkl\.corrupt-[0-9-]+$")

75def _rag_cache_root() -> Path:

76 return get_cache_directory() / _RAG_CACHE_SUBDIR

79def _is_live_pkl(path: Path) -> bool:

80 """A live legacy docstore ``<hash>.pkl`` (has a re-keyable sibling index)."""

81 return path.is_file() and path.name.endswith(".pkl")

84def _is_quarantined_pkl(path: Path) -> bool:

85 """A quarantined docstore ``<hash>.pkl.corrupt-<ns>`` (dead, no sidecar)."""

86 return path.is_file() and _QUARANTINE_RE.search(path.name) is not None

89def _sidecar_for(pkl_path: Path) -> Path:

90 # <hash>.pkl -> <hash>.idmap.json

91 return pkl_path.with_name(pkl_path.stem + _SIDECAR_SUFFIX)

94def _find_legacy_docstores(root: Path) -> Tuple[List[Path], List[str]]:

95 """Return ``(found_docstores, unscannable_dirs)``.

97 A non-empty ``unscannable_dirs`` means ``os.walk`` could not read a

98 directory, so a plaintext ``.pkl`` inside it may survive UNDETECTED — the

99 caller must treat that as an INCOMPLETE migration, not clean success (a

100 re-scan is blind to the very same unreadable directory)."""

101 if not root.exists(): 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true

102 return [], []

103

104 scan_errors: List[str] = []

105

106 def _onerror(exc: OSError) -> None:

107 # os.walk (like rglob) SILENTLY skips a directory it can't scan. This

108 # migration's whole guarantee is "no plaintext .pkl left on disk", so a

109 # scan error must be LOUD and must fail the completeness check —

110 # otherwise an unreadable subdir could hide an un-deleted .pkl while the

111 # run reports clean success.

112 target = getattr(exc, "filename", "?")

113 scan_errors.append(str(target))

114 logger.error(

115 "RAG plaintext migration could not scan "

116 f"{target} ({exc}); a legacy .pkl there may "

117 "NOT have been removed — fix permissions and re-run, or delete it "

118 "manually."

119 )

120

121 found: List[Path] = []

122 for dirpath, dirs, files in os.walk(str(root), onerror=_onerror):

123 # os.walk does NOT descend into symlinked subdirectories

124 # (followlinks=False) and does NOT call onerror for them — they are

125 # silently skipped. A legacy .pkl inside one would survive undetected,

126 # so flag each symlinked subdir as unscanned (same as a permission

127 # error) rather than miss it. followlinks=True is avoided deliberately:

128 # a symlink cycle would hang the migration.

129 for d in dirs:

130 sub = Path(dirpath) / d

131 # A stat error (PermissionError, etc.) on a SINGLE entry must never

132 # abort the whole walk — that would leave EVERY .pkl on disk while

133 # the caller logs "migration failed". Flag the entry and keep going.

134 try:

135 is_link = sub.is_symlink()

136 except OSError:

137 scan_errors.append(str(sub))

138 logger.exception(

139 f"RAG plaintext migration could not stat directory {sub}; "

140 "a legacy .pkl there may NOT have been removed."

141 )

142 continue

143 if is_link:

144 scan_errors.append(str(sub))

145 logger.error(

146 "RAG plaintext migration did not descend into symlinked "

147 f"directory {sub}; a legacy .pkl inside may NOT have been "

148 "removed — replace the symlink with a real directory and "

149 "re-run, or delete the .pkl manually."

150 )

151 for name in files:

152 p = Path(dirpath) / name

153 try:

154 is_pkl = _is_live_pkl(p) or _is_quarantined_pkl(p)

155 except OSError:

156 scan_errors.append(str(p))

157 logger.exception(

158 f"RAG plaintext migration could not stat file {p}; "

159 "a legacy .pkl there may NOT have been removed."

160 )

161 continue

162 if is_pkl:

163 # A symlinked .pkl is a trap: _is_*_pkl / is_file() follow the

164 # link, so it looks migratable, but _delete_file's unlink()

165 # removes only the symlink — the real plaintext file it points

166 # to (potentially OUTSIDE the cache root) survives untouched

167 # while the run logs clean success. Don't follow it: flag it as

168 # unresolved (incomplete migration) so an operator removes the

169 # real file. We deliberately do NOT auto-delete the target — it

170 # can be an arbitrary path we must not unlink blindly. A stat

171 # error here is treated as "can't tell" -> unresolved.

172 try:

173 is_link = p.is_symlink()

174 except OSError:

175 is_link = True

176 if is_link:

177 scan_errors.append(str(p))

178 logger.error(

179 f"RAG plaintext migration found a SYMLINKED (or "

180 f"un-statable) docstore {p}; deleting the link would "

181 "leave the real plaintext file intact. Remove the real "

182 "target manually, then re-run."

183 )

184 continue

185 found.append(p)

186 return sorted(found), scan_errors

187

188

189def _delete_file(path: Path) -> None:

190 """Delete ``path`` cross-platform (Windows + POSIX).

191

192 Windows raises ``PermissionError`` when a file carries the read-only

193 attribute; clear it and retry once. On POSIX ``unlink`` succeeds even while a

194 file is open. A still-failing delete propagates to the caller, which records

195 it — the re-scan then reports it.

196 """

197 try:

198 path.unlink()

199 return

200 except FileNotFoundError:

201 return

202 except PermissionError:

203 pass

204 try:

205 os.chmod(path, stat.S_IWRITE | stat.S_IREAD)

206 except OSError:

207 pass

208 path.unlink()

209

210

211def _write_sidecar_atomic(sidecar: Path, id_map: Dict[str, str]) -> None:

212 """Atomically write the text-free position->uuid map as JSON."""

213 # Lazy import: this module runs at startup, before the app is fully

214 # wired up, so avoid a top-level import cycle.

215 from ..security.directory_creation import create_directory

216

217 create_directory(sidecar.parent, context="RAG legacy sidecar directory")

218 try:

219 sidecar.parent.chmod(0o700)

220 except OSError:

221 logger.warning(f"Could not chmod {sidecar.parent} to 0o700")

222 # Sweep stale sidecar temps from a previously hard-killed migration. Phase-1

223 # normally runs single-threaded at startup, but under a multi-process server

224 # (gunicorn -w N) two workers can start it concurrently — only sweep temps

225 # OLDER than a generous threshold so a concurrent worker's just-created,

226 # still-in-flight temp for the same sidecar is never mistaken for an orphan

227 # and unlinked (which would crash its write). Mirrors FaissVectorStore.persist.

228 now = time.time()

229 for stale in sidecar.parent.glob(f"{sidecar.name}.*.tmp"):

230 try:

231 if now - stale.stat().st_mtime < _STALE_TMP_AGE_SECONDS:

232 continue

233 stale.unlink()

234 except OSError:

235 logger.warning(f"Could not remove stale sidecar temp {stale}")

236 # Safe: this writes only the text-free ``.idmap.json`` sidecar — a plain

237 # position->uuid JSON map (see module docstring / `_ID_RE`), never chunk

238 # text/metadata. Not sensitive data at rest.

239 fd, tmp_name = tempfile.mkstemp(

240 dir=str(sidecar.parent), prefix=f"{sidecar.name}.", suffix=".tmp"

241 )

242 tmp = Path(tmp_name)

243 try:

244 # Safe: same text-free id-map as above, written via the fd from

245 # mkstemp — no chunk text/metadata ever passes through this path.

246 with os.fdopen(fd, "w", encoding="utf-8") as f:

247 json.dump(id_map, f)

248 f.flush()

249 os.fsync(

250 f.fileno()

251 ) # durable before the rename (crash-consistency)

252 os.replace(tmp, sidecar)

253 # fsync the directory so the rename itself survives a crash (mirrors

254 # FaissVectorStore.persist); unsupported on some platforms.

255 try:

256 dir_fd = os.open(str(sidecar.parent), os.O_RDONLY)

257 try:

258 os.fsync(dir_fd)

259 finally:

260 os.close(dir_fd)

261 except OSError:

262 pass

263 finally:

264 if tmp.exists(): 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true

265 try:

266 tmp.unlink()

267 except OSError:

268 logger.warning(f"Could not remove temp sidecar {tmp}")

269

270

271def _extract_map(pkl_path: Path) -> bool:

272 """Extract the text-free position->uuid map to a sidecar next to ``pkl_path``.

273

274 Returns True if the sidecar was written (so the login re-key can rebuild

275 without re-embedding). Returns False if the ``.pkl`` couldn't be read — the

276 caller still deletes the ``.pkl`` (plaintext removal is not optional) and

277 that collection will fall back to a full reindex.

278 """

279 try:

280 raw = load_index_to_docstore_id(str(pkl_path))

281 except Exception as exc:

282 logger.warning(

283 f"Could not read legacy docstore {pkl_path} "

284 f"({type(exc).__name__}); that collection will reindex"

285 )

286 return False

287 # Validate shape + content. A malformed/tampered map (non-dict, non-int

288 # positions, or non-uuid values that could smuggle text into the "text-free"

289 # sidecar) falls to the reindex fallback — it must never crash the whole

290 # migration (a DoS leaving other users' plaintext on disk) nor leak content.

291 if not isinstance(raw, dict):

292 logger.warning(

293 f"Legacy docstore {pkl_path} has an unexpected id-map shape; "

294 "that collection will reindex"

295 )

296 return False

297 id_map: Dict[str, str] = {}

298 for pos, key in raw.items():

299 # Require key to be a real str and match it DIRECTLY (never str(key)).

300 # A tampered .pkl can put an arbitrary unpickled object here; calling

301 # str() on a maliciously deep/shared nested structure is a pickle

302 # "billion laughs" amplification that can hang/OOM the whole startup

303 # migration. A legitimate index_to_docstore_id value is always a uuid

304 # string, so a non-str (or non-uuid) value is simply rejected to the

305 # reindex fallback — it can never reach str().

306 # pos must be a plausible faiss array position. Bounding it also stops

307 # str(pos) below from raising ValueError on a maliciously huge int

308 # (Python's int->str digit limit), which would otherwise escape this

309 # loop uncaught and abort the whole startup migration — leaving other

310 # users' plaintext .pkl files on disk. bool is an int subclass; the

311 # range check harmlessly admits True/False (1/0), rejected downstream.

312 if (

313 not isinstance(pos, int)

314 or isinstance(pos, bool) # bool is an int subclass; True/False

315 # would serialize as "True"/"False" and never match a str(pos)

316 # lookup at rekey time, silently orphaning that chunk's vector.

317 or not (0 <= pos < 2**31)

318 or not isinstance(key, str)

319 or not _ID_RE.fullmatch(key)

320 ):

321 logger.warning(

322 f"Legacy docstore {pkl_path} has an unexpected id-map entry; "

323 "that collection will reindex"

324 )

325 return False

326 id_map[str(pos)] = key

327 try:

328 _write_sidecar_atomic(_sidecar_for(pkl_path), id_map)

329 except OSError:

330 logger.exception(

331 f"Could not write id-map sidecar for {pkl_path}; that "

332 "collection will reindex"

333 )

334 return False

335 return True

336

337

338def migrate_legacy_docstores() -> Dict[str, int]:

339 """Extract each legacy ``.pkl``'s id-map to a text-free sidecar, then delete

340 the ``.pkl``, removing all plaintext docstores from disk.

341

342 Returns ``{"found", "extracted", "deleted", "reindex_fallback",

343 "remaining"}``. Re-scans afterward and logs a prominent error recommending

344 manual deletion if any docstore remains — plaintext must never pass

345 silently. Idempotent (a clean tree is a no-op).

346 """

347 root = _rag_cache_root()

348 # SECURITY: refuse to operate on a SYMLINKED cache root. os.walk (and

349 # Path.chmod below) follow a symlinked TOP-LEVEL path unconditionally —

350 # followlinks=False only guards descent into symlinked CHILDREN found

351 # mid-walk (which _find_legacy_docstores already flags). A symlinked

352 # rag_indices/ would otherwise make this migration chmod 0o700 and

353 # pattern-delete .pkl files OUTSIDE the intended tree (an admin who

354 # symlinked the cache onto a bigger volume, or a hostile account that can

355 # write the cache parent), reporting clean success. Treat it exactly like a

356 # symlinked subdir: log loudly and refuse rather than following it. Uses

357 # is_symlink() (an lstat, does NOT follow) so it is safe to call first.

358 if root.is_symlink():

359 logger.error(

360 f"RAG cache root {root} is a symlink; refusing to migrate through "

361 "it (following it would chmod/delete files outside the intended "

362 "cache tree). Replace it with a real directory and restart."

363 )

364 return {

365 "found": 0,

366 "extracted": 0,

367 "deleted": 0,

368 "reindex_fallback": 0,

369 "remaining": 0,

370 "scan_errors": 1,

371 }

372 # Harden the RAG cache ROOT at the earliest point it is touched so another

373 # local OS account can't traverse into any user's vector caches. Both the

374 # leaf-only chmod in _write_sidecar_atomic and _get_index_path otherwise

375 # leave the root at the process umask (typically world-traversable) between

376 # startup and the first index op.

377 if root.exists(): 377 ↛ 382line 377 didn't jump to line 382 because the condition on line 377 was always true

378 try:

379 root.chmod(0o700)

380 except OSError:

381 logger.warning(f"Could not chmod RAG cache root {root} to 0o700")

382 found, scan_errors = _find_legacy_docstores(root)

383 if not found and not scan_errors:

384 logger.debug(f"RAG cache: no legacy docstore files under {root}")

385 return {

386 "found": 0,

387 "extracted": 0,

388 "deleted": 0,

389 "reindex_fallback": 0,

390 "remaining": 0,

391 "scan_errors": 0,

392 }

393

394 extracted = 0

395 deleted = 0

396 reindex_fallback = 0

397 for pkl in found:

398 # Live .pkl: preserve the re-key map before deleting. Quarantined

399 # (.pkl.corrupt-*): dead — just delete, no sidecar.

400 if _is_live_pkl(pkl):

401 if _extract_map(pkl):

402 extracted += 1

403 else:

404 reindex_fallback += 1

405 try:

406 _delete_file(pkl)

407 deleted += 1

408 except OSError:

409 logger.exception(f"RAG cache: could not delete {pkl}")

410

411 # Re-scan is the source of truth (partial failures, POSIX unlink-of-open).

412 # A scan error is NOT clean success: an unreadable directory hides its

413 # contents from BOTH the delete pass above and this re-scan, so a plaintext

414 # .pkl inside it could survive while `remaining` reads empty. Treat any

415 # unscannable dir as an unresolved leftover.

416 remaining, rescan_errors = _find_legacy_docstores(root)

417 if remaining or rescan_errors:

418 parts = []

419 if remaining: 419 ↛ 420line 419 didn't jump to line 420 because the condition on line 419 was never true

420 parts.append(

421 f"{len(remaining)} legacy docstore file(s) remain on disk: "

422 + ", ".join(str(p) for p in remaining)

423 )

424 if rescan_errors: 424 ↛ 430line 424 didn't jump to line 430 because the condition on line 424 was always true

425 parts.append(

426 f"{len(rescan_errors)} director(ies) could not be scanned "

427 "(a plaintext .pkl inside may NOT have been removed): "

428 + ", ".join(rescan_errors)

429 )

430 logger.error(

431 "RAG cache: plaintext migration did NOT complete cleanly — "

432 + "; ".join(parts)

433 + ". Fix permissions and re-run, or delete the file(s) manually."

434 )

435 else:

436 logger.info(

437 f"RAG cache: extracted {extracted} id-map(s), removed {deleted} "

438 f"legacy docstore file(s) from {root}"

439 + (

440 f" ({reindex_fallback} unreadable -> will reindex)"

441 if reindex_fallback

442 else ""

443 )

444 )

445 return {

446 "found": len(found),

447 "extracted": extracted,

448 "deleted": deleted,

449 "reindex_fallback": reindex_fallback,

450 "remaining": len(remaining),

451 "scan_errors": len(rescan_errors),

452 }

453

454

455def rekey_index_file(

456 faiss_path: Path,

457 sidecar_path: Path,

458 uuid_to_id: Mapping[str, int],

459 *,

460 dimension: int,

461 index_type: str,

462 metric: str,

463 normalize: bool,

464) -> Dict[str, int]:

465 """Convert one old-format index to the new ``IndexIDMap2`` keyed by int id,

466 WITHOUT re-embedding — the login-time (phase-2) half of the migration.

467

468 Reads the raw old ``.faiss`` (reconstructs each vector by position), maps

469 position -> uuid (from the ``.idmap.json`` sidecar) -> int ``DocumentChunk.id``

470 (via ``uuid_to_id``), builds a fresh ``IndexIDMap2`` keyed by those ints,

471 writes it atomically over ``faiss_path``, and **reads it back and verifies**.

472 Positions whose uuid has no DB row are orphans and are dropped.

473

474 Does NOT delete the sidecar (nor acquire the write lock, nor re-record file

475 integrity) — the caller owns all three. The caller deletes the sidecar ONLY

476 after recording the new file's integrity, so a crash/failure between persist

477 and integrity-record leaves the sidecar as a retry signal; the ``IndexIDMap2``

478 guard above then makes that retry a clean no-op. Raises on a failed read-back

479 verification (the caller quarantines + falls back to a normal reindex).

480 """

481 # Import here to avoid a module-load cycle (implementations import base only).

482 from .implementations.faiss_store import FaissVectorStore

483

484 faiss_path = Path(faiss_path)

485 sidecar_path = Path(sidecar_path)

486

487 old = read_index(str(faiss_path))

488

489 # Crash-consistency guard. rekey persists the new IndexIDMap2 (atomic

490 # replace) and only unlinks the sidecar AFTER a successful read-back verify.

491 # If the process dies in that window, faiss_path is ALREADY the new format

492 # but the sidecar survives, so the next login re-enters here on an

493 # already-migrated file. Running the reconstruct path below on an

494 # IndexIDMap2 is not merely wrong — reconstruct_n(0, ntotal) would treat

495 # 0..ntotal-1 as *ids* to look up; a missing id raises a C++ faiss

496 # exception that SWIG does NOT surface as a catchable Python exception on

497 # this call path, aborting the whole process (and, if the int ids happen to

498 # fall in 0..ntotal-1, silently reconstructing the WRONG vectors instead).

499 # Either way the re-key is already done — report success WITHOUT deleting

500 # the sidecar. The CALLER owns sidecar deletion (only after it has recorded

501 # file integrity), so an interrupted prior run that persisted the new index

502 # but hadn't yet recorded integrity is recoverable: this re-entry succeeds,

503 # the caller (re)records integrity, THEN drops the sidecar. Deleting it here

504 # would lose that retry signal.

505 if isinstance(old, IndexIDMap2):

506 # Same dimension guard the fresh re-key path enforces below: an

507 # already-migrated index whose width no longer matches the configured

508 # embedding model (model changed post-migration) must be rebuilt, not

509 # silently accepted. Raise so the caller quarantines + reindexes now,

510 # rather than deferring to an opaque dimension failure at first search.

511 if old.d != dimension: 511 ↛ 512line 511 didn't jump to line 512 because the condition on line 511 was never true

512 raise ValueError(

513 f"rekey {faiss_path}: already-migrated index dim {old.d} != "

514 f"expected {dimension} (embedding model changed) — reindex"

515 )

516 # Cross-check the on-disk ids against what the CURRENT sidecar + DB

517 # resolution says they should be — the fresh path below does this via a

518 # read-back verify; this short-circuit must too. Without it, a stale or

519 # FOREIGN IndexIDMap2 sitting at this path (e.g. two collections/users

520 # colliding on the pre-per-user-scoping shared cache path, or a retry

521 # after the uuid->id resolution changed) would be silently adopted, and

522 # its baked-in ids would rehydrate the WRONG document's text on search.

523 try:

524 pos_to_uuid = json.loads(sidecar_path.read_text(encoding="utf-8"))

525 except FileNotFoundError:

526 # The sidecar is the caller's COMPLETION SIGNAL: it is deleted only

527 # AFTER a successful re-key records the new index's integrity (see

528 # rekey_user_indexes Phase B). A concurrent worker — another process

529 # (gunicorn -w N / the scheduler) with no cross-process mutex, or a

530 # thread whose sidecar check passed before that worker won the lock —

531 # can finalize and delete it between our caller's sidecar check and

532 # this read. A missing sidecar next to an already-IndexIDMap2 file of

533 # the right dimension therefore means "already finalized by another

534 # worker" — the SAME signature rekey_user_indexes treats as done at

535 # its top-of-loop no-sidecar branch (skipped_no_sidecar). Report

536 # success rather than letting the FileNotFoundError bubble up and

537 # QUARANTINE a healthy, just-rekeyed index (forcing a needless full

538 # re-embed). No id cross-check: there is no sidecar left to check

539 # against, exactly as the caller's no-sidecar branch also skips it.

540 logger.info(

541 f"Re-key skipped for {faiss_path.name}: already migrated to "

542 "IndexIDMap2; its sidecar was removed by a concurrent finalizer"

543 )

544 return {

545 "reconstructed": int(old.ntotal),

546 "kept": int(old.ntotal),

547 "orphaned": 0,

548 }

549 expected_ids = {

550 int(uuid_to_id[u]) for u in pos_to_uuid.values() if u in uuid_to_id

551 }

552 on_disk_id_list = (

553 vector_to_array(old.id_map).astype("int64").tolist()

554 if old.ntotal

555 else []

556 )

557 # Multiplicity, not just membership: a duplicate id (e.g. [5, 5, 6, 7])

558 # has the SAME set as {5, 6, 7}, so the set comparison below cannot see

559 # it (sets dedupe) — exactly the corruption the fresh-build path rejects

560 # at ``len(set(kept_ids)) != len(kept_ids)``. A duplicate baked into an

561 # already-migrated IndexIDMap2 double-counts/mis-scores that chunk in

562 # search, so reject it here too rather than silently adopt it.

563 if len(set(on_disk_id_list)) != len(on_disk_id_list):

564 raise ValueError(

565 f"rekey {faiss_path}: already-migrated index has duplicate "

566 "DocumentChunk.id(s) — refusing to adopt an id-collided index; "

567 "quarantine + reindex"

568 )

569 if set(on_disk_id_list) != expected_ids:

570 raise ValueError(

571 f"rekey {faiss_path}: already-migrated index ids do not match "

572 "the sidecar/DB resolution — refusing to adopt a stale/foreign "

573 "index; quarantine + reindex"

574 )

575 logger.info(

576 f"Re-key skipped for {faiss_path.name}: already migrated to "

577 "IndexIDMap2 (interrupted prior run — caller will finalize)"

578 )

579 return {

580 "reconstructed": int(old.ntotal),

581 "kept": int(old.ntotal),

582 "orphaned": 0,

583 }

584

585 # The crash hazard is the WHOLE IndexIDMap family, not just IndexIDMap2:

586 # reconstruct_n on a plain IndexIDMap (our own format's superclass) also

587 # raises an untranslated C++ faiss exception that aborts the process

588 # (SIGABRT), bypassing every try/except. A plain IndexIDMap is never a

589 # format we write (we only ever persist IndexIDMap2) nor the pre-migration

590 # raw base index we mean to reconstruct here — so it is foreign/corrupt.

591 # Raise a CATCHABLE ValueError so the caller quarantines + reindexes,

592 # instead of falling through to the process-killing reconstruct_n below.

593 if isinstance(old, IndexIDMap):

594 raise ValueError(

595 f"rekey {faiss_path}: on-disk index is {type(old).__name__} "

596 "(IndexIDMap family but not IndexIDMap2) — refusing to reconstruct "

597 "a foreign/corrupt index; quarantine + reindex instead"

598 )

599

600 # Positive whitelist. The IndexIDMap guard above is only one of MANY faiss

601 # classes that abort the process on reconstruct_n: any index type that isn't

602 # a plain reconstruct-safe base index (e.g. a corrupt/tampered/foreign file

603 # that happens to decode as IndexIVFPQFastScan) segfaults at reconstruct_n

604 # below — an untranslated C++ abort that bypasses the caller's try/except

605 # and, on the login-time rekey path, crash-loops the worker. Only the base

606 # classes _build_base_index actually produces are reconstruct-safe here;

607 # reject anything else with a CATCHABLE ValueError so the caller quarantines

608 # + reindexes.

609 if not isinstance(old, (IndexFlatL2, IndexFlatIP, IndexHNSWFlat)):

610 raise ValueError(

611 f"rekey {faiss_path}: on-disk index is {type(old).__name__}, not a "

612 "reconstruct-safe base index (IndexFlatL2/IndexFlatIP/IndexHNSWFlat)"

613 " — refusing a foreign/corrupt index; quarantine + reindex instead"

614 )

615

616 ntotal = int(old.ntotal)

617 if old.d != dimension: 617 ↛ 618line 617 didn't jump to line 618 because the condition on line 617 was never true

618 raise ValueError(

619 f"rekey {faiss_path}: index dim {old.d} != expected {dimension}"

620 )

621

622 try:

623 pos_to_uuid = json.loads(sidecar_path.read_text(encoding="utf-8"))

624 except FileNotFoundError:

625 # A concurrent worker (multi-process: gunicorn -w N / the scheduler,

626 # with no cross-process lock) may have finished re-keying this SAME

627 # index between our read_index() above and this read — persisting a new

628 # IndexIDMap2 and deleting the sidecar as its completion signal

629 # (legacy_rekey Phase B, which persists the new index BEFORE unlinking

630 # the sidecar). Re-read the on-disk file: if it is now an IndexIDMap2 of

631 # the right dimension, the migration is already done, so report success

632 # rather than letting FileNotFoundError bubble up — the caller treats

633 # that as un-rekeyable and QUARANTINES a perfectly good, just-migrated

634 # index (forcing a needless full re-embed). Only if the file is still

635 # NOT migrated (a genuinely missing sidecar next to an old-format file)

636 # do we re-raise so the caller quarantines. This mirrors the

637 # already-migrated short-circuit's own FileNotFoundError handling above.

638 now_on_disk = None

639 try:

640 now_on_disk = read_index(str(faiss_path))

641 except Exception:

642 # Probe only: a failed/partial read just means "not a valid

643 # migrated index yet" — now_on_disk stays None and we re-raise the

644 # original FileNotFoundError below.

645 logger.debug(

646 f"Concurrent-migration re-read probe failed for {faiss_path.name}"

647 )

648 if isinstance(now_on_disk, IndexIDMap2) and now_on_disk.d == dimension: 648 ↛ 658line 648 didn't jump to line 658 because the condition on line 648 was always true

649 logger.info(

650 f"Re-key skipped for {faiss_path.name}: a concurrent worker "

651 "migrated it (sidecar already removed)"

652 )

653 return {

654 "reconstructed": int(now_on_disk.ntotal),

655 "kept": int(now_on_disk.ntotal),

656 "orphaned": 0,

657 }

658 raise

659

660 # Reconstruct ALL vectors by position in one bounded call. reconstruct_n only

661 # touches positions 0..ntotal-1, avoiding the out-of-range garbage that a

662 # hand-rolled reconstruct(i) loop risks (faiss #4413).

663 if ntotal: 663 ↛ 666line 663 didn't jump to line 666 because the condition on line 663 was always true

664 vectors = np.asarray(old.reconstruct_n(0, ntotal), dtype="float32")

665 else:

666 vectors = np.empty((0, dimension), dtype="float32")

667

668 kept_ids: List[int] = []

669 kept_rows: List[int] = []

670 orphaned = 0

671 for pos in range(ntotal):

672 uid = pos_to_uuid.get(str(pos))

673 int_id = uuid_to_id.get(uid) if uid is not None else None

674 if int_id is None:

675 orphaned += 1

676 continue

677 kept_ids.append(int(int_id))

678 kept_rows.append(pos)

679

680 # Two legacy positions must never map to the SAME DocumentChunk.id: that

681 # would bake a duplicate id into the IndexIDMap2 (which tolerates dupes),

682 # producing phantom, mis-scored search hits — and the set()-based read-back

683 # verify below cannot see it (sets dedupe). A duplicate means the sidecar /

684 # uuid->id resolution is inconsistent, so refuse and let the caller

685 # quarantine + reindex cleanly from the DB rather than persist corruption.

686 if len(set(kept_ids)) != len(kept_ids):

687 raise ValueError(

688 f"rekey {faiss_path}: {len(kept_ids) - len(set(kept_ids))} duplicate "

689 "DocumentChunk.id(s) from the legacy id-map — refusing to persist an "

690 "id-collided index; quarantine + reindex instead"

691 )

692

693 kept_vecs = (

694 vectors[kept_rows]

695 if kept_rows

696 else np.empty((0, dimension), dtype="float32")

697 )

698

699 store = FaissVectorStore.create(

700 dimension=dimension,

701 index_type=index_type,

702 metric=metric,

703 normalize=normalize,

704 )

705 if kept_ids: 705 ↛ 707line 705 didn't jump to line 707 because the condition on line 705 was always true

706 store.add(kept_ids, kept_vecs)

707 store.persist(faiss_path) # atomic temp + os.replace

708

709 # Read-back verify BEFORE removing the sidecar (the only remaining re-key

710 # input): the new file must load as IndexIDMap2 with exactly the kept ids.

711 verify = FaissVectorStore.load(

712 faiss_path,

713 dimension=dimension,

714 index_type=index_type,

715 metric=metric,

716 normalize=normalize,

717 )

718 if set(verify.live_ids()) != set(kept_ids): 718 ↛ 719line 718 didn't jump to line 719 because the condition on line 718 was never true

719 raise RuntimeError(

720 f"rekey {faiss_path}: read-back verify failed "

721 f"({verify.count()} ids on disk != {len(kept_ids)} expected)"

722 )

723

724 # NOTE: the sidecar is deliberately NOT deleted here. The caller deletes it

725 # ONLY after it has recorded the new file's integrity, so a crash/failure

726 # between persist and integrity-record leaves the sidecar as a retry signal

727 # (the IndexIDMap2 guard at the top makes the retry a clean no-op re-key).

728 logger.info(

729 f"Re-keyed {faiss_path.name}: {len(kept_ids)} vectors "

730 f"({orphaned} orphaned, {ntotal} reconstructed)"

731 )

732 return {

733 "reconstructed": ntotal,

734 "kept": len(kept_ids),

735 "orphaned": orphaned,

736 }

Coverage for src/local_deep_research/vector_stores/legacy_cleanup.py: 81%

238 statements