fix: recover embeddings after transient health check failure (#471)
Some checks failed
CI / Python tests / Unit Tests - py3.12 (push) Has been cancelled
CI / Python tests / Unit Tests - py3.13 (push) Has been cancelled
CI / TypeScript integrations / Type-check, test, and pack (push) Has been cancelled
CI / Windows / CLI smoke - py3.11 (push) Has been cancelled
Deploy / Documentation / Build documentation (push) Has been cancelled
Security / CodeQL / Analyze javascript-typescript (push) Has been cancelled
Security / CodeQL / Analyze python (push) Has been cancelled
CI / Documentation / Test and build documentation (push) Has been cancelled
CI / Python quality / Pre-commit (push) Has been cancelled
CI / Python tests / Unit Tests - py3.11 (push) Has been cancelled
Deploy / Documentation / deploy (push) Has been cancelled

* fix: recover embedding after transient health failure

* refactor(embedding_store): remove provider_success_count and simplify health recovery logic

- Deleted provider_success_count attribute and related methods across embedding and file stores
- Updated _recover_after_real_request to rely solely on is_healthy flag for recovery decisions
- Removed redundant counting logic for provider successes during embedding operations
- Cleaned up health status management to streamline provider recovery detection
- Adjusted unit tests to align with removal of provider_success_count and maintain health checks consistency

* refactor(embedding_store): use default health check timeout

* fix(embedding_store): ensure is_healthy remains unchanged on cache hits

- Updated get_embeddings docstring to clarify cache hits must not alter is_healthy state
- Improved code comment for embedding dimension matching method

* fix(file_store): make embedding recovery race-safe

* ci: use default CodeQL query suite

* fix(file_store): preserve queued embedding rebuilds

* fix(file_store): preserve verified recovery without chunks
This commit is contained in:
jinliyl 2026-08-21 13:58:51 +08:00 committed by GitHub
parent 8416fd3ac9
commit c8e1248769
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 541 additions and 49 deletions

View file

@ -37,7 +37,6 @@ jobs:
with:
languages: ${{ matrix.language }}
build-mode: none
queries: security-and-quality
- name: Perform CodeQL analysis
uses: github/codeql-action/analyze@v4

View file

@ -117,6 +117,14 @@ Out of the box, search therefore uses primarily BM25 plus link expansion. After
`SearchStep` runs vector and keyword recall together. Additionally, switching the `file_store` `backend` from `local` to
`faiss` upgrades vector retrieval from a linear scan to a FAISS HNSW index, offering faster recall at scale.
The embedding store accepts `health_check_timeout` for its startup probe. A temporary failure skips the current vector
backfill while keeping BM25 available; a later successful provider request resumes the missing-vector backfill
automatically.
Embedded integrations that have already verified a provider can call `resume_embedding(verified=True)`. When changing
the embedding vector space, pass `rebuild=True`; persisted vectors are invalidated before a serial background rebuild,
and vector search remains unavailable until the rebuilt vectors are safely persisted.
## How to Search
The `search` Job is also configured in `default.yaml`:

View file

@ -106,6 +106,12 @@ file_store:
所以开箱搜索主要是 BM25 + 链接展开。把 `embedding_store: default` 打开后,`SearchStep` 会同时跑向量召回和关键词召回。此时若将
`file_store``backend``local` 改为 `faiss`,向量检索会从线性扫描升级为 FAISS HNSW 索引,在大规模 chunk 场景下召回效率更高。
Embedding store 可通过 `health_check_timeout` 配置启动探测。临时失败只会跳过本次向量回填BM25 仍可使用;
后续真实请求成功后会自动恢复缺失向量的回填。
已经完成真实服务验证的嵌入式集成可以调用 `resume_embedding(verified=True)`。切换 Embedding 向量空间时应同时传入
`rebuild=True`ReMe 会先使旧向量失效,再串行后台重建,并在新向量安全持久化前暂停向量搜索。
## 怎么搜索
`search` Job 也是在 `default.yaml` 中配置:

View file

@ -25,6 +25,7 @@ class BaseEmbeddingStore(BaseComponent):
max_input_length: int = 8192,
max_retries: int = 3,
quota_retry_delay: float | None = None,
health_check_timeout: float = 15.0,
**kwargs,
):
super().__init__(**kwargs)
@ -32,6 +33,7 @@ class BaseEmbeddingStore(BaseComponent):
self.max_input_length = max_input_length
self.max_retries = max_retries
self.quota_retry_delay = quota_retry_delay
self.health_check_timeout = health_check_timeout
self.is_healthy: bool = True
def _truncate(self, text: str) -> str:
@ -57,7 +59,7 @@ class BaseEmbeddingStore(BaseComponent):
return text
@abstractmethod
async def health_check(self, timeout: float = 2.0) -> bool:
async def health_check(self, timeout: float | None = None) -> bool:
"""Probe the provider; sets and returns is_healthy."""
async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray | None:
@ -67,7 +69,7 @@ class BaseEmbeddingStore(BaseComponent):
@abstractmethod
async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]:
"""Get embeddings for texts."""
"""Get embeddings; cache hits must not change is_healthy."""
def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool:
"""Return whether an embedding matches the configured model dimension."""

View file

@ -68,8 +68,12 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
async def _close(self) -> None:
await self.dump()
async def health_check(self, timeout: float = 5.0) -> bool:
async def health_check(self, timeout: float | None = None) -> bool:
timeout = self.health_check_timeout if timeout is None else timeout
if not isinstance(timeout, (int, float)) or not np.isfinite(timeout) or timeout <= 0:
raise ValueError("timeout must be finite and greater than 0")
tag = f"[EMBEDDING HEALTH CHECK] name={self.name} workspace_dir={self.workspace_path}"
started_at = asyncio.get_running_loop().time()
try:
# Provider construction may synchronously import an SDK and build
# its HTTP client. Keep that one-time work outside the request
@ -82,13 +86,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
if len(result[0]) != self.dimensions:
raise RuntimeError(f"embedding dimension mismatch: {len(result[0])} != {self.dimensions}")
self.is_healthy = True
self.logger.info(f"{tag} -> OK")
elapsed = asyncio.get_running_loop().time() - started_at
self.logger.info(f"{tag} -> OK timeout={timeout}s elapsed={elapsed:.3f}s")
except asyncio.TimeoutError:
self.is_healthy = False
self.logger.error(f"{tag} -> FAIL timeout({timeout}s)")
except Exception as e:
elapsed = asyncio.get_running_loop().time() - started_at
self.logger.error(f"{tag} -> FAIL timeout={timeout}s elapsed={elapsed:.3f}s error=timeout({timeout}s)")
except Exception as exc: # Provider SDKs expose many exception types.
self.is_healthy = False
self.logger.error(f"{tag} -> FAIL {type(e).__name__}: {e}")
elapsed = asyncio.get_running_loop().time() - started_at
self.logger.error(
f"{tag} -> FAIL timeout={timeout}s elapsed={elapsed:.3f}s error={type(exc).__name__}: {exc}",
)
return self.is_healthy
# -- Public API --
@ -143,6 +152,10 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
if bad_dims:
details = ", ".join(f"{count} with dim {dim}" for dim, count in sorted(bad_dims.items()))
self.logger.error(f"Embedding dimension mismatch in batch: expected {self.dimensions}; rejected {details}")
if out:
self.is_healthy = True
else:
self.is_healthy = False
return out
async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None:
@ -166,7 +179,9 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
await asyncio.sleep(self.quota_retry_delay)
continue
self.logger.exception("Embedding request failed")
self.is_healthy = False
return None
self.is_healthy = False
return None
@staticmethod

View file

@ -252,6 +252,11 @@ class FaissLocalFileStore(LocalFileStore):
self._add_to_index([c.id for c in to_add], vectors)
self._compact_if_needed()
async def _reset_vector_index(self) -> None:
"""Discard all vectors before rebuilding a changed vector space."""
await self._stop_reindex_worker()
self._rebuild_index()
# -- async reindex ----------------------------------------------------
def _submit_reindex(self) -> None:
@ -625,26 +630,29 @@ class FaissLocalFileStore(LocalFileStore):
# -- search -----------------------------------------------------------
async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
index_empty = self._faiss_index is None or self._faiss_index.ntotal == 0
embedding_unavailable = self.embedding_store is None or self._embedding_rebuild_pending
if (
self.embedding_store is None
embedding_unavailable
or not query
or limit <= 0
or self._faiss_index is None
or self._faiss_index.ntotal == 0
or (index_empty and getattr(self.embedding_store, "is_healthy", True))
):
return []
query_embedding = None
was_healthy = bool(getattr(self.embedding_store, "is_healthy", True))
try:
query_embedding = await self.embedding_store.get_embedding(query)
except Exception as e:
self._disable_embedding(f"search: {type(e).__name__}: {e}")
self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}")
if query_embedding is None or not self._embedding_dim_matches(query_embedding):
if query_embedding is not None:
self._disable_embedding(
self._mark_embedding_unhealthy(
f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}",
)
return []
await self._recover_after_real_request(was_healthy)
# get_embedding above yielded control; a concurrent clear() drops the
# index to None once embedding is disabled, and a reindex may have swapped

View file

@ -67,10 +67,14 @@ class LocalFileStore(BaseFileStore):
self.file_chunks: dict[str, FileChunk] = {}
self.chunks_path = self.component_metadata_path / f"file_chunks_{self.name}_{self.store_version}.jsonl.zst"
self._embedding_backfill_task: asyncio.Task | None = None
self._embedding_backfill_pending: tuple[bool, bool] | None = None
self._embedding_rebuild_pending = False
self._closing = False
# -- lifecycle ------------------------------------------------------------
async def _start(self) -> None:
self._closing = False
started_at = time.monotonic()
self.component_metadata_path.mkdir(parents=True, exist_ok=True)
await super()._start()
@ -94,17 +98,45 @@ class LocalFileStore(BaseFileStore):
)
async def _close(self) -> None:
self._closing = True
await self._cancel_embedding_backfill()
await self.dump()
self.file_chunks.clear()
await super()._close()
def _disable_embedding(self, reason: str) -> None:
"""Drop embedding after a runtime failure; keyword search still works."""
def _mark_embedding_unhealthy(self, reason: str) -> None:
"""Record a temporary provider failure while preserving the component."""
if self.embedding_store is None:
return
self.logger.error(f"{self.name}: embedding disabled, {reason}")
self.embedding_store = None
self.embedding_store.is_healthy = False
self.logger.error(f"{self.name}: embedding unavailable, {reason}; keyword search remains active")
async def _recover_after_real_request(self, was_healthy: bool) -> None:
"""Schedule repair when a real, non-cache provider request recovers."""
if self.embedding_store is None or was_healthy or not getattr(self.embedding_store, "is_healthy", True):
return
self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill")
await self.resume_embedding(verified=True)
async def resume_embedding(self, *, verified: bool = False, rebuild: bool = False) -> bool:
"""Resume a configured provider and schedule a deduplicated repair.
Embedded applications may pass ``verified=True`` after they have already
made a successful real provider request, avoiding a redundant ping. Pass
``rebuild=True`` when the active vector space changed; existing vectors
are derived data and are discarded before a full background rebuild.
"""
if self.embedding_store is None or self._closing:
return False
if verified:
self.embedding_store.is_healthy = True
if rebuild:
await self._prepare_embedding_rebuild()
if not self.file_chunks:
self._embedding_rebuild_pending = False
return True
self._start_embedding_backfill(skip_health_check=verified, rebuild=rebuild)
return True
def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool:
"""Return whether an index embedding matches the active embedding model."""
@ -228,29 +260,40 @@ class LocalFileStore(BaseFileStore):
return
self._drop_stale_embeddings(self.file_chunks.values(), "load")
def _start_embedding_backfill(self) -> None:
def _start_embedding_backfill(self, *, skip_health_check: bool = False, rebuild: bool = False) -> None:
"""Schedule startup embedding repair without delaying component readiness."""
started_at = time.monotonic()
if self._closing:
self.logger.info(f"{self.name}: embedding backfill skipped: reason=closing")
return
if not self.embedding_store:
self.logger.info(
f"{self.name}: embedding backfill skipped: reason=embedding_disabled, "
f"elapsed={time.monotonic() - started_at:.3f}s",
)
return
if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done():
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",
)
return
if not self.file_chunks:
self.logger.info(
f"{self.name}: embedding backfill skipped: reason=no_chunks, "
f"elapsed={time.monotonic() - started_at:.3f}s",
)
return
if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done():
self.logger.info(
f"{self.name}: embedding backfill scheduling skipped: reason=already_running, "
f"elapsed={time.monotonic() - started_at:.3f}s",
)
return
self._embedding_backfill_task = asyncio.create_task(
self._backfill_missing_embeddings(),
self._run_embedding_backfill(skip_health_check=skip_health_check, rebuild=rebuild),
name=f"embedding-backfill:{self.name}",
)
self.logger.info(
@ -258,10 +301,37 @@ class LocalFileStore(BaseFileStore):
f"elapsed={time.monotonic() - started_at:.3f}s",
)
async def _run_embedding_backfill(self, *, skip_health_check: bool, rebuild: bool) -> None:
"""Run one repair and honor a verified request queued behind it."""
current_task = asyncio.current_task()
try:
if rebuild:
# A task that was already running when rebuild was requested
# may have written a stale provider result after the first
# invalidation. Clear once more at the queue boundary.
await self._prepare_embedding_rebuild()
await self._backfill_missing_embeddings(skip_health_check=skip_health_check)
finally:
if self._embedding_backfill_task is current_task:
self._embedding_backfill_task = None
pending = self._embedding_backfill_pending
self._embedding_backfill_pending = None
if pending is not None and not self._closing and self.embedding_store is not None:
pending_verified, pending_rebuild = pending
if pending_verified:
self.embedding_store.is_healthy = True
if pending_rebuild:
self._embedding_rebuild_pending = True
self._start_embedding_backfill(
skip_health_check=pending_verified,
rebuild=pending_rebuild,
)
async def _cancel_embedding_backfill(self) -> None:
"""Cancel and collect the startup repair task during component shutdown."""
task = self._embedding_backfill_task
self._embedding_backfill_task = None
self._embedding_backfill_pending = None
if task is None:
return
if not task.done():
@ -285,9 +355,22 @@ class LocalFileStore(BaseFileStore):
next_percent += _PROGRESS_LOG_PERCENT_STEP
return next_percent
async def _backfill_missing_embeddings(self) -> None:
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: "
@ -313,18 +396,18 @@ class LocalFileStore(BaseFileStore):
batch_size = max(1, int(getattr(self.embedding_store, "max_batch_size", 10)))
self.logger.info(f"{self.name}: embedding backfill started: total={total}, batch_size={batch_size}")
try:
health_check_started_at = time.monotonic()
is_healthy = await self.embedding_store.health_check()
self.logger.info(
f"{self.name}: embedding health check complete: healthy={is_healthy}, "
f"elapsed={time.monotonic() - health_check_started_at:.3f}s",
)
if not is_healthy:
self._disable_embedding("backfill health check failed")
self.logger.warning(
f"{self.name}: embedding backfill failed: processed=0/{total}, reason=health check failed",
if not skip_health_check:
health_check_started_at = time.monotonic()
is_healthy = await self.embedding_store.health_check()
self.logger.info(
f"{self.name}: embedding health check complete: healthy={is_healthy}, "
f"elapsed={time.monotonic() - health_check_started_at:.3f}s",
)
return
if not is_healthy:
self.logger.warning(
f"{self.name}: embedding backfill skipped: processed=0/{total}, reason=health check failed",
)
return
processed = 0
batch_count = 0
@ -348,8 +431,9 @@ class LocalFileStore(BaseFileStore):
f"{total}, elapsed={elapsed:.2f}s",
)
raise
except Exception as e:
self._disable_embedding(f"backfill: {type(e).__name__}: {e}")
self._mark_embedding_unhealthy(f"backfill: {type(e).__name__}: {e}")
elapsed = time.monotonic() - started_at
self.logger.exception(
f"{self.name}: embedding backfill failed: processed={processed if 'processed' in locals() else 0}/"
@ -362,13 +446,24 @@ class LocalFileStore(BaseFileStore):
self.logger.info(
f"{self.name}: embedding backfill complete: filled={filled}/{total}, elapsed={elapsed:.2f}s",
)
if filled:
if filled and not self._embedding_rebuild_pending:
try:
await self._after_embedding_backfill()
await self.dump()
except Exception:
self.logger.exception(f"{self.name}: failed to persist completed embedding backfill")
async def _prepare_embedding_rebuild(self) -> None:
"""Invalidate and persist vectors from the previous vector space."""
self._embedding_rebuild_pending = True
for chunk in self.file_chunks.values():
chunk.embedding = None
await self._reset_vector_index()
await self.dump()
async def _reset_vector_index(self) -> None:
"""Drop a derived vector index before rebuilding a changed vector space."""
async def _after_embedding_backfill(self) -> None:
"""Backend hook for refreshing derived vector indexes after backfill."""
@ -540,12 +635,15 @@ class LocalFileStore(BaseFileStore):
async def _embed_pending(self, chunks: list[FileChunk]) -> None:
if not (chunks and self.embedding_store):
return
was_healthy = bool(getattr(self.embedding_store, "is_healthy", True))
try:
await self.embedding_store.get_node_embeddings(chunks)
except Exception as e:
self._disable_embedding(f"upsert: {type(e).__name__}: {e}")
self._mark_embedding_unhealthy(f"upsert: {type(e).__name__}: {e}")
return
self._drop_stale_embeddings(chunks, "upsert")
if any(chunk.embedding is not None for chunk in chunks):
await self._recover_after_real_request(was_healthy)
async def delete(self, path: str | list[str]) -> None:
assert self.file_graph is not None
@ -601,21 +699,23 @@ class LocalFileStore(BaseFileStore):
# -- search ---------------------------------------------------------------
async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
if self.embedding_store is None or not query or limit <= 0:
if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0:
return []
was_healthy = bool(getattr(self.embedding_store, "is_healthy", True))
try:
query_embedding = await self.embedding_store.get_embedding(query)
except Exception as e:
self._disable_embedding(f"search: {type(e).__name__}: {e}")
self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}")
return []
if query_embedding is None:
return []
if not self._embedding_dim_matches(query_embedding):
self._disable_embedding(
self._mark_embedding_unhealthy(
f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}",
)
return []
await self._recover_after_real_request(was_healthy)
top: list[tuple[float, int, FileChunk]] = []
candidates: list[FileChunk] = []

View file

@ -176,6 +176,10 @@ class ZvecLocalFileStore(LocalFileStore):
]
self._upsert_docs(to_add)
async def _reset_vector_index(self) -> None:
"""Discard all vectors before rebuilding a changed vector space."""
self._collection = self._create_collection()
# -- maintenance ------------------------------------------------------
async def optimize_index(self) -> None:
@ -412,20 +416,25 @@ class ZvecLocalFileStore(LocalFileStore):
# -- search -----------------------------------------------------------
async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
if self.embedding_store is None or not query or limit <= 0 or self._collection is None or not self._indexed_ids:
if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0:
return []
index_empty = self._collection is None or not self._indexed_ids
if index_empty and getattr(self.embedding_store, "is_healthy", True):
return []
query_embedding = None
was_healthy = bool(getattr(self.embedding_store, "is_healthy", True))
try:
query_embedding = await self.embedding_store.get_embedding(query)
except Exception as e:
self._disable_embedding(f"search: {type(e).__name__}: {e}")
self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}")
if query_embedding is None or not self._embedding_dim_matches(query_embedding):
if query_embedding is not None:
self._disable_embedding(
self._mark_embedding_unhealthy(
f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}",
)
return []
await self._recover_after_real_request(was_healthy)
# get_embedding above yielded control; a concurrent clear() may have
# swapped or dropped the collection. Re-read before dereferencing.

View file

@ -16,6 +16,7 @@ import pytest
from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore
from reme.components.file_store import local_file_store as local_file_store_module
from reme.components.embedding_store import LocalEmbeddingStore
from reme.schema import FileChunk, FileNode
from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
@ -67,6 +68,7 @@ class CountingFakeEmbeddingStore(FakeEmbeddingStore):
def __init__(self):
self.node_embedding_calls: list[list[str]] = []
self.is_healthy = True
async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]:
self.node_embedding_calls.append([node.id for node in nodes])
@ -76,10 +78,35 @@ class CountingFakeEmbeddingStore(FakeEmbeddingStore):
class UnhealthyCountingEmbeddingStore(CountingFakeEmbeddingStore):
"""Fake embedding store that fails the backfill health gate."""
def __init__(self):
super().__init__()
self.is_healthy = False
async def health_check(self, _timeout: float = 2.0) -> bool:
return False
class RecoveringEmbeddingStore(CountingFakeEmbeddingStore):
"""Fake provider that starts unhealthy and records real recoveries."""
def __init__(self):
super().__init__()
self.is_healthy = False
self.health_calls = 0
async def health_check(self, _timeout: float = 2.0) -> bool:
self.health_calls += 1
return False
async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray:
self.is_healthy = True
return await super().get_embedding(input_text, **kwargs)
async def get_node_embeddings(self, nodes: list[FileChunk], **kwargs) -> list[FileChunk]:
self.is_healthy = True
return await super().get_node_embeddings(nodes, **kwargs)
class HealthCountingEmbeddingStore(FakeEmbeddingStore):
"""Fake provider that records eager health checks."""
@ -104,6 +131,44 @@ class BlockingEmbeddingStore(FakeEmbeddingStore):
return await super().get_node_embeddings(nodes, **kwargs)
class CancellationResistantHealthStore(CountingFakeEmbeddingStore):
"""Startup probe that completes stale after cancellation is requested."""
def __init__(self):
super().__init__()
self.is_healthy = True
self.health_started = asyncio.Event()
self.release_health = asyncio.Event()
async def health_check(self, _timeout: float = 2.0) -> bool:
self.health_started.set()
try:
await self.release_health.wait()
except asyncio.CancelledError:
await self.release_health.wait()
self.is_healthy = False
return False
class DelayedOldVectorStore(CountingFakeEmbeddingStore):
"""First batch returns an old-space vector after rebuild was requested."""
def __init__(self):
super().__init__()
self.first_batch_started = asyncio.Event()
self.release_first_batch = asyncio.Event()
async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]:
self.node_embedding_calls.append([node.id for node in nodes])
if len(self.node_embedding_calls) == 1:
self.first_batch_started.set()
await self.release_first_batch.wait()
for chunk_node in nodes:
chunk_node.embedding = np.array([0.0, 1.0], dtype=np.float16)
return nodes
return await FakeEmbeddingStore.get_node_embeddings(self, nodes)
class WrongDimEmbeddingStore(FakeEmbeddingStore):
"""Fake embedding store that returns vectors with the wrong dimension."""
@ -153,6 +218,15 @@ def _new_local_store(name, **kwargs):
return LocalFileStore(name=name, embedding_store="", **kwargs)
def _new_faiss_store(name, **kwargs):
"""Construct a FAISS store when the optional backend is installed."""
try:
store = FaissLocalFileStore(name=name, embedding_store="", **kwargs)
except ImportError:
pytest.skip("faiss is not installed")
return store
def _new_zvec_store(name, **kwargs):
"""Construct a zvec store with embedding disabled at bind time."""
try:
@ -593,7 +667,7 @@ def test_background_embedding_backfill_uses_provider_batch_size():
def test_load_skips_backfill_when_embedding_health_check_fails():
"""Background backfill disables embeddings before batching when the provider is unhealthy."""
"""Background backfill preserves an unhealthy provider for later recovery."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
@ -610,13 +684,249 @@ def test_load_skips_backfill_when_embedding_health_check_fails():
await store._embedding_backfill_task
assert not fake.node_embedding_calls
assert store.embedding_store is None
assert store.embedding_store is fake
assert fake.is_healthy is False
assert store.file_chunks["a"].embedding is None
await store.close()
run(go())
def test_verified_resume_supersedes_inflight_startup_health_check():
"""A stale startup probe cannot consume or overwrite verified recovery."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = _new_local_store("t_embedding_verified_resume_race")
await store.start()
await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")})
fake = CancellationResistantHealthStore()
store.embedding_store = fake
store._start_embedding_backfill()
startup_task = store._embedding_backfill_task
await fake.health_started.wait()
recovery = asyncio.create_task(store.resume_embedding(verified=True))
await asyncio.sleep(0)
assert await recovery is True
assert store._embedding_backfill_pending == (True, False)
fake.release_health.set()
await startup_task
assert store._embedding_backfill_task is not startup_task
if store._embedding_backfill_task is not None:
await store._embedding_backfill_task
assert fake.is_healthy is True
assert fake.node_embedding_calls == [["a"]]
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
await store.close()
run(go())
def test_verified_resume_without_chunks_supersedes_inflight_health_check():
"""A newer verified state survives a stale probe even after chunks are cleared."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = _new_local_store("t_embedding_empty_verified_resume_race")
await store.start()
await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")})
fake = CancellationResistantHealthStore()
store.embedding_store = fake
store._start_embedding_backfill()
startup_task = store._embedding_backfill_task
await fake.health_started.wait()
await store.clear()
assert await store.resume_embedding(verified=True) is True
assert store._embedding_backfill_pending == (True, False)
fake.release_health.set()
await startup_task
assert fake.is_healthy is True
assert not fake.node_embedding_calls
assert store._embedding_backfill_task is None
await store.close()
run(go())
@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store])
def test_verified_rebuild_discards_same_dimension_vectors_before_backfill(store_factory):
"""A changed vector space never searches compatible-shaped stale vectors."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = store_factory("t_embedding_verified_rebuild")
await store.start()
stale = chunk("a", "a.md", "alpha text")
stale.embedding = np.array([0.0, 1.0], dtype=np.float16)
await set_chunks_with_graph(store, {"a": stale})
fake = CountingFakeEmbeddingStore()
fake.is_healthy = False
store.embedding_store = fake
if isinstance(store, FaissLocalFileStore):
store._rebuild_index()
elif isinstance(store, ZvecLocalFileStore):
store._rebuild_collection()
assert await store.resume_embedding(verified=True, rebuild=True) is True
assert store._embedding_rebuild_pending is True
assert store.file_chunks["a"].embedding is None
assert await store.vector_search("alpha", 5, {}) == []
await store._embedding_backfill_task
assert store._embedding_rebuild_pending is False
assert fake.node_embedding_calls == [["a"]]
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
assert [item.id for item in await store.vector_search("alpha", 5, {})] == ["a"]
await store.close()
run(go())
def test_verified_rebuild_discards_late_result_from_previous_vector_space():
"""A queued rebuild clears old-space vectors written by an in-flight batch."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = _new_local_store("t_embedding_verified_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(verified=True, rebuild=True) is True
assert store._embedding_backfill_pending == (True, 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())
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."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = store_factory(name="t_embedding_search_recovery")
await store.start()
await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")})
fake = RecoveringEmbeddingStore()
store.embedding_store = fake
assert await store.vector_search("alpha", 5, {}) == []
await store._embedding_backfill_task
assert fake.is_healthy is True
assert fake.health_calls == 0
assert fake.node_embedding_calls == [["a"]]
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
assert [item.id for item in await store.vector_search("alpha", 5, {})] == ["a"]
await store.close()
run(go())
def test_cache_only_search_does_not_mark_provider_recovered():
"""Cached vectors do not prove that the remote provider is available."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
embedding_store = LocalEmbeddingStore(name="t_embedding_cache_only_recovery")
embedding_store.as_embedding = type(
"CachedProvider",
(),
{
"dimensions": 2,
"vector_space_id": "cached-provider",
"__call__": lambda self, texts, **_kwargs: asyncio.sleep(
0,
result=[[1.0, 0.0] for _ in texts],
),
},
)()
await embedding_store.get_embedding("alpha")
embedding_store.is_healthy = False
store = LocalFileStore(name="t_embedding_cache_only_recovery", embedding_store="")
await store.start()
await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "historical text")})
store.embedding_store = embedding_store
assert await store.vector_search("alpha", 5, {}) == []
assert embedding_store.is_healthy is False
assert store._embedding_backfill_task is None
await store.close()
run(go())
@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store])
def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(store_factory):
"""Loading persisted chunks re-embeds vectors that do not match current dimensions."""
@ -692,7 +1002,8 @@ def test_upsert_drops_wrong_dimension_from_custom_embedding_store():
assert store.file_chunks["a"].embedding is None
assert await store.vector_search("alpha", 5, {}) == []
assert store.embedding_store is None
assert store.embedding_store is not None
assert store.embedding_store.is_healthy is False
await store.close()
run(go())

View file

@ -39,6 +39,23 @@ class BadHealthAsEmbedding:
return [[1.0]]
class FailingHealthAsEmbedding:
"""Fake provider that records a failed health-check attempt."""
dimensions = 2
vector_space_id = "fakespace000"
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")
class FakeProviderModel:
"""Stand-in for a constructed AgentScope embedding model object."""
@ -186,6 +203,23 @@ def test_health_check_starts_timeout_after_provider_initialization(monkeypatch):
run(go())
def test_health_check_makes_one_attempt():
"""A failed startup probe does not add hidden retries."""
async def go():
provider = FailingHealthAsEmbedding()
store = LocalEmbeddingStore(
name="t_local_embedding_health_retry",
health_check_timeout=3.0,
)
store.as_embedding = provider
assert await store.health_check() is False
assert provider.calls == 1
run(go())
def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch):
"""Quota exhaustion uses the dedicated delay before ReMe retries."""