From 4b8212a447a10b3f5944181b8d578f325947e338 Mon Sep 17 00:00:00 2001 From: Ian Date: Wed, 7 Oct 2026 11:47:18 -0400 Subject: [PATCH] refactor(llm): ban timed-out providers in the existing retry path --- strix/config/models.py | 84 +++++++++---------------- strix/core/execution.py | 7 ++- strix/llm/request_log.py | 15 +---- tests/test_execution_transient_retry.py | 19 ++++++ tests/test_stream_idle_timeout.py | 79 ++--------------------- 5 files changed, 61 insertions(+), 143 deletions(-) diff --git a/strix/config/models.py b/strix/config/models.py index cb24c2db..a4396f22 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -264,23 +264,19 @@ class _TurnGuardModel(Model): 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. + the first event and the whole stream. """ 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 @@ -371,32 +367,15 @@ class _TurnGuardModel(Model): conversation_id=conversation_id, prompt=prompt, ) - 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 + 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 self._log_dropped(limiter) @@ -427,9 +406,9 @@ async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None: await stream.aclose() -class _StreamTimeoutError(TimeoutError): - def __init__(self, reason: str, message: str = "") -> None: - super().__init__(message or reason) +class StreamTimeoutError(TimeoutError): + def __init__(self, reason: str, message: str) -> None: + super().__init__(message) self.reason = reason @@ -443,37 +422,32 @@ async def _with_timeouts( """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]] = [] # Without a first-event bound, the idle bound covers the first event too. - on_first = first and bool(first_event) - gap = first_event if on_first else idle - if gap: - limits.append( - (gap, gap, "stream_first_event_timeout" if on_first else "stream_idle_timeout") - ) + if first_event: + limits.append((first_event, first_event, "stream_first_event_timeout")) + elif idle: + limits.append((idle, idle, "stream_idle_timeout")) if total: limits.append((max(0.0, deadline - time.monotonic()), total, "stream_total_timeout")) try: - if not limits: - event = await iterator.__anext__() + if limits: + event = await _next_event(iterator, *min(limits)) else: - wait, limit, reason = min(limits) - event = await _next_event(iterator, wait, reason) + event = await iterator.__anext__() except StopAsyncIteration: return - except _StreamTimeoutError: + except StreamTimeoutError as exc: await _aclose(stream) - message = f"model stream hit {reason} ({limit:.0f}s)" - logger.warning("%s; abandoning the turn", message) - raise _StreamTimeoutError(reason, message) from None - first = False + logger.warning("%s; abandoning the turn", exc) + raise + first_event = 0.0 yield event async def _next_event( - iterator: AsyncIterator[TResponseStreamEvent], timeout: float, reason: str + iterator: AsyncIterator[TResponseStreamEvent], wait: float, limit: 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. @@ -487,12 +461,15 @@ async def _next_event( expired = True task.cancel(msg=f"strix:{reason}") - handle = asyncio.get_running_loop().call_later(timeout, expire) + handle = asyncio.get_running_loop().call_later(wait, 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 + error = StreamTimeoutError(reason, f"model stream hit {reason} ({limit:.0f}s)") + # Keep the attempt's reply: it names the provider to ban. + setattr(error, request_log._FAILED_REPLY_ATTR, request_log.failed_reply(exc)) + raise error from None raise finally: handle.cancel() @@ -678,7 +655,6 @@ class StrixProvider(MultiProvider): 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, diff --git a/strix/core/execution.py b/strix/core/execution.py index c1d9b621..da186c47 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -22,6 +22,7 @@ from openai import ( from strix.config import codex from strix.config.loader import load_settings +from strix.config.models import StreamTimeoutError from strix.core.hooks import ( BudgetExceededError, BudgetPausedError, @@ -162,7 +163,11 @@ async def _maybe_ban_provider(exc: BaseException, run_config: RunConfig) -> None return provider = reply.upstream_provider code = _model_error_status_code(exc) - if reply.upstream_error_type == "provider_unavailable": + if isinstance(exc, StreamTimeoutError): + reason = exc.reason + elif reply.upstream_error_type == "timeout": + reason = "upstream_timeout" + elif reply.upstream_error_type == "provider_unavailable": reason = "provider_unavailable" elif reply.upstream_rejected: reason = f"upstream_rejected_{code}" diff --git a/strix/llm/request_log.py b/strix/llm/request_log.py index 53294780..8f48dc80 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, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import asdict, dataclass, replace from datetime import UTC, date, datetime @@ -198,17 +198,6 @@ 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. @@ -1291,7 +1280,7 @@ class RequestLoggingModel(Model): raw_response: object = None first_event_mono: float | None = None finish_reason: str | None = None - reply = _http_reply.get() or HttpReply() + reply = HttpReply() token = _http_reply.set(reply) try: async for event in self._inner.stream_response( diff --git a/tests/test_execution_transient_retry.py b/tests/test_execution_transient_retry.py index aa43fb50..1f6da726 100644 --- a/tests/test_execution_transient_retry.py +++ b/tests/test_execution_transient_retry.py @@ -17,6 +17,7 @@ from openai import ( ) from strix.config import codex +from strix.config.models import StreamTimeoutError from strix.core import execution from strix.core.agents import AgentCoordinator from strix.llm import provider_bans, request_log @@ -149,6 +150,24 @@ async def test_bans_upstream_after_repeated_mid_stream_4xx( assert bans == [("openrouter/z-ai/glm-5.3", "Relace", "repeated_4xx")] +@pytest.mark.asyncio +async def test_bans_upstream_right_away_on_timeouts(bans: list[tuple[str, str, str]]) -> None: + await execution._maybe_ban_provider( + _upstream_error(408, "Relace", error_type="timeout"), _OPENROUTER_RUN_CONFIG + ) + stalled = StreamTimeoutError("stream_idle_timeout", "stalled") + setattr( + stalled, + request_log._FAILED_REPLY_ATTR, + request_log.HttpReply(upstream_provider="InferenceNet"), + ) + await execution._maybe_ban_provider(stalled, _OPENROUTER_RUN_CONFIG) + assert bans == [ + ("openrouter/z-ai/glm-5.3", "Relace", "upstream_timeout"), + ("openrouter/z-ai/glm-5.3", "InferenceNet", "stream_idle_timeout"), + ] + + class _FakeStream: def __init__(self, exc: BaseException | None = None) -> None: self._exc = exc diff --git a/tests/test_stream_idle_timeout.py b/tests/test_stream_idle_timeout.py index 2f247d08..4f347496 100644 --- a/tests/test_stream_idle_timeout.py +++ b/tests/test_stream_idle_timeout.py @@ -88,11 +88,7 @@ def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]: base_url=base_url, ) guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout) - return _call(guarded) - - -def _call(model: Model) -> AsyncIterator[Any]: - return model.stream_response( + return guarded.stream_response( None, "go", ModelSettings(), @@ -124,12 +120,14 @@ async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str request_log.register_sink(logged.append) started = time.monotonic() try: - with pytest.raises(TimeoutError, match="stream_idle_timeout"): + with pytest.raises(TimeoutError, match="stream_idle_timeout") as raised: await _drain(stalling_gateway, idle_timeout=1) finally: request_log.unregister_sink(logged.append) assert time.monotonic() - started < _STALL_SECONDS + # The attempt's reply survives for the provider ban. + assert request_log.failed_reply(raised.value) is not None # 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" @@ -167,75 +165,6 @@ async def test_first_event_and_whole_stream_are_bounded() -> None: [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 (