scriptoriumadmin/strix
1
0
Fork 0
mirror of https://github.com/usestrix/strix.git synced 2026-10-14 04:08:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

refactor(llm): ban timed-out providers in the existing retry path

Browse source
This commit is contained in:
Ian 2026-10-07 11:47:18 -04:00
parent 58d25f447d
commit 4b8212a447
5 changed files with 61 additions and 143 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
strix/config/models.py
View file

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

7
strix/core/execution.py
View file

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

15
strix/llm/request_log.py
View file

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

19
tests/test_execution_transient_retry.py
View file

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

79
tests/test_stream_idle_timeout.py
View file

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