From 59e645eba5666dce7f75d595cd1823a0afc26cbc Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 6 Oct 2026 15:28:36 -0400 Subject: [PATCH] feat(llm): ban OpenRouter providers that time out Add STRIX_PROVIDER_TTFT_TIMEOUT (first event) and STRIX_PROVIDER_REQUEST_TIMEOUT (whole stream) next to LLM_STREAM_IDLE_TIMEOUT. When any of them fires, or OpenRouter reports error_type=timeout, ban the provider that served the attempt (when a chunk named it), then raise as before so the transient retry replays the turn. The timers cancel with a strix: message, so the request log tells a stream timeout from a shutdown. --- strix/config/models.py | 116 ++++++++++++++++++++++---- strix/config/settings.py | 4 + strix/llm/request_log.py | 18 +++- tests/test_stream_idle_timeout.py | 131 ++++++++++++++++++++++++++++-- 4 files changed, 242 insertions(+), 27 deletions(-) diff --git a/strix/config/models.py b/strix/config/models.py index 063794fd..b84e2aa2 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -263,18 +263,28 @@ class _TurnGuardModel(Model): not covered by the request timeout, which resets on any byte (keepalives included). ``LLM_STREAM_IDLE_TIMEOUT`` bounds the gap between events so the turn fails instead of hanging, and the existing retry path replays it. + ``LLM_STREAM_FIRST_EVENT_TIMEOUT`` and ``LLM_STREAM_TOTAL_TIMEOUT`` bound + the first event and the whole stream. When one of these fires, or OpenRouter + reports an upstream timeout, the OpenRouter provider serving the attempt is + banned (if a chunk named it) before the retry. """ def __init__( self, inner: Model, *, + model_name: str | None = None, max_tool_calls_per_turn: int = 0, stream_idle_timeout: float = 0.0, + first_event_timeout: float = 0.0, + total_timeout: float = 0.0, ) -> None: self._inner = inner + self._model_name = model_name self._max_tool_calls_per_turn = max_tool_calls_per_turn self._stream_idle_timeout = stream_idle_timeout + self._first_event_timeout = first_event_timeout + self._total_timeout = total_timeout def _limiter(self) -> TurnToolCallLimiter: return TurnToolCallLimiter(self._max_tool_calls_per_turn) @@ -361,10 +371,32 @@ class _TurnGuardModel(Model): conversation_id=conversation_id, prompt=prompt, ) - async for event in _with_idle_timeout(stream, self._stream_idle_timeout): - guarded = _guard_event(event, rewriter, limiter) - if guarded is not None: - yield guarded + with request_log.capture_reply() as reply: + try: + async for event in _with_timeouts( + stream, + idle=self._stream_idle_timeout, + first_event=self._first_event_timeout, + total=self._total_timeout, + ): + guarded = _guard_event(event, rewriter, limiter) + if guarded is not None: + yield guarded + except Exception as exc: + if isinstance(exc, _StreamTimeoutError): + reason = exc.reason + elif reply.upstream_error_type == "timeout": + reason = "upstream_timeout" + else: + raise + if reply.upstream_provider and self._model_name: + await asyncio.to_thread( + provider_bans.ban_provider, + self._model_name, + reply.upstream_provider, + reason, + ) + raise self._log_dropped(limiter) @@ -395,28 +427,75 @@ async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None: await stream.aclose() -async def _with_idle_timeout( - stream: AsyncIterator[TResponseStreamEvent], timeout: float -) -> AsyncIterator[TResponseStreamEvent]: - if timeout <= 0: - async for event in stream: - yield event - return +class _StreamTimeoutError(TimeoutError): + def __init__(self, reason: str, message: str = "") -> None: + super().__init__(message or reason) + self.reason = reason + +async def _with_timeouts( + stream: AsyncIterator[TResponseStreamEvent], + *, + idle: float, + first_event: float = 0.0, + total: float = 0.0, +) -> AsyncIterator[TResponseStreamEvent]: + """Bound the first event, the gap between events and the whole stream; 0 is no bound.""" iterator = stream.__aiter__() + deadline = time.monotonic() + total + first = True while True: + limits: list[tuple[float, float, str]] = [] + gap = first_event if first else idle + if gap: + limits.append( + (gap, gap, "stream_first_event_timeout" if first else "stream_idle_timeout") + ) + if total: + limits.append((max(0.0, deadline - time.monotonic()), total, "stream_total_timeout")) try: - event = await asyncio.wait_for(iterator.__anext__(), timeout) + if not limits: + event = await iterator.__anext__() + else: + wait, limit, reason = min(limits) + event = await _next_event(iterator, wait, reason) except StopAsyncIteration: return - except TimeoutError: + except _StreamTimeoutError: await _aclose(stream) - message = f"model stream produced no event for {timeout:.0f}s" + message = f"model stream hit {reason} ({limit:.0f}s)" logger.warning("%s; abandoning the turn", message) - raise TimeoutError(message) from None + raise _StreamTimeoutError(reason, message) from None + first = False yield event +async def _next_event( + iterator: AsyncIterator[TResponseStreamEvent], timeout: float, reason: str +) -> TResponseStreamEvent: + # asyncio.timeout() cancels without a message, so the request log could not + # tell this from a shutdown; this is asyncio.timeout() with a message. + task = asyncio.current_task() + assert task is not None + cancelling = task.cancelling() + expired = False + + def expire() -> None: + nonlocal expired + expired = True + task.cancel(msg=f"strix:{reason}") + + handle = asyncio.get_running_loop().call_later(timeout, expire) + try: + return await iterator.__anext__() + except BaseException as exc: + if expired and task.uncancel() <= cancelling and isinstance(exc, asyncio.CancelledError): + raise _StreamTimeoutError(reason) from None + raise + finally: + handle.cancel() + + def _guard_event( event: TResponseStreamEvent, rewriter: TurnCallIdRewriter, limiter: TurnToolCallLimiter ) -> TResponseStreamEvent | None: @@ -552,6 +631,8 @@ class StrixProvider(MultiProvider): llm = load_settings().llm slug = codex.subscription_model(model_name) idle_timeout = float(llm.stream_idle_timeout) + first_event_timeout = float(llm.stream_first_event_timeout) + total_timeout = float(llm.stream_total_timeout) if slug: # The ChatGPT subscription backend is always streamed; it has no # non-streaming mode to fall back to, so LLM_DISABLE_STREAMING @@ -592,11 +673,14 @@ class StrixProvider(MultiProvider): # The wrapper emits its single event only once the whole request # is done, so an idle gap is meaningless here; the request # timeout bounds it instead. - idle_timeout = 0.0 + idle_timeout = first_event_timeout = total_timeout = 0.0 return _TurnGuardModel( model, + model_name=model_name or llm.model, max_tool_calls_per_turn=llm.max_tool_calls_per_turn, stream_idle_timeout=idle_timeout, + first_event_timeout=first_event_timeout, + total_timeout=total_timeout, ) diff --git a/strix/config/settings.py b/strix/config/settings.py index a471286e..ab082de6 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -81,6 +81,10 @@ class LlmSettings(BaseSettings): timeout: int = Field(default=300, alias="LLM_TIMEOUT") preflight_timeout: int = Field(default=30, ge=1, alias="LLM_PREFLIGHT_TIMEOUT") stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT") + stream_first_event_timeout: int = Field( + default=120, ge=0, alias="LLM_STREAM_FIRST_EVENT_TIMEOUT" + ) + stream_total_timeout: int = Field(default=600, ge=0, alias="LLM_STREAM_TOTAL_TIMEOUT") max_tool_calls_per_turn: int = Field( default=32, ge=0, diff --git a/strix/llm/request_log.py b/strix/llm/request_log.py index 99f313e1..53294780 100644 --- a/strix/llm/request_log.py +++ b/strix/llm/request_log.py @@ -30,7 +30,7 @@ import logging import re import time import uuid -from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import asdict, dataclass, replace from datetime import UTC, date, datetime @@ -198,6 +198,17 @@ def failed_reply(exc: BaseException) -> HttpReply | None: return reply if isinstance(reply, HttpReply) else None +@contextlib.contextmanager +def capture_reply() -> Iterator[HttpReply]: + """Open the attempt's reply here so it outlives the logging wrapper inside.""" + reply = HttpReply() + token = _http_reply.set(reply) + try: + yield reply + finally: + _reset_http_reply(token) + + async def record_http_reply(response: Response) -> None: """httpx ``response`` event hook: remember the reply for the attempt in flight. @@ -1280,7 +1291,7 @@ class RequestLoggingModel(Model): raw_response: object = None first_event_mono: float | None = None finish_reason: str | None = None - reply = HttpReply() + reply = _http_reply.get() or HttpReply() token = _http_reply.set(reply) try: async for event in self._inner.stream_response( @@ -1363,7 +1374,8 @@ def _is_abandonment(exc: BaseException | None) -> bool: def _abandonment_message(exc: BaseException) -> str | None: if isinstance(exc, asyncio.CancelledError): - return "attempt cancelled before the reply was consumed (stream idle timeout or shutdown)" + reason = f": {exc.args[0]}" if exc.args else "" + return f"attempt cancelled before the reply was consumed{reason}" if isinstance(exc, GeneratorExit): return "stream closed by the caller before it finished" return None diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py index 9d618978..8ee47cba 100644 --- a/tests/test_stream_idle_timeout.py +++ b/tests/test_stream_idle_timeout.py @@ -23,7 +23,8 @@ from openai import AsyncOpenAI from strix.config import loader from strix.config.loader import load_settings -from strix.config.models import StrixProvider, _TurnGuardModel, _with_idle_timeout +from strix.config.models import StrixProvider, _TurnGuardModel, _with_timeouts +from strix.llm import request_log if TYPE_CHECKING: @@ -80,9 +81,18 @@ def stalling_gateway() -> Iterator[str]: def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]: client = AsyncOpenAI(api_key="tok", base_url=base_url, max_retries=0, timeout=_STALL_SECONDS) - inner: Model = OpenAIChatCompletionsModel(model="gw-model", openai_client=client) + inner: Model = request_log.RequestLoggingModel( + OpenAIChatCompletionsModel(model="gw-model", openai_client=client), + model_name="gw-model", + provider="openai", + base_url=base_url, + ) guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout) - return guarded.stream_response( + return _call(guarded) + + +def _call(model: Model) -> AsyncIterator[Any]: + return model.stream_response( None, "go", ModelSettings(), @@ -110,11 +120,20 @@ async def test_stalled_stream_hangs_without_the_watchdog(stalling_gateway: str) @pytest.mark.asyncio async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str) -> None: + logged: list[request_log.LlmRequestEvent] = [] + request_log.register_sink(logged.append) started = time.monotonic() - with pytest.raises(TimeoutError, match="produced no event"): - await _drain(stalling_gateway, idle_timeout=1) + try: + with pytest.raises(TimeoutError, match="stream_idle_timeout"): + await _drain(stalling_gateway, idle_timeout=1) + finally: + request_log.unregister_sink(logged.append) assert time.monotonic() - started < _STALL_SECONDS + # The request log can tell our timeout from a shutdown. + assert [event.error_message for event in logged] == [ + "attempt cancelled before the reply was consumed: strix:stream_idle_timeout" + ] @pytest.mark.asyncio @@ -124,14 +143,106 @@ async def test_events_keep_flowing_while_the_stream_is_alive() -> None: await asyncio.sleep(0.05) yield f"event-{i}" - seen: list[Any] = [event async for event in _with_idle_timeout(_live(), 1.0)] + seen: list[Any] = [event async for event in _with_timeouts(_live(), idle=1.0)] assert seen == [f"event-{i}" for i in range(5)] +@pytest.mark.asyncio +async def test_first_event_and_whole_stream_are_bounded() -> None: + async def _slow_start() -> AsyncIterator[Any]: + await asyncio.sleep(_STALL_SECONDS) + yield "late" + + async def _endless() -> AsyncIterator[Any]: + while True: + await asyncio.sleep(0.05) + yield "event" + + with pytest.raises(TimeoutError, match="stream_first_event_timeout"): + [event async for event in _with_timeouts(_slow_start(), idle=5, first_event=0.2)] + with pytest.raises(TimeoutError, match="stream_total_timeout"): + [event async for event in _with_timeouts(_endless(), idle=5, total=0.3)] + + +class _UpstreamInner(Model): + """Records the upstream provider like the OpenRouter chunk parser, then fails.""" + + def __init__(self, provider: str | None, error_type: str | None = None) -> None: + self._provider = provider + self._error_type = error_type + + async def get_response(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + async def stream_response(self, *_: Any, **__: Any) -> AsyncIterator[Any]: + if self._error_type: + request_log.record_upstream_provider(self._provider, self._error_type) + raise RuntimeError("upstream error chunk") + await asyncio.sleep(0.05) + request_log.record_upstream_provider(self._provider) + yield "event" + await asyncio.sleep(_STALL_SECONDS) + + +@pytest.fixture +def bans(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, str, str]]: + calls: list[tuple[str, str, str]] = [] + monkeypatch.setattr( + "strix.llm.provider_bans.ban_provider", lambda *args: calls.append(args) or True + ) + return calls + + +async def _drain_guarded(inner: Model, **timeouts: float) -> None: + logging_model = request_log.RequestLoggingModel( + inner, model_name="openrouter/z-ai/glm-5.3", provider="openrouter", base_url=None + ) + guarded = _TurnGuardModel(logging_model, model_name="openrouter/z-ai/glm-5.3", **timeouts) + [event async for event in _call(guarded)] + + +@pytest.mark.asyncio +async def test_stream_timeout_bans_the_provider_that_served_it( + bans: list[tuple[str, str, str]], +) -> None: + with pytest.raises(TimeoutError): + await _drain_guarded(_UpstreamInner("InferenceNet"), stream_idle_timeout=0.2) + + assert bans == [("openrouter/z-ai/glm-5.3", "InferenceNet", "stream_idle_timeout")] + + +@pytest.mark.asyncio +async def test_openrouter_upstream_timeout_bans_the_provider( + bans: list[tuple[str, str, str]], +) -> None: + with pytest.raises(RuntimeError): + await _drain_guarded(_UpstreamInner("InferenceNet", "timeout")) + with pytest.raises(RuntimeError): + await _drain_guarded(_UpstreamInner("InferenceNet", "provider_unavailable")) + + assert bans == [("openrouter/z-ai/glm-5.3", "InferenceNet", "upstream_timeout")] + + +@pytest.mark.asyncio +async def test_timeout_without_a_known_provider_bans_nothing( + bans: list[tuple[str, str, str]], +) -> None: + with pytest.raises(TimeoutError, match="stream_first_event_timeout"): + await _drain_guarded(_UpstreamInner(None), first_event_timeout=0.01) + + assert bans == [] + + @pytest.fixture def _reset_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - for key in ("STRIX_LLM", "LLM_DISABLE_STREAMING", "LLM_STREAM_IDLE_TIMEOUT"): + for key in ( + "STRIX_LLM", + "LLM_DISABLE_STREAMING", + "LLM_STREAM_IDLE_TIMEOUT", + "LLM_STREAM_FIRST_EVENT_TIMEOUT", + "LLM_STREAM_TOTAL_TIMEOUT", + ): monkeypatch.delenv(key, raising=False) monkeypatch.setattr(loader, "_cached", None) monkeypatch.setattr(loader, "_override", None) @@ -151,11 +262,15 @@ def test_idle_timeout_is_configurable( ) -> None: monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel()) monkeypatch.setenv("LLM_STREAM_IDLE_TIMEOUT", "45") + monkeypatch.setenv("LLM_STREAM_FIRST_EVENT_TIMEOUT", "30") + monkeypatch.setenv("LLM_STREAM_TOTAL_TIMEOUT", "400") load_settings() model = StrixProvider().get_model("openai/gpt-4o-mini") assert isinstance(model, _TurnGuardModel) assert model._stream_idle_timeout == 45 + assert model._first_event_timeout == 30 + assert model._total_timeout == 400 def test_idle_timeout_is_off_without_streaming( @@ -170,4 +285,4 @@ def test_idle_timeout_is_off_without_streaming( model = StrixProvider().get_model("openai/gpt-4o-mini") assert isinstance(model, _TurnGuardModel) - assert model._stream_idle_timeout == 0 + assert model._stream_idle_timeout == model._first_event_timeout == model._total_timeout == 0