diff --git a/src/embedding_client.py b/src/embedding_client.py index 07197f43..208d9d71 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -182,10 +182,13 @@ class _EmbeddingClient: if self.transport == "gemini": if not config.api_key: raise ValueError("Gemini API key is required") - http_options = ( - genai_types.HttpOptions(base_url=config.base_url) - if config.base_url - else None + # 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. + http_options = genai_types.HttpOptions( + base_url=config.base_url, + timeout=600_000, ) self.client: genai.Client | AsyncOpenAI = genai.Client( api_key=config.api_key, diff --git a/src/llm/registry.py b/src/llm/registry.py index 6dd5eecc..48c8191d 100644 --- a/src/llm/registry.py +++ b/src/llm/registry.py @@ -32,6 +32,14 @@ from .history_adapters import ( ) from .types import ProviderClient +# Default client-level HTTP timeouts. Anthropic accepts seconds (float); +# google-genai's HttpOptions.timeout is an int in milliseconds, so the Gemini +# value is kept separately. Both default to 10 minutes to match the existing +# Anthropic behavior — long enough for slow streamed responses, short enough +# that a stalled socket can no longer wedge the deriver worker (see #785). +_ANTHROPIC_TIMEOUT_S = 600.0 +_GEMINI_TIMEOUT_MS = 600_000 + # Client-level ``default_headers`` applied to OpenAI-compatible clients, keyed by # base-URL prefix. Currently only OpenRouter, which uses them for app attribution # (https://openrouter.ai/docs/app-attribution); add a prefix here to tag another @@ -54,13 +62,28 @@ def _default_headers_for(base_url: str | None) -> dict[str, str]: return {} +def _build_gemini_http_options(base_url: str | None) -> genai_types.HttpOptions: + """Build Gemini ``HttpOptions`` carrying a default HTTP timeout. + + google-genai's ``HttpOptions.timeout`` is an int in milliseconds. A stalled + Gemini socket without this value wedges the entire deriver process because + all deriver workers share one uvloop event loop (see #785). Keep the + timeout even when no ``base_url`` is configured — that's the path the + default ``get_gemini_client`` takes and it's the one that was hanging. + """ + return genai_types.HttpOptions( + base_url=base_url, + timeout=_GEMINI_TIMEOUT_MS, + ) + + @lru_cache(maxsize=1) def get_anthropic_client() -> AsyncAnthropic: """Default Anthropic client built from settings.LLM.ANTHROPIC_API_KEY.""" return AsyncAnthropic( api_key=settings.LLM.ANTHROPIC_API_KEY, base_url=settings.LLM.ANTHROPIC_BASE_URL, - timeout=600.0, + timeout=_ANTHROPIC_TIMEOUT_S, ) @@ -77,12 +100,10 @@ def get_openai_client() -> AsyncOpenAI: @lru_cache(maxsize=1) def get_gemini_client() -> genai.Client: """Default Gemini client built from settings.LLM.GEMINI_API_KEY.""" - http_options = ( - genai_types.HttpOptions(base_url=settings.LLM.GEMINI_BASE_URL) - if settings.LLM.GEMINI_BASE_URL - else None + return genai.Client( + api_key=settings.LLM.GEMINI_API_KEY, + http_options=_build_gemini_http_options(settings.LLM.GEMINI_BASE_URL), ) - return genai.Client(api_key=settings.LLM.GEMINI_API_KEY, http_options=http_options) # Bounded cache — in practice the (base_url, api_key) key space is small @@ -105,7 +126,9 @@ def get_anthropic_override_client( api_key: str | None, ) -> AsyncAnthropic: """Anthropic client for a specific (base_url, api_key) pair. Cached by key.""" - return AsyncAnthropic(api_key=api_key, base_url=base_url, timeout=600.0) + return AsyncAnthropic( + api_key=api_key, base_url=base_url, timeout=_ANTHROPIC_TIMEOUT_S + ) @lru_cache(maxsize=128) @@ -113,8 +136,10 @@ def get_gemini_override_client( base_url: str | None, api_key: str | None ) -> genai.Client: """Gemini client for a specific (base_url, api_key) pair. Cached by key.""" - http_options = genai_types.HttpOptions(base_url=base_url) if base_url else None - return genai.Client(api_key=api_key, http_options=http_options) + return genai.Client( + api_key=api_key, + http_options=_build_gemini_http_options(base_url), + ) # Module-level default-client registry, populated at import time. Tests patch @@ -125,7 +150,7 @@ if settings.LLM.ANTHROPIC_API_KEY: CLIENTS["anthropic"] = AsyncAnthropic( api_key=settings.LLM.ANTHROPIC_API_KEY, base_url=settings.LLM.ANTHROPIC_BASE_URL, - timeout=600.0, + timeout=_ANTHROPIC_TIMEOUT_S, ) if settings.LLM.OPENAI_API_KEY: @@ -136,14 +161,9 @@ if settings.LLM.OPENAI_API_KEY: ) if settings.LLM.GEMINI_API_KEY: - http_options = ( - genai_types.HttpOptions(base_url=settings.LLM.GEMINI_BASE_URL) - if settings.LLM.GEMINI_BASE_URL - else None - ) CLIENTS["gemini"] = genai.Client( api_key=settings.LLM.GEMINI_API_KEY, - http_options=http_options, + http_options=_build_gemini_http_options(settings.LLM.GEMINI_BASE_URL), ) diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index a4642159..48587d43 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest @@ -141,6 +141,12 @@ async def test_gemini_embedding_client_uses_output_dimensionality( embedding = await client.embed("hello world") assert embedding == [0.2] * 12 + # 10-minute HTTP timeout, in lockstep with the LLM registry's Gemini client + # (see #785). Without this, a stalled Gemini embedding socket wedges the + # deriver worker — the same failure mode the LLM fix addresses. + gemini_client = cast(Any, client.client) + assert gemini_client.http_options.base_url == "https://gemini-proxy.example/v1beta" + assert gemini_client.http_options.timeout == 600_000 assert calls == [ { "model": "gemini-embedding-001", @@ -150,6 +156,49 @@ async def test_gemini_embedding_client_uses_output_dimensionality( ] +@pytest.mark.asyncio +async def test_gemini_embedding_client_keeps_timeout_without_base_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No-base-url Gemini embedding client must still carry an HTTP timeout.""" + + class FakeGeminiModels: + async def embed_content( + self, + *, + model: str, + contents: str, + config: dict[str, Any], + ) -> SimpleNamespace: + return SimpleNamespace( + embeddings=[SimpleNamespace(values=[0.1] * 8)], + ) + + class FakeGeminiClient: + def __init__(self, *, api_key: str | None, http_options: Any) -> None: + self.api_key = api_key + self.http_options = http_options + self.aio = SimpleNamespace(models=FakeGeminiModels()) + + monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="gemini", + model="gemini-embedding-001", + api_key="gemini-key", + ), + 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.base_url is None + assert gemini_client.http_options.timeout == 600_000 + + def _build_openai_client( monkeypatch: pytest.MonkeyPatch, *, diff --git a/tests/llm/test_registry.py b/tests/llm/test_registry.py index 71f3302e..b6df137b 100644 --- a/tests/llm/test_registry.py +++ b/tests/llm/test_registry.py @@ -1,22 +1,139 @@ -"""Tests for src.llm.registry helpers.""" +"""Tests for the provider-client registry in src/llm/registry.py. + +Locks the HTTP-timeout behavior added for the Gemini transport (#785) and +pins the existing 600s timeout on the Anthropic clients so regressions on +either side are caught at unit-test time. +""" from __future__ import annotations -from src.llm.registry import _default_headers_for # pyright: ignore[reportPrivateUsage] +from collections.abc import Iterator +from unittest.mock import patch + +import pytest +from google.genai import types as genai_types + +from src import config as app_config +from src.llm import registry as registry_module + +# Gemini's HttpOptions.timeout is an int in milliseconds; keep it in lockstep +# with the Anthropic client's 600s timeout to match the rest of the registry. +_GEMINI_TIMEOUT_MS = 600_000 +_ANTHROPIC_TIMEOUT_S = 600.0 -def test_default_headers_for_openrouter_base_url() -> None: - """OpenRouter base URLs get the app-attribution headers.""" - headers = _default_headers_for("https://openrouter.ai/api/v1") - assert headers["HTTP-Referer"] == "https://honcho.dev" - assert headers["X-Openrouter-Title"] == "Honcho" +@pytest.fixture(autouse=True) +def _patch_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + """Default the LLM settings so the registry reads valid values.""" + monkeypatch.setenv("PYTHON_DOTENV_DISABLED", "1") + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-anthropic-key") + monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-key") + yield -def test_default_headers_for_non_openrouter_base_url() -> None: - """Other OpenAI-compatible providers get no extra headers.""" - assert _default_headers_for("https://api.openai.com/v1") == {} +@pytest.fixture +def _fresh_lru_caches() -> Iterator[None]: + """Drop lru_cache state so each test exercises a fresh client build.""" + registry_module.get_anthropic_client.cache_clear() + registry_module.get_gemini_client.cache_clear() + registry_module.get_anthropic_override_client.cache_clear() + registry_module.get_gemini_override_client.cache_clear() + yield + registry_module.get_anthropic_client.cache_clear() + registry_module.get_gemini_client.cache_clear() + registry_module.get_anthropic_override_client.cache_clear() + registry_module.get_gemini_override_client.cache_clear() -def test_default_headers_for_none_base_url() -> None: - """A missing base URL (default OpenAI) gets no extra headers.""" - assert _default_headers_for(None) == {} +def test_get_gemini_client_sets_http_timeout( + _fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Default Gemini client must carry an HttpOptions timeout, not None.""" + monkeypatch.setattr(app_config.settings.LLM, "GEMINI_BASE_URL", None) + + with patch("src.llm.registry.genai.Client") as mock_client: + registry_module.get_gemini_client() + + assert mock_client.call_count == 1 + http_options = mock_client.call_args.kwargs["http_options"] + assert isinstance(http_options, genai_types.HttpOptions) + assert http_options.timeout == _GEMINI_TIMEOUT_MS + + +def test_get_gemini_client_preserves_custom_base_url( + _fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Base URL and timeout must coexist on the default Gemini client.""" + monkeypatch.setattr( + app_config.settings.LLM, "GEMINI_BASE_URL", "https://gemini-proxy.example.com" + ) + + with patch("src.llm.registry.genai.Client") as mock_client: + registry_module.get_gemini_client() + + http_options = mock_client.call_args.kwargs["http_options"] + assert isinstance(http_options, genai_types.HttpOptions) + assert http_options.base_url == "https://gemini-proxy.example.com" + assert http_options.timeout == _GEMINI_TIMEOUT_MS + + +def test_get_gemini_override_client_sets_http_timeout( + _fresh_lru_caches: None, +) -> None: + """Override Gemini client must also carry a timeout.""" + with patch("src.llm.registry.genai.Client") as mock_client: + registry_module.get_gemini_override_client( + "https://gemini-proxy.example.com", "sk-override" + ) + + http_options = mock_client.call_args.kwargs["http_options"] + assert isinstance(http_options, genai_types.HttpOptions) + assert http_options.base_url == "https://gemini-proxy.example.com" + assert http_options.timeout == _GEMINI_TIMEOUT_MS + + +def test_get_gemini_override_client_handles_missing_base_url( + _fresh_lru_caches: None, +) -> None: + """Override Gemini client with no base URL still carries a timeout.""" + with patch("src.llm.registry.genai.Client") as mock_client: + registry_module.get_gemini_override_client(None, "sk-override") + + http_options = mock_client.call_args.kwargs["http_options"] + assert isinstance(http_options, genai_types.HttpOptions) + assert http_options.timeout == _GEMINI_TIMEOUT_MS + + +def test_get_anthropic_client_keeps_600s_timeout( + _fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Anthropic timeout is the established behavior — lock it.""" + monkeypatch.setattr(app_config.settings.LLM, "ANTHROPIC_BASE_URL", None) + + with patch("src.llm.registry.AsyncAnthropic") as mock_anthropic: + registry_module.get_anthropic_client() + + assert mock_anthropic.call_args.kwargs["timeout"] == _ANTHROPIC_TIMEOUT_S + + +def test_get_anthropic_override_client_keeps_600s_timeout( + _fresh_lru_caches: None, +) -> None: + """Override Anthropic client also keeps the 600s timeout.""" + with patch("src.llm.registry.AsyncAnthropic") as mock_anthropic: + registry_module.get_anthropic_override_client(None, "sk-override") + + assert mock_anthropic.call_args.kwargs["timeout"] == _ANTHROPIC_TIMEOUT_S + + +def test_gemini_http_options_builder_applies_timeout() -> None: + """The shared helper must always set a timeout, even with no base_url.""" + options = registry_module._build_gemini_http_options(None) + assert isinstance(options, genai_types.HttpOptions) + assert options.timeout == _GEMINI_TIMEOUT_MS + assert options.base_url is None + + options = registry_module._build_gemini_http_options("https://example.com") + assert isinstance(options, genai_types.HttpOptions) + assert options.timeout == _GEMINI_TIMEOUT_MS + assert options.base_url == "https://example.com"