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():

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 # nosemgrep: semgrep.rules.path-traversal-risk, reason: tmp is created by this call's mkstemp beside the fixed-suffix sidecar derived from a non-symlink legacy pickle discovered under the operator cache root

253 os.replace(tmp, sidecar)

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

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

256 try:

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

258 try:

259 os.fsync(dir_fd)

260 finally:

261 os.close(dir_fd)

262 except OSError:

263 pass

264 finally:

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

266 try:

267 # nosemgrep: semgrep.rules.path-traversal-risk, reason: tmp is this sidecar write's mkstemp file; final cleanup removes only that generated temporary name

268 tmp.unlink()

269 except OSError:

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

271

272

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

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

275

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

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

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

279 that collection will fall back to a full reindex.

280 """

281 try:

282 raw = load_index_to_docstore_id(str(pkl_path))

283 except Exception as exc:

284 logger.warning(

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

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

287 )

288 return False

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

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

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

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

293 if not isinstance(raw, dict):

294 logger.warning(

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

296 "that collection will reindex"

297 )

298 return False

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

314 if (

315 not isinstance(pos, int)

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

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

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

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

320 or not isinstance(key, str)

321 or not _ID_RE.fullmatch(key)

322 ):

323 logger.warning(

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

325 "that collection will reindex"

326 )

327 return False

328 id_map[str(pos)] = key

329 try:

330 _write_sidecar_atomic(_sidecar_for(pkl_path), id_map)

331 except OSError:

332 logger.exception(

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

334 "collection will reindex"

335 )

336 return False

337 return True

338

339

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

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

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

343

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

345 "remaining", "scan_errors", "hardening_errors"}``. Re-scans afterward and

346 logs a prominent error recommending manual deletion if any docstore

347 remains — plaintext must never pass silently. ``hardening_errors`` is a

348 0/1 flag (not a count): 1 when the cache root could not be restricted to

349 ``0o700``, reported on its own ERROR line independent of the plaintext

350 outcome. Idempotent (a clean tree is a no-op).

351 """

352 root = _rag_cache_root()

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

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

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

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

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

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

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

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

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

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

363 if root.is_symlink():

364 logger.error(

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

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

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

368 )

369 return {

370 "found": 0,

371 "extracted": 0,

372 "deleted": 0,

373 "reindex_fallback": 0,

374 "remaining": 0,

375 "scan_errors": 1,

376 "hardening_errors": 0,

377 }

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

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

380 # leaf-only chmod in _write_sidecar_atomic and _get_index_path otherwise

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

382 # startup and the first index op.

383 hardening_errors = 0

384 if root.exists():

385 try:

386 root.chmod(0o700)

387 except OSError:

388 hardening_errors = 1

389 # Logged here, at the point of failure and at ERROR, so the

390 # severity is the same whether or not any legacy docstore

391 # happens to exist below. Kept as its own line rather than

392 # folded into the plaintext-migration summary: a chmod failure

393 # (e.g. on a mount that doesn't support it) must not read as

394 # "plaintext may remain", which is a different incident.

395 logger.exception(

396 f"RAG cache: root {root} could not be restricted to 0o700 — "

397 "other local OS accounts may still traverse it. Fix the "

398 "directory permissions manually."

399 )

400 found, scan_errors = _find_legacy_docstores(root)

401 if not found and not scan_errors:

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

403 return {

404 "found": 0,

405 "extracted": 0,

406 "deleted": 0,

407 "reindex_fallback": 0,

408 "remaining": 0,

409 "scan_errors": 0,

410 "hardening_errors": hardening_errors,

411 }

412

413 extracted = 0

414 deleted = 0

415 reindex_fallback = 0

416 for pkl in found:

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

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

419 if _is_live_pkl(pkl):

420 if _extract_map(pkl):

421 extracted += 1

422 else:

423 reindex_fallback += 1

424 try:

425 _delete_file(pkl)

426 deleted += 1

427 except OSError:

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

429

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

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

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

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

434 # unscannable dir as an unresolved leftover.

435 remaining, rescan_errors = _find_legacy_docstores(root)

436 if remaining or rescan_errors:

437 parts = []

438 if remaining:

439 parts.append(

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

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

442 )

443 if rescan_errors:

444 parts.append(

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

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

447 + ", ".join(rescan_errors)

448 )

449 logger.error(

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

451 + "; ".join(parts)

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

453 )

454 else:

455 logger.info(

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

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

458 + (

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

460 if reindex_fallback

461 else ""

462 )

463 )

464 return {

465 "found": len(found),

466 "extracted": extracted,

467 "deleted": deleted,

468 "reindex_fallback": reindex_fallback,

469 "remaining": len(remaining),

470 "scan_errors": len(rescan_errors),

471 "hardening_errors": hardening_errors,

472 }

473

474

475def rekey_index_file(

476 faiss_path: Path,

477 sidecar_path: Path,

478 uuid_to_id: Mapping[str, int],

479 *,

480 dimension: int,

481 index_type: str,

482 metric: str,

483 normalize: bool,

484) -> Dict[str, int]:

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

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

487

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

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

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

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

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

493

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

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

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

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

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

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

500 """

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

502 from .implementations.faiss_store import FaissVectorStore

503

504 faiss_path = Path(faiss_path)

505 sidecar_path = Path(sidecar_path)

506

507 old = read_index(str(faiss_path))

508

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

524 # would lose that retry signal.

525 if isinstance(old, IndexIDMap2):

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

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

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

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

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

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

532 raise ValueError(

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

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

535 )

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

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

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

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

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

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

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

543 try:

544 # nosemgrep: semgrep.rules.path-traversal-risk, reason: production caller rekey_user_indexes derives this sidecar from LibraryRAGService's generated index path in the authenticated user's DB; requests cannot assign that path

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

546 except FileNotFoundError:

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

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

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

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

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

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

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

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

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

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

557 # success rather than letting the FileNotFoundError bubble up and

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

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

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

561 logger.info(

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

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

564 )

565 return {

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

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

568 "orphaned": 0,

569 }

570 expected_ids = {

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

572 }

573 on_disk_id_list = (

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

575 if old.ntotal

576 else []

577 )

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

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

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

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

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

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

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

585 raise ValueError(

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

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

588 "quarantine + reindex"

589 )

590 if set(on_disk_id_list) != expected_ids:

591 raise ValueError(

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

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

594 "index; quarantine + reindex"

595 )

596 logger.info(

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

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

599 )

600 return {

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

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

603 "orphaned": 0,

604 }

605

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

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

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

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

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

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

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

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

614 if isinstance(old, IndexIDMap):

615 raise ValueError(

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

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

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

619 )

620

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

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

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

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

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

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

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

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

629 # + reindexes.

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

631 raise ValueError(

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

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

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

635 )

636

637 ntotal = int(old.ntotal)

638 if old.d != dimension:

639 raise ValueError(

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

641 )

642

643 try:

644 # nosemgrep: semgrep.rules.path-traversal-risk, reason: production caller rekey_user_indexes derives this sidecar from LibraryRAGService's generated index path in the authenticated user's DB; requests cannot assign that path

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

646 except FileNotFoundError:

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

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

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

650 # IndexIDMap2 and deleting the sidecar as its completion signal

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

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

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

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

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

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

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

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

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

660 now_on_disk = None

661 try:

662 now_on_disk = read_index(str(faiss_path))

663 except Exception:

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

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

666 # original FileNotFoundError below.

667 logger.debug(

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

669 )

670 if isinstance(now_on_disk, IndexIDMap2) and now_on_disk.d == dimension:

671 logger.info(

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

673 "migrated it (sidecar already removed)"

674 )

675 return {

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

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

678 "orphaned": 0,

679 }

680 raise

681

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

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

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

685 if ntotal:

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

687 else:

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

689

690 kept_ids: List[int] = []

691 kept_rows: List[int] = []

692 orphaned = 0

693 for pos in range(ntotal):

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

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

696 if int_id is None:

697 orphaned += 1

698 continue

699 kept_ids.append(int(int_id))

700 kept_rows.append(pos)

701

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

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

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

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

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

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

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

709 raise ValueError(

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

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

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

713 )

714

715 kept_vecs = (

716 vectors[kept_rows]

717 if kept_rows

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

719 )

720

721 store = FaissVectorStore.create(

722 dimension=dimension,

723 index_type=index_type,

724 metric=metric,

725 normalize=normalize,

726 )

727 if kept_ids:

728 store.add(kept_ids, kept_vecs)

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

730

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

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

733 verify = FaissVectorStore.load(

734 faiss_path,

735 dimension=dimension,

736 index_type=index_type,

737 metric=metric,

738 normalize=normalize,

739 )

740 if set(verify.live_ids()) != set(kept_ids):

741 raise RuntimeError(

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

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

744 )

745

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

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

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

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

750 logger.info(

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

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

753 )

754 return {

755 "reconstructed": ntotal,

756 "kept": len(kept_ids),

757 "orphaned": orphaned,

758 }

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

240 statements