fix(llm): forward provider_params.timeout to the OpenAI-compatible embedding client (#1024)

* fix(llm): forward provider_params.timeout to the OpenAI-compatible embedding client

#832 and #903 added a configurable request timeout for the LLM registry
and the Gemini embedding client respectively, but the OpenAI-compatible
embedding client (src/embedding_client.py) was never wired up. It
constructed AsyncOpenAI with no timeout at all, so a stalled socket
against a slow or contended OpenAI-compatible backend (e.g. a
self-hosted embedding model under load) wedges the deriver worker's
event loop indefinitely — the exact failure #785/#903 describe, just
via a code path #903 didn't cover.

EmbeddingModelConfig now carries provider_params through from
resolve_embedding_model_config, mirroring how resolve_model_config
already does it for ModelConfig, and the OpenAI branch of
_EmbeddingClient.__init__ extracts `timeout` via the existing
request_timeout_from_extra_params helper. Unset stays unset — no
existing behavior changes.

Reproduced and verified against a real self-hosted deployment (local
Ollama backend under load): before this fix, a single stuck embedding
call blocked all deriver queue processing for 20+ minutes with no
error logged, twice in one session.

* fix(embedding): use first-class timeout on embedding model config

provider_params is the LLM per-request escape hatch; embedding timeouts are
client-construction knobs and belong next to max_batch_size. Wire the field
for OpenAI and Gemini, omit the OpenAI kwarg when unset so the SDK default
stays, and keep Gemini's 10-minute floor when unset.

* test(embedding): live coverage for first-class embedding timeout

Exercise EmbeddingModelConfig.timeout on one representative OpenAI and
Gemini model: configured timeout lands on the SDK client, and a near-zero
timeout aborts before the provider answers.

---------

Co-authored-by: Aakash Kattelu <aakash@plasticlabs.ai>

config.toml.example

@@ -77,6 +77,8 @@ model = "text-embedding-3-small"
 # Optional provider request input cap. Useful for OpenAI-compatible embedding
 # APIs with smaller limits, such as DashScope text-embedding-v4.
 # max_batch_size = 10
+# Optional client HTTP timeout in seconds (OpenAI + Gemini).
+# timeout = 90.0
 
 # Optional module-level endpoint overrides
 # [embedding.model_config.overrides]

docs/v3/contributing/configuration.mdx

@@ -267,6 +267,7 @@ EMBEDDING_MAX_TOKENS_PER_REQUEST=300000
 EMBEDDING_MODEL_CONFIG__TRANSPORT=openai  # openai, gemini
 EMBEDDING_MODEL_CONFIG__MODEL=text-embedding-3-small
 EMBEDDING_MODEL_CONFIG__MAX_BATCH_SIZE=10  # optional per-request input cap
+EMBEDDING_MODEL_CONFIG__TIMEOUT=90.0       # optional client HTTP timeout (seconds)
 
 # Optional endpoint overrides
 EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=http://localhost:8000/v1
@@ -279,6 +280,13 @@ document a per-request limit. Set it when an OpenAI-compatible embedding
 provider accepts fewer inputs per request, such as DashScope
 `text-embedding-v4` with a limit of 10.
 
+`EMBEDDING_MODEL_CONFIG__TIMEOUT` is an optional client HTTP timeout in
+seconds. OpenAI-compatible transports receive it as the SDK `timeout` kwarg
+(omitted when unset, so the SDK default applies). Gemini converts it to
+milliseconds on `http_options.timeout`, and keeps its existing 10-minute
+default when unset. The value is validated at config load the same way as
+LLM `provider_params.timeout` (positive, finite number of seconds).
+
 Forwarding `dimensions=` to OpenAI-compatible providers is controlled by `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE`:
 
 - `auto` (default): forwards `dimensions=` when **the operator has explicitly set `EMBEDDING_VECTOR_DIMENSIONS`** — provenance, not value — and the configured model is not on the known-rejecting list (currently `text-embedding-ada-002`). Explicit `EMBEDDING_VECTOR_DIMENSIONS=1536` *does* trigger the forward; this is how `text-embedding-3-large` truncation to 1536 is expressed. Deployments that leave the setting unset get their existing behavior (`dimensions=` is not forwarded).

src/config.py

@@ -394,6 +394,15 @@ class ConfiguredEmbeddingModelSettings(BaseModel):
     dimensions_mode: EmbeddingDimensionsMode = "auto"
     encoding_format_mode: EmbeddingEncodingFormatMode = "auto"
     max_batch_size: Annotated[int, Field(gt=0)] | None = None
+    # Client HTTP timeout in seconds. OpenAI receives seconds; Gemini converts to ms.
+    timeout: float | None = None
+
+    @field_validator("timeout", mode="before")
+    @classmethod
+    def _validate_timeout(cls, v: Any) -> float | None:
+        if v is None:
+            return None
+        return coerce_provider_timeout(v)
 
     @model_validator(mode="before")
     @classmethod
@@ -431,6 +440,15 @@ class EmbeddingModelConfig(BaseModel):
     api_key: str | None = None
     base_url: str | None = None
     max_batch_size: Annotated[int, Field(gt=0)] | None = None
+    # Client HTTP timeout in seconds. OpenAI receives seconds; Gemini converts to ms.
+    timeout: float | None = None
+
+    @field_validator("timeout", mode="before")
+    @classmethod
+    def _validate_timeout(cls, v: Any) -> float | None:
+        if v is None:
+            return None
+        return coerce_provider_timeout(v)
 
     @model_validator(mode="before")
     @classmethod
@@ -556,6 +574,7 @@ def resolve_embedding_model_config(
         api_key=api_key,
         base_url=configured.overrides.base_url,
         max_batch_size=configured.max_batch_size,
+        timeout=configured.timeout,
     )
 
 

src/embedding_client.py

@@ -195,13 +195,13 @@ class _EmbeddingClient:
             from google import genai
             from google.genai import types as genai_types
 
-            # 10-minute HTTP timeout, in lockstep with the LLM registry's Gemini
-            # client (`src/llm/registry.py:_build_gemini_http_options`). Without
-            # this, a stalled Gemini embedding socket wedges the deriver worker
-            # exactly the way #785 describes for the LLM client.
+            # Default 10-minute HTTP timeout matches the LLM registry Gemini client.
+            timeout_ms = (
+                int(config.timeout * 1000) if config.timeout is not None else 600_000
+            )
             http_options = genai_types.HttpOptions(
                 base_url=config.base_url,
-                timeout=600_000,
+                timeout=timeout_ms,
             )
             self.client: genai.Client | AsyncOpenAI = genai.Client(
                 api_key=config.api_key,
@@ -216,10 +216,14 @@ class _EmbeddingClient:
                 raise ValueError("OpenAI API key is required")
             from openai import AsyncOpenAI
 
-            self.client = AsyncOpenAI(
-                api_key=config.api_key,
-                base_url=config.base_url,
-            )
+            # Omit timeout when unset so the OpenAI SDK keeps its own default.
+            client_kwargs: dict[str, Any] = {
+                "api_key": config.api_key,
+                "base_url": config.base_url,
+            }
+            if config.timeout is not None:
+                client_kwargs["timeout"] = config.timeout
+            self.client = AsyncOpenAI(**client_kwargs)
             self.max_embedding_tokens = max_input_tokens
             self.max_batch_size = config.max_batch_size or 2048
 

tests/live_llm/README.md

@@ -68,5 +68,5 @@ Coverage by provider:
 - OpenAI transport → OpenRouter non-reasoning models (e.g. `inception/mercury-2`): non-chat / diffusion architectures must stay on `max_tokens`, no `reasoning_effort`, tool-calling parameter-schema compatibility is the canary for exotic OR-served providers
 - Gemini 2.5/3.0 classes: structured outputs, cached-content reuse, thought signatures, multi-turn tool replay
 - Gemini 3.1 class: thinking and tool replay coverage by default; structured-output/caching coverage should only be added once Google documents support for that path
-- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it
+- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it. Also covers first-class `EmbeddingModelConfig.timeout` plumbing (one representative model per transport): configured timeout lands on the SDK client, and a near-zero timeout aborts before the provider answers
 - OpenAI-compatible embedding providers (e.g. OpenRouter's `google/gemini-embedding-001`): the #932 surface. Those providers reject a base64 embedding request outright (HTTP 400) or answer HTTP 200 with empty data, so the whole matrix fails without `encoding_format="float"`. Real OpenAI accepts base64 happily, so only a third-party provider catches it. Note that OpenRouter load-balances across upstreams, so the base64 failure is per-attempt rather than guaranteed: a retry can land on an endpoint that accepts it. `test_live_openai_float_encoding_matches_base64` covers the other side, that the float switch must not move vectors on real OpenAI

tests/live_llm/conftest.py

@@ -66,14 +66,23 @@ def require_embedding_key(spec: LiveEmbeddingSpec) -> str:
     return key
 
 
+_EMBEDDING_CONFIG_OVERRIDE_KEYS = frozenset({"timeout", "max_batch_size"})
+
+
 def make_embedding_client(
     spec: LiveEmbeddingSpec, **overrides: Any
 ) -> _EmbeddingClient:
     """Build a live embedding client for one matrix entry.
 
     Bypasses the `EmbeddingClient` singleton so each spec gets its own client
-    without mutating global settings.
+    without mutating global settings. `timeout` and `max_batch_size` land on
+    `EmbeddingModelConfig`; remaining kwargs go to `_EmbeddingClient`.
     """
+    config_overrides = {
+        key: overrides.pop(key)
+        for key in _EMBEDDING_CONFIG_OVERRIDE_KEYS
+        if key in overrides
+    }
     kwargs: dict[str, Any] = {
         "vector_dimensions": spec.dimensions,
         "max_input_tokens": 2048,
@@ -90,6 +99,7 @@ def make_embedding_client(
             model=spec.model,
             api_key=require_embedding_key(spec),
             base_url=spec.base_url,
+            **config_overrides,
         ),
         **kwargs,
     )

tests/live_llm/test_live_embeddings.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import time
 from typing import Any, cast
 
+import httpx
+import openai
 import pytest
 from openai import AsyncOpenAI
 
+from src.config import EmbeddingTransport
+
 from .conftest import cosine_similarity, make_embedding_client
 from .embedding_matrix import LiveEmbeddingSpec, get_live_embedding_specs
 
@@ -25,6 +30,52 @@ OPENAI_NATIVE_SPECS = tuple(
     spec for spec in ALL_SPECS if spec.family == "openai_embedding"
 )
 
+GENEROUS_TIMEOUT_SECONDS = 120.0
+TIGHT_TIMEOUT_SECONDS = 0.01
+# Well under the client defaults; generous enough to absorb SDK retries.
+TIGHT_TIMEOUT_WALL_CLOCK_LIMIT_SECONDS = 30
+
+EMBEDDING_TIMEOUT_EXCEPTIONS: dict[
+    EmbeddingTransport, tuple[type[BaseException], ...]
+] = {
+    "openai": (openai.APITimeoutError,),
+    # google-genai raises httpx or aiohttp timeouts depending on its transport;
+    # aiohttp surfaces as asyncio.TimeoutError (== builtins.TimeoutError).
+    "gemini": (httpx.TimeoutException, TimeoutError),
+}
+
+TRANSPORT_MARKS = {
+    "openai": pytest.mark.requires_openai,
+    "gemini": pytest.mark.requires_gemini,
+}
+
+
+def representative_embedding_specs() -> list[Any]:
+    """One spec per transport — timeout plumbing is client-level, not model-level."""
+    params: list[Any] = []
+    for transport in ("openai", "gemini"):
+        specs = get_live_embedding_specs(transport=transport)
+        if not specs:
+            continue
+        # Prefer the native family over openai-compatible proxies.
+        family = f"{transport}_embedding"
+        native = next((s for s in specs if s.family == family), specs[0])
+        params.append(
+            pytest.param(native, marks=TRANSPORT_MARKS[transport], id=native.id)
+        )
+    return params
+
+
+def assert_embedding_timeout_on_client(
+    client: Any, transport: EmbeddingTransport, timeout_seconds: float
+) -> None:
+    if transport == "gemini":
+        http_options = client.client._api_client._http_options
+        assert http_options.timeout == int(timeout_seconds * 1000)
+        return
+    openai_client = cast(AsyncOpenAI, client.client)
+    assert openai_client.timeout == timeout_seconds
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id)
@@ -160,3 +211,36 @@ async def test_live_gemini_batch_embed_survives_batch_split(
 
     assert len(embeddings) == len(BATCH_TEXTS)
     assert len({tuple(embedding) for embedding in embeddings}) == len(BATCH_TEXTS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("spec", representative_embedding_specs())
+async def test_live_embedding_timeout_reaches_the_client(
+    spec: LiveEmbeddingSpec,
+) -> None:
+    """Configured embedding timeout lands on the provider SDK client."""
+    client = make_embedding_client(spec, timeout=GENEROUS_TIMEOUT_SECONDS)
+
+    embedding = await client.embed(BATCH_TEXTS[0])
+
+    assert len(embedding) == spec.dimensions
+    assert_embedding_timeout_on_client(client, spec.transport, GENEROUS_TIMEOUT_SECONDS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("spec", representative_embedding_specs())
+async def test_live_tight_embedding_timeout_aborts_request(
+    spec: LiveEmbeddingSpec,
+) -> None:
+    """A near-zero embedding timeout aborts before the provider can answer."""
+    client = make_embedding_client(spec, timeout=TIGHT_TIMEOUT_SECONDS)
+
+    started = time.monotonic()
+    with pytest.raises(EMBEDDING_TIMEOUT_EXCEPTIONS[spec.transport]):
+        await client.embed(BATCH_TEXTS[0])
+    elapsed = time.monotonic() - started
+
+    assert elapsed < TIGHT_TIMEOUT_WALL_CLOCK_LIMIT_SECONDS, (
+        f"tight embedding timeout took {elapsed:.1f}s — client timeout "
+        f"likely not applied"
+    )

tests/llm/test_embedding_client.py

@@ -65,7 +65,13 @@ async def test_openai_embedding_client_uses_configured_model_and_dimensions(
     fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 8)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.api_key: str | None = api_key
             self.base_url: str | None = base_url
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
@@ -104,7 +110,13 @@ async def test_openai_embedding_client_rejects_dimension_mismatch(
     fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 7)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
 
     monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
@@ -220,6 +232,118 @@ async def test_gemini_embedding_client_keeps_timeout_without_base_url(
     assert gemini_client.http_options.timeout == 600_000
 
 
+@pytest.mark.asyncio
+async def test_openai_embedding_client_forwards_timeout(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Configured embedding timeout reaches the OpenAI-compatible client."""
+
+    class FakeOpenAIClient:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
+            self.api_key: str | None = api_key
+            self.base_url: str | None = base_url
+            self.timeout: float | None = timeout
+            self.embeddings: FakeOpenAIEmbeddingsAPI = FakeOpenAIEmbeddingsAPI(
+                [0.1] * 8
+            )
+
+    monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="openai",
+            model="text-embedding-3-small",
+            api_key="test-key",
+            timeout=45,
+        ),
+        vector_dimensions=8,
+        max_input_tokens=8192,
+        max_tokens_per_request=300_000,
+        send_dimensions=False,
+    )
+
+    openai_client = cast(Any, client.client)
+    assert openai_client.timeout == 45.0
+
+
+@pytest.mark.asyncio
+async def test_openai_embedding_client_omits_timeout_when_unset(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Unset timeout omits the kwarg so the OpenAI SDK keeps its default."""
+
+    missing = object()
+
+    class FakeOpenAIClient:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: object = missing,
+        ) -> None:
+            self.api_key: str | None = api_key
+            self.base_url: str | None = base_url
+            self.timeout: object = timeout
+            self.embeddings: FakeOpenAIEmbeddingsAPI = FakeOpenAIEmbeddingsAPI(
+                [0.1] * 8
+            )
+
+    monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="openai",
+            model="text-embedding-3-small",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=8192,
+        max_tokens_per_request=300_000,
+        send_dimensions=False,
+    )
+
+    openai_client = cast(Any, client.client)
+    assert openai_client.timeout is missing
+
+
+@pytest.mark.asyncio
+async def test_gemini_embedding_client_forwards_timeout(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Configured embedding timeout reaches Gemini as milliseconds."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str | None, http_options: Any) -> None:
+            self.api_key: str | None = api_key
+            self.http_options: Any = http_options
+            self.aio: Any = SimpleNamespace(models=SimpleNamespace())
+
+    monkeypatch.setattr("google.genai.Client", FakeGeminiClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-001",
+            api_key="gemini-key",
+            timeout=45,
+        ),
+        vector_dimensions=8,
+        max_input_tokens=4096,
+        max_tokens_per_request=300_000,
+        send_dimensions=False,
+    )
+
+    gemini_client = cast(Any, client.client)
+    assert gemini_client.http_options.timeout == 45_000
+
+
 def _build_openai_client(
     monkeypatch: pytest.MonkeyPatch,
     *,
@@ -233,7 +357,13 @@ def _build_openai_client(
     fake_embeddings = FakeOpenAIEmbeddingsAPI(embedding)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.api_key: str | None = api_key
             self.base_url: str | None = base_url
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
@@ -707,7 +837,13 @@ async def test_simple_batch_embed_respects_token_budget_per_request(
     fake_embeddings = FakeOpenAIEmbeddingsAPI([0.5] * 4)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
 
     monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
@@ -745,7 +881,13 @@ async def test_simple_batch_embed_rejects_oversized_input(
     fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
 
     monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
@@ -775,7 +917,13 @@ def test_prepare_chunks_returns_ordered_chunks(
     fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4)
 
     class FakeOpenAIClient:
-        def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
+        def __init__(
+            self,
+            *,
+            api_key: str | None,
+            base_url: str | None,
+            timeout: float | None = None,
+        ) -> None:
             self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
 
     monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient)
@@ -818,6 +966,31 @@ def test_embedding_model_config_parses_max_batch_size_from_env(
     assert resolved.max_batch_size == 10
 
 
+def test_embedding_model_config_parses_timeout_from_env(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    s = _build_embedding_settings(
+        {"EMBEDDING_MODEL_CONFIG__TIMEOUT": "90.0"},
+        monkeypatch,
+    )
+
+    assert s.MODEL_CONFIG.timeout == 90.0
+
+    resolved = resolve_embedding_model_config(s.MODEL_CONFIG)
+    assert resolved.timeout == 90.0
+
+
+def test_embedding_model_config_rejects_invalid_timeout() -> None:
+    with pytest.raises(
+        ValueError, match=r"provider_params\.timeout must be a positive number"
+    ):
+        EmbeddingModelConfig(
+            transport="openai",
+            model="text-embedding-3-small",
+            timeout=-1,
+        )
+
+
 @pytest.mark.asyncio
 async def test_gemini_process_batch_wraps_contents_as_content_part(
     monkeypatch: pytest.MonkeyPatch,