fix(llm): set HTTP timeout on Gemini clients (#903)

* fix(llm): set HTTP timeout on Gemini clients (#785)

* fix(embedding): set HTTP timeout on Gemini embedding client (#785)

Same wedge-class failure as the LLM client: a stalled Gemini embedding
socket hangs the in-process reconciler, which shares the deriver worker's
uvloop event loop. Apply the same 10-minute timeout here, in lockstep
with src/llm/registry.py's _build_gemini_http_options.

* style(test): drop extra blank line in test_registry imports
This commit is contained in:
PK 2026-08-04 13:29:32 -07:00 committed by GitHub
parent d815c8b8dc
commit 5c32bd10ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 223 additions and 34 deletions

View File

@ -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,

View File

@ -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),
)

View File

@ -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,
*,

View File

@ -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"