fix(embedding): retry vector space changes per request

docs/en/plugin_management.md

@@ -173,11 +173,8 @@ curl -s http://127.0.0.1:2333/auto_fin \
 
 When the application uses an MCP service, service-enabled plugin Jobs appear as MCP tools instead.
 
-To add the plugin to another application config, select it explicitly:
-
-```bash
-reme start config=demo plugins='["auto-fin"]'
-```
+Custom application configs must provide the plugin's runtime dependencies, including an `agent_wrapper.default` and
+the `search` and `read` Jobs used by Auto Fin.
 
 ## Uninstall a plugin
 

docs/zh/plugin_management.md

@@ -167,11 +167,7 @@ curl -s http://127.0.0.1:2333/auto_fin \
 
 当应用使用 MCP service 时，允许对外服务的插件 Job 会显示为 MCP tool。
 
-如果需要将插件叠加到其他应用配置，则显式选择该配置：
-
-```bash
-reme start config=demo plugins='["auto-fin"]'
-```
+自定义应用配置需要提供插件的运行依赖，包括 `agent_wrapper.default`，以及 Auto Fin 使用的 `search` 和 `read` Jobs。
 
 ## 卸载插件
 

plugins/auto-fin/README.md

@@ -59,11 +59,7 @@ reme start plugins='["auto-fin"]' \
   service.backend=http
 ```
 
-To add Auto Fin to another application instead, select that config explicitly, for example:
-
-```bash
-reme start config=demo plugins='["auto-fin"]'
-```
+Custom application configs must provide `agent_wrapper.default` and the `search` and `read` Jobs used by Auto Fin.
 
 ## Pipeline
 

plugins/auto-fin/README_ZH.md

@@ -54,11 +54,7 @@ reme start plugins='["auto-fin"]' \
   service.backend=http
 ```
 
-如果需要将 Auto Fin 叠加到其他应用，则显式选择相应配置，例如：
-
-```bash
-reme start config=demo plugins='["auto-fin"]'
-```
+自定义应用配置需要提供 `agent_wrapper.default`，以及 Auto Fin 使用的 `search` 和 `read` Jobs。
 
 ## 流程
 

reme/components/embedding_store/local_embedding_store.py

@@ -104,12 +104,25 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
     # -- Public API --
 
     async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]:
-        await self._sync_cache_space()
         texts = [self._truncate(t) for t in input_text]
-        results, misses = self._partition_by_cache(texts)
-        if misses:
-            await self._fill_misses(misses, results, **kwargs)
-        return results
+        for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1):
+            await self._sync_cache_space()
+            vector_space_id = self._cache_space
+            results, misses = self._partition_by_cache(texts)
+            stable = not misses or await self._fill_misses(misses, results, vector_space_id, **kwargs)
+            if stable and vector_space_id == self.vector_space_id == self._cache_space:
+                return results
+            if attempt == _MAX_VECTOR_SPACE_ATTEMPTS:
+                self.logger.warning(
+                    f"Embedding vector space kept changing while computing a request; "
+                    f"discarding all result(s) after {attempt} attempts",
+                )
+            else:
+                self.logger.info(
+                    f"Embedding vector space changed while computing a request; "
+                    f"discarding all result(s) and retrying ({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})",
+                )
+        return [None] * len(texts)
 
     # -- Batching --
 
@@ -125,31 +138,26 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
                 misses.append((idx, text, key))
         return results, misses
 
-    async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None:
+    async def _fill_misses(
+        self,
+        misses: list[Miss],
+        results: list[np.ndarray | None],
+        vector_space_id: str,
+        **kwargs,
+    ) -> bool:
+        """Fill every miss only while the request remains in one vector space."""
         size = self.max_batch_size
         for start in range(0, len(misses), size):
+            if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space:
+                return False
             batch = misses[start : start + size]
-            for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1):
-                await self._sync_cache_space()
-                vector_space_id = self._cache_space
-                computed = await self._compute_batch(batch, **kwargs)
-                if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space:
-                    if attempt == _MAX_VECTOR_SPACE_ATTEMPTS:
-                        self.logger.warning(
-                            f"Embedding vector space kept changing while computing a batch; "
-                            f"discarding {len(computed)} stale result(s) after {attempt} attempts",
-                        )
-                    else:
-                        self.logger.info(
-                            f"Embedding vector space changed while computing a batch; "
-                            f"discarding {len(computed)} stale result(s) and retrying "
-                            f"({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})",
-                        )
-                    continue
-                for idx, key, emb in computed:
-                    results[idx] = emb
-                    self._cache_put(key, emb)
-                break
+            computed = await self._compute_batch(batch, **kwargs)
+            if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space:
+                return False
+            for idx, key, emb in computed:
+                results[idx] = emb
+                self._cache_put(key, emb)
+        return True
 
     async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]:
         texts = [text for _, text, _ in batch]

tests/unit/test_local_embedding_store.py

@@ -509,41 +509,79 @@ def test_cache_space_is_rechecked_after_async_load(monkeypatch, tmp_path):
     run(go())
 
 
-def test_completed_request_retries_after_vector_space_changes():
-    """A request completed by the old provider must not escape into the new vector space."""
+def test_whole_request_retries_after_vector_space_changes_between_batches():
+    """Completed batches must be discarded when a later batch changes vector space."""
 
     async def go():
-        embedding = OpenAIAsEmbedding(name="t_space_write_race", backend="openai", model="v3", dimensions=2)
-        store = LocalEmbeddingStore(name="t_local_write_race")
+        embedding = FakeAsEmbedding()
+        embedding.vector_space_id = "v3"
+        store = LocalEmbeddingStore(name="t_local_write_race", max_batch_size=1, enable_cache=False)
         store.as_embedding = embedding
         store._cache_space = embedding.vector_space_id
         calls = 0
 
-        async def compute_after_round_trip(_batch, **_kwargs):
+        async def switch_during_second_batch(batch, **_kwargs):
             nonlocal calls
             calls += 1
-            if calls == 1:
-                embedding.model = FakeProviderModel("v4")
-                return [(0, "key", np.array([3.0, 0.0], dtype=np.float16))]
-            return [(0, "key", np.array([4.0, 0.0], dtype=np.float16))]
+            if calls == 2:
+                embedding.vector_space_id = "v4"
+            idx, _text, key = batch[0]
+            version = 3.0 if calls < 3 else 4.0
+            return [(idx, key, np.array([version, 0.0], dtype=np.float16))]
 
-        store._compute_batch = compute_after_round_trip
-        results = [None]
-        await store._fill_misses([(0, "text", "key")], results)
+        store._compute_batch = switch_during_second_batch
+        results = await store.get_embeddings(["first", "second"])
 
-        assert calls == 2
+        assert calls == 4
         assert store._cache_space == embedding.vector_space_id
-        np.testing.assert_array_equal(results[0], np.array([4.0, 0.0], dtype=np.float16))
-        np.testing.assert_array_equal(store._cache["key"], np.array([4.0, 0.0], dtype=np.float16))
+        for result in results:
+            np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))
 
     run(go())
 
 
-def test_completed_request_stops_retrying_when_vector_space_keeps_changing():
-    """Continuous configuration churn must leave the batch empty instead of blocking forever."""
+def test_whole_request_rereads_cache_after_vector_space_changes(monkeypatch, tmp_path):
+    """A cache hit from the old space must not survive a later provider switch."""
 
     async def go():
-        embedding = OpenAIAsEmbedding(name="t_space_write_churn", backend="openai", model="v3", dimensions=2)
+        monkeypatch.setattr(
+            LocalEmbeddingStore,
+            "component_metadata_path",
+            property(lambda _self: tmp_path),
+        )
+        embedding = FakeAsEmbedding()
+        embedding.vector_space_id = "v3"
+        store = LocalEmbeddingStore(name="t_local_cache_race", enable_cache=True)
+        store.as_embedding = embedding
+        store._cache_space = embedding.vector_space_id
+        first_key = store._cache_key("first")
+        store._cache[first_key] = np.array([3.0, 0.0], dtype=np.float16)
+        calls = 0
+
+        async def switch_on_miss(batch, **_kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 1:
+                embedding.vector_space_id = "v4"
+            version = 3.0 if calls == 1 else 4.0
+            return [(idx, key, np.array([version, 0.0], dtype=np.float16)) for idx, _text, key in batch]
+
+        store._compute_batch = switch_on_miss
+        results = await store.get_embeddings(["first", "second"])
+
+        assert calls == 2
+        for result in results:
+            np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))
+
+    run(go())
+
+
+def test_whole_request_stops_retrying_when_vector_space_keeps_changing():
+    """Continuous configuration churn must discard the whole request instead of blocking forever."""
+
+    async def go():
+        embedding = FakeAsEmbedding()
+        embedding.vector_space_id = "v3"
         store = LocalEmbeddingStore(name="t_local_write_churn")
         store.as_embedding = embedding
         store._cache_space = embedding.vector_space_id
@@ -552,12 +590,11 @@ def test_completed_request_stops_retrying_when_vector_space_keeps_changing():
         async def change_space_every_time(_batch, **_kwargs):
             nonlocal calls
             calls += 1
-            embedding.model = FakeProviderModel(f"v{calls + 3}")
+            embedding.vector_space_id = f"v{calls + 3}"
             return [(0, "key", np.array([float(calls), 0.0], dtype=np.float16))]
 
         store._compute_batch = change_space_every_time
-        results = [None]
-        await store._fill_misses([(0, "text", "key")], results)
+        results = await store.get_embeddings(["text"])
 
         assert calls == 3
         assert results == [None]