From f44f52d91918f426bdd30edab61fbe539b8d50ba Mon Sep 17 00:00:00 2001
From: jinliyl <6469360+jinliyl@users.noreply.github.com>
Date: Fri, 21 Aug 2026 11:23:14 +0800
Subject: [PATCH] fix(embedding): exclude provider init from health timeout
 (#484)

---
 reme/components/as_embedding/__init__.py      |  9 +++++
 .../embedding_store/local_embedding_store.py  |  5 +++
 tests/unit/test_local_embedding_store.py      | 39 +++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/reme/components/as_embedding/__init__.py b/reme/components/as_embedding/__init__.py
index 908249ef..0e079bc8 100644
--- a/reme/components/as_embedding/__init__.py
+++ b/reme/components/as_embedding/__init__.py
@@ -121,6 +121,15 @@ class BaseAsEmbedding(BaseComponent):
         response = await self.model(inputs, **kwargs)  # pylint: disable=not-callable
         return response.embeddings
 
+    def initialize_model(self) -> None:
+        """Construct the provider without making a remote request.
+
+        Callers that apply their own request timeout can initialize first so
+        one-time SDK imports and client construction do not consume that
+        timeout budget. Normal embedding calls remain lazily initialized.
+        """
+        self._ensure_model()
+
     async def _start(self) -> None:
         """Defer provider construction until the first remote embedding call."""
         return None
diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py
index 990b276e..930d9af4 100644
--- a/reme/components/embedding_store/local_embedding_store.py
+++ b/reme/components/embedding_store/local_embedding_store.py
@@ -71,6 +71,11 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
     async def health_check(self, timeout: float = 5.0) -> bool:
         tag = f"[EMBEDDING HEALTH CHECK] name={self.name} workspace_dir={self.workspace_path}"
         try:
+            # Provider construction may synchronously import an SDK and build
+            # its HTTP client. Keep that one-time work outside the request
+            # timeout so the full budget applies to the initialized provider
+            # call instead of being consumed before a request can be sent.
+            self.as_embedding.initialize_model()
             result = await asyncio.wait_for(self.as_embedding(["ping"]), timeout=timeout)
             if not result or result[0] is None:
                 raise RuntimeError("empty embedding")
diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py
index 8ff851f4..5ac1d925 100644
--- a/tests/unit/test_local_embedding_store.py
+++ b/tests/unit/test_local_embedding_store.py
@@ -19,6 +19,9 @@ class FakeAsEmbedding:
     dimensions = 2
     vector_space_id = "fakespace000"
 
+    def initialize_model(self):
+        """Mirror the real component's idempotent initialization hook."""
+
     async def __call__(self, texts: list[str], **_kwargs):
         return [[1.0] if text == "bad" else [1.0, 0.0] for text in texts]
 
@@ -29,6 +32,9 @@ class BadHealthAsEmbedding:
     dimensions = 2
     vector_space_id = "fakespace000"
 
+    def initialize_model(self):
+        """Mirror the real component's idempotent initialization hook."""
+
     async def __call__(self, _texts: list[str], **_kwargs):
         return [[1.0]]
 
@@ -147,6 +153,39 @@ def test_health_check_rejects_wrong_dimension():
     run(go())
 
 
+def test_health_check_starts_timeout_after_provider_initialization(monkeypatch):
+    """One-time client construction must not consume the request timeout."""
+
+    async def go():
+        events = []
+
+        class InitializingAsEmbedding(FakeAsEmbedding):
+            """Record initialization and provider-call ordering."""
+
+            def initialize_model(self):
+                events.append("initialized")
+
+            async def __call__(self, texts: list[str], **_kwargs):
+                events.append("remote request")
+                return [[1.0, 0.0] for _ in texts]
+
+        original_wait_for = asyncio.wait_for
+
+        async def checked_wait_for(awaitable, timeout):
+            assert events == ["initialized"]
+            assert timeout == 5.0
+            return await original_wait_for(awaitable, timeout)
+
+        store = LocalEmbeddingStore(name="t_local_embedding_health_timeout_scope")
+        store.as_embedding = InitializingAsEmbedding()
+        monkeypatch.setattr(asyncio, "wait_for", checked_wait_for)
+
+        assert await store.health_check(timeout=5.0) is True
+        assert events == ["initialized", "remote request"]
+
+    run(go())
+
+
 def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch):
     """Quota exhaustion uses the dedicated delay before ReMe retries."""