fix(llm): abandon a model stream that stops producing events (#978)

Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

strix/config/models.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import inspect
 import logging
 import os
 import time
+from collections.abc import AsyncGenerator
 from typing import TYPE_CHECKING, Any, cast
 
 from agents import (
@@ -252,11 +254,23 @@ class _TurnGuardModel(Model):
     Tool-call volume: a degenerate response can queue hundreds of calls that
     the run loop then honours one by one. Only the first
     ``LLM_MAX_TOOL_CALLS_PER_TURN`` calls of a response are kept.
+
+    Stalled streams: a turn that emits a few tokens and then goes silent is
+    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.
     """
 
-    def __init__(self, inner: Model, *, max_tool_calls_per_turn: int = 0) -> None:
+    def __init__(
+        self,
+        inner: Model,
+        *,
+        max_tool_calls_per_turn: int = 0,
+        stream_idle_timeout: float = 0.0,
+    ) -> None:
         self._inner = inner
         self._max_tool_calls_per_turn = max_tool_calls_per_turn
+        self._stream_idle_timeout = stream_idle_timeout
 
     def _limiter(self) -> TurnToolCallLimiter:
         return TurnToolCallLimiter(self._max_tool_calls_per_turn)
@@ -337,13 +351,41 @@ class _TurnGuardModel(Model):
             conversation_id=conversation_id,
             prompt=prompt,
         )
-        async for event in stream:
+        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
         self._log_dropped(limiter)
 
 
+async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None:
+    if isinstance(stream, AsyncGenerator):
+        with contextlib.suppress(Exception):
+            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
+
+    iterator = stream.__aiter__()
+    while True:
+        try:
+            event = await asyncio.wait_for(iterator.__anext__(), timeout)
+        except StopAsyncIteration:
+            return
+        except TimeoutError:
+            await _aclose(stream)
+            message = f"model stream produced no event for {timeout:.0f}s"
+            logger.warning("%s; abandoning the turn", message)
+            raise TimeoutError(message) from None
+        yield event
+
+
 def _guard_event(
     event: TResponseStreamEvent, rewriter: TurnCallIdRewriter, limiter: TurnToolCallLimiter
 ) -> TResponseStreamEvent | None:
@@ -429,6 +471,7 @@ class StrixProvider(MultiProvider):
     def get_model(self, model_name: str | None) -> Model:
         llm = load_settings().llm
         slug = codex.subscription_model(model_name)
+        idle_timeout = float(llm.stream_idle_timeout)
         if slug:
             # The ChatGPT subscription backend is always streamed; it has no
             # non-streaming mode to fall back to, so LLM_DISABLE_STREAMING
@@ -442,7 +485,15 @@ class StrixProvider(MultiProvider):
             model = super().get_model(model_name)
             if llm.disable_streaming:
                 model = _NonStreamingModel(model)
-        return _TurnGuardModel(model, max_tool_calls_per_turn=llm.max_tool_calls_per_turn)
+                # 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
+        return _TurnGuardModel(
+            model,
+            max_tool_calls_per_turn=llm.max_tool_calls_per_turn,
+            stream_idle_timeout=idle_timeout,
+        )
 
 
 DEFAULT_MODEL_RETRY = ModelRetrySettings(

strix/config/settings.py

@@ -57,6 +57,7 @@ class LlmSettings(BaseSettings):
         alias="LLM_DISABLE_STREAMING",
     )
     timeout: int = Field(default=300, alias="LLM_TIMEOUT")
+    stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT")
     max_tool_calls_per_turn: int = Field(
         default=32,
         ge=0,