diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index e87ca36e..5e738d4f 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -279,11 +279,14 @@ class LocalFileStore(BaseFileStore): ) return if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done(): - if skip_health_check: - pending_rebuild = rebuild or bool( - self._embedding_backfill_pending and self._embedding_backfill_pending[1], - ) - self._embedding_backfill_pending = (True, pending_rebuild) + pending_verified = skip_health_check or bool( + self._embedding_backfill_pending and self._embedding_backfill_pending[0], + ) + pending_rebuild = rebuild or bool( + self._embedding_backfill_pending and self._embedding_backfill_pending[1], + ) + if pending_verified or pending_rebuild: + self._embedding_backfill_pending = (pending_verified, pending_rebuild) self.logger.info( f"{self.name}: embedding backfill scheduling skipped: reason=already_running, " f"elapsed={time.monotonic() - started_at:.3f}s", @@ -355,6 +358,19 @@ class LocalFileStore(BaseFileStore): async def _backfill_missing_embeddings(self, *, skip_health_check: bool = False) -> None: """Background-repair persisted chunks that do not have usable vectors.""" started_at = time.monotonic() + try: + await self._backfill_missing_embeddings_inner(skip_health_check=skip_health_check, started_at=started_at) + finally: + if self._embedding_rebuild_pending and not self._closing: + try: + await self._after_embedding_backfill() + await self.dump() + self._embedding_rebuild_pending = False + except Exception: + self.logger.exception(f"{self.name}: failed to finalize embedding rebuild") + + async def _backfill_missing_embeddings_inner(self, *, skip_health_check: bool, started_at: float) -> None: + """Perform one backfill pass; the caller owns rebuild finalization.""" if not self.embedding_store or not self.file_chunks: self.logger.info( f"{self.name}: embedding backfill finished without work: " @@ -370,10 +386,6 @@ class LocalFileStore(BaseFileStore): f"missing={len(missing)}, elapsed={time.monotonic() - scan_started_at:.3f}s", ) if not missing: - if self._embedding_rebuild_pending: - await self._after_embedding_backfill() - await self.dump() - self._embedding_rebuild_pending = False self.logger.info( f"{self.name}: embedding backfill complete: filled=0/0, " f"elapsed={time.monotonic() - started_at:.3f}s", @@ -434,11 +446,10 @@ class LocalFileStore(BaseFileStore): self.logger.info( f"{self.name}: embedding backfill complete: filled={filled}/{total}, elapsed={elapsed:.2f}s", ) - if filled or self._embedding_rebuild_pending: + if filled and not self._embedding_rebuild_pending: try: await self._after_embedding_backfill() await self.dump() - self._embedding_rebuild_pending = False except Exception: self.logger.exception(f"{self.name}: failed to persist completed embedding backfill") diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 83a21096..0f09cd47 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -787,6 +787,59 @@ def test_verified_rebuild_discards_late_result_from_previous_vector_space(): run(go()) +def test_unverified_rebuild_is_queued_behind_inflight_backfill(): + """An unverified rebuild request cannot be lost while another batch is running.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_unverified_rebuild_race") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = DelayedOldVectorStore() + store.embedding_store = fake + store._start_embedding_backfill(skip_health_check=True) + old_task = store._embedding_backfill_task + await fake.first_batch_started.wait() + + assert await store.resume_embedding(rebuild=True) is True + assert store._embedding_backfill_pending == (False, True) + fake.release_first_batch.set() + + await old_task + if store._embedding_backfill_task is not None: + await store._embedding_backfill_task + assert fake.node_embedding_calls == [["a"], ["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert store._embedding_rebuild_pending is False + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_clear_during_scheduled_rebuild_finishes_rebuild_state(store_factory): + """Clearing all chunks before the worker scan must not disable vector search forever.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_embedding_clear_during_rebuild") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + store.embedding_store = CountingFakeEmbeddingStore() + + assert await store.resume_embedding(verified=True, rebuild=True) is True + task = store._embedding_backfill_task + await store.clear() + if task is not None: + await task + + assert store.file_chunks == {} + assert store._embedding_rebuild_pending is False + await store.close() + + run(go()) + + @pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) def test_search_recovery_schedules_backfill_without_another_health_check(store_factory): """A successful real search request repairs historical missing vectors.""" diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 6dde0a8c..713afaa8 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -48,6 +48,9 @@ class FailingHealthAsEmbedding: def __init__(self): self.calls = 0 + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + async def __call__(self, _texts: list[str], **_kwargs): self.calls += 1 raise ConnectionError("not ready") @@ -197,6 +200,8 @@ def test_health_check_starts_timeout_after_provider_initialization(monkeypatch): assert await store.health_check(timeout=5.0) is True assert events == ["initialized", "remote request"] + run(go()) + def test_health_check_makes_one_attempt(): """A failed startup probe does not add hidden retries."""