fix(models): split first-event and idle stream timeouts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

strix/config/models.py

@@ -61,6 +61,10 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class ModelStreamTimeoutError(TimeoutError):
+    """Raised when a model stream exceeds its event timeout."""
+
+
 def request_timeout_extra_args(timeout_s: float | None) -> dict[str, float] | None:
     """Per-request model timeout; a plain float so ``ModelSettings.to_json_dict()`` stays serializable."""  # noqa: E501
     if not timeout_s or timeout_s <= 0:
@@ -267,10 +271,12 @@ class _TurnGuardModel(Model):
         *,
         max_tool_calls_per_turn: int = 0,
         stream_idle_timeout: float = 0.0,
+        stream_first_event_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
+        self._stream_first_event_timeout = stream_first_event_timeout
 
     def _limiter(self) -> TurnToolCallLimiter:
         return TurnToolCallLimiter(self._max_tool_calls_per_turn)
@@ -351,7 +357,11 @@ class _TurnGuardModel(Model):
             conversation_id=conversation_id,
             prompt=prompt,
         )
-        async for event in _with_idle_timeout(stream, self._stream_idle_timeout):
+        async for event in _with_idle_timeout(
+            stream,
+            self._stream_idle_timeout,
+            self._stream_first_event_timeout,
+        ):
             guarded = _guard_event(event, rewriter, limiter)
             if guarded is not None:
                 yield guarded
@@ -365,24 +375,39 @@ async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None:
 
 
 async def _with_idle_timeout(
-    stream: AsyncIterator[TResponseStreamEvent], timeout: float
+    stream: AsyncIterator[TResponseStreamEvent],
+    timeout: float,
+    first_event_timeout: float = 0.0,
 ) -> AsyncIterator[TResponseStreamEvent]:
-    if timeout <= 0:
+    effective_first_event_timeout = first_event_timeout if first_event_timeout > 0 else timeout
+    if timeout <= 0 and effective_first_event_timeout <= 0:
         async for event in stream:
             yield event
         return
 
     iterator = stream.__aiter__()
+    yielded_event = False
     while True:
+        event_timeout = timeout if yielded_event else effective_first_event_timeout
         try:
-            event = await asyncio.wait_for(iterator.__anext__(), timeout)
+            if event_timeout > 0:
+                event = await asyncio.wait_for(iterator.__anext__(), event_timeout)
+            else:
+                event = await iterator.__anext__()
         except StopAsyncIteration:
             return
         except TimeoutError:
             await _aclose(stream)
-            message = f"model stream produced no event for {timeout:.0f}s"
+            if yielded_event:
+                message = f"model stream produced no event for {timeout:.0f}s"
+            else:
+                message = (
+                    f"model stream produced no first event within "
+                    f"{effective_first_event_timeout:.0f}s"
+                )
             logger.warning("%s; abandoning the turn", message)
-            raise TimeoutError(message) from None
+            raise ModelStreamTimeoutError(message) from None
+        yielded_event = True
         yield event
 
 
@@ -472,6 +497,7 @@ 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)
         if slug:
             # The ChatGPT subscription backend is always streamed; it has no
             # non-streaming mode to fall back to, so LLM_DISABLE_STREAMING
@@ -489,10 +515,12 @@ class StrixProvider(MultiProvider):
                 # is done, so an idle gap is meaningless here; the request
                 # timeout bounds it instead.
                 idle_timeout = 0.0
+                first_event_timeout = 0.0
         return _TurnGuardModel(
             model,
             max_tool_calls_per_turn=llm.max_tool_calls_per_turn,
             stream_idle_timeout=idle_timeout,
+            stream_first_event_timeout=first_event_timeout,
         )
 
 

strix/config/settings.py

@@ -58,6 +58,12 @@ class LlmSettings(BaseSettings):
     )
     timeout: int = Field(default=300, alias="LLM_TIMEOUT")
     stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT")
+    # Time allowed for the first stream event (0 = use stream_idle_timeout).
+    stream_first_event_timeout: int = Field(
+        default=0,
+        ge=0,
+        alias="LLM_STREAM_FIRST_EVENT_TIMEOUT",
+    )
     max_tool_calls_per_turn: int = Field(
         default=32,
         ge=0,

strix/core/execution.py

@@ -20,6 +20,7 @@ from openai import (
 )
 
 from strix.config import codex
+from strix.config.models import ModelStreamTimeoutError
 from strix.core.hooks import (
     BudgetExceededError,
     BudgetPausedError,
@@ -120,6 +121,7 @@ async def _compact_session(
 
 
 _MAX_TRANSIENT_MODEL_RETRIES = 5
+_MAX_STREAM_TIMEOUT_RETRIES = 2
 _TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0
 _TRANSIENT_MODEL_RETRY_MAX_DELAY_S = 90.0
 
@@ -657,6 +659,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
     image_strips = 0
     compactions = 0
     model_retries = 0
+    stream_timeout_retries = 0
     while True:
         stream: Any = None
         pre_run_items: list[Any] = []
@@ -771,15 +774,27 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                     )
                     input_data = []
                     continue
-            if model_retries < _MAX_TRANSIENT_MODEL_RETRIES and _is_transient_model_error(exc):
-                model_retries += 1
-                delay = _transient_model_retry_delay(model_retries)
+            is_stream_timeout = isinstance(exc, ModelStreamTimeoutError)
+            if is_stream_timeout:
+                retry_count = stream_timeout_retries
+                retry_limit = _MAX_STREAM_TIMEOUT_RETRIES
+            else:
+                retry_count = model_retries
+                retry_limit = _MAX_TRANSIENT_MODEL_RETRIES
+                stream_timeout_retries = 0
+            if retry_count < retry_limit and (is_stream_timeout or _is_transient_model_error(exc)):
+                retry_count += 1
+                if is_stream_timeout:
+                    stream_timeout_retries = retry_count
+                else:
+                    model_retries = retry_count
+                delay = _transient_model_retry_delay(retry_count)
                 logger.warning(
                     "transient model/provider error for %s; replaying turn "
                     "(attempt %d/%d, backoff %.1fs): %r",
                     agent_id,
-                    model_retries,
-                    _MAX_TRANSIENT_MODEL_RETRIES,
+                    retry_count,
+                    retry_limit,
                     delay,
                     exc,
                 )

tests/test_execution_transient_retry.py

@@ -16,6 +16,7 @@ from openai import (
 )
 
 from strix.config import codex
+from strix.config.models import ModelStreamTimeoutError
 from strix.core import execution
 from strix.core.agents import AgentCoordinator
 
@@ -162,6 +163,16 @@ async def test_run_cycle_gives_up_after_max_retries(
         await _run_once(monkeypatch, streams)
 
 
+@pytest.mark.asyncio
+async def test_run_cycle_gives_up_after_max_stream_timeout_retries(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    timeout = ModelStreamTimeoutError("model stream produced no event for 1s")
+    streams = [_FakeStream(exc=timeout) for _ in range(execution._MAX_STREAM_TIMEOUT_RETRIES + 1)]
+    with pytest.raises(ModelStreamTimeoutError):
+        await _run_once(monkeypatch, streams)
+
+
 @pytest.mark.asyncio
 async def test_run_cycle_does_not_retry_permanent_error(
     monkeypatch: pytest.MonkeyPatch,

tests/test_stream_idle_timeout.py

@@ -23,7 +23,12 @@ 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 (
+    ModelStreamTimeoutError,
+    StrixProvider,
+    _TurnGuardModel,
+    _with_idle_timeout,
+)
 
 
 if TYPE_CHECKING:
@@ -64,6 +69,24 @@ class _StallingHandler(BaseHTTPRequestHandler):
         self.stop.wait(_STALL_SECONDS)
 
 
+class _FirstEventStallingHandler(BaseHTTPRequestHandler):
+    """Sends stream headers, then waits before sending the first event."""
+
+    stop = threading.Event()
+
+    def log_message(self, *args: Any) -> None:
+        pass
+
+    def do_POST(self) -> None:
+        length = int(self.headers.get("Content-Length", 0))
+        self.rfile.read(length)
+        self.send_response(200)
+        self.send_header("Content-Type", "text/event-stream")
+        self.end_headers()
+        self.wfile.flush()
+        self.stop.wait(_STALL_SECONDS)
+
+
 @pytest.fixture
 def stalling_gateway() -> Iterator[str]:
     _StallingHandler.stop.clear()
@@ -78,10 +101,33 @@ def stalling_gateway() -> Iterator[str]:
         server.server_close()
 
 
-def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
+@pytest.fixture
+def first_event_stalling_gateway() -> Iterator[str]:
+    _FirstEventStallingHandler.stop.clear()
+    server = HTTPServer(("127.0.0.1", 0), _FirstEventStallingHandler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield f"http://127.0.0.1:{server.server_address[1]}/v1"
+    finally:
+        _FirstEventStallingHandler.stop.set()
+        server.shutdown()
+        server.server_close()
+
+
+def _stream(
+    base_url: str,
+    *,
+    idle_timeout: float,
+    first_event_timeout: float = 0.0,
+) -> 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)
-    guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout)
+    guarded = _TurnGuardModel(
+        inner,
+        stream_idle_timeout=idle_timeout,
+        stream_first_event_timeout=first_event_timeout,
+    )
     return guarded.stream_response(
         None,
         "go",
@@ -96,8 +142,20 @@ def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
     )
 
 
-async def _drain(base_url: str, *, idle_timeout: float) -> list[Any]:
-    return [event async for event in _stream(base_url, idle_timeout=idle_timeout)]
+async def _drain(
+    base_url: str,
+    *,
+    idle_timeout: float,
+    first_event_timeout: float = 0.0,
+) -> list[Any]:
+    return [
+        event
+        async for event in _stream(
+            base_url,
+            idle_timeout=idle_timeout,
+            first_event_timeout=first_event_timeout,
+        )
+    ]
 
 
 @pytest.mark.asyncio
@@ -117,6 +175,33 @@ async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str
     assert time.monotonic() - started < _STALL_SECONDS
 
 
+@pytest.mark.asyncio
+async def test_first_event_timeout_abandons_silent_stream(
+    first_event_stalling_gateway: str,
+) -> None:
+    started = time.monotonic()
+    with pytest.raises(ModelStreamTimeoutError, match="no first event within 1s"):
+        await _drain(
+            first_event_stalling_gateway,
+            idle_timeout=10,
+            first_event_timeout=1,
+        )
+
+    assert time.monotonic() - started < _STALL_SECONDS
+
+
+@pytest.mark.asyncio
+async def test_first_event_timeout_does_not_replace_idle_timeout(
+    stalling_gateway: str,
+) -> None:
+    with pytest.raises(ModelStreamTimeoutError, match="produced no event for 1s"):
+        await _drain(
+            stalling_gateway,
+            idle_timeout=1,
+            first_event_timeout=10,
+        )
+
+
 @pytest.mark.asyncio
 async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
     async def _live() -> AsyncIterator[Any]:
@@ -131,7 +216,12 @@ async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
 
 @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",
+    ):
         monkeypatch.delenv(key, raising=False)
     monkeypatch.setattr(loader, "_cached", None)
     monkeypatch.setattr(loader, "_override", None)
@@ -156,6 +246,20 @@ def test_idle_timeout_is_configurable(
     model = StrixProvider().get_model("openai/gpt-4o-mini")
     assert isinstance(model, _TurnGuardModel)
     assert model._stream_idle_timeout == 45
+    assert model._stream_first_event_timeout == 0
+
+
+def test_first_event_timeout_is_configurable(
+    monkeypatch: pytest.MonkeyPatch, _reset_settings: None
+) -> 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", "12")
+    load_settings()
+
+    model = StrixProvider().get_model("openai/gpt-4o-mini")
+    assert isinstance(model, _TurnGuardModel)
+    assert model._stream_first_event_timeout == 12
 
 
 def test_idle_timeout_is_off_without_streaming(
@@ -171,3 +275,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_first_event_timeout == 0