fix(file_store): preserve queued embedding rebuilds

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")
 

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."""

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."""