Reduce stream idle timeout overhead

litellm/router.py

@@ -3004,24 +3004,46 @@ class Router:
                         first_token_received = True
 
             if stream_idle_timeout is not None:
-                while True:
+                idle_timeout = stream_idle_timeout
+                stream_idle_timed_out = False
+                current_task = asyncio.current_task()
+                idle_timeout_handle: asyncio.TimerHandle | None = None
+
+                def _cancel_on_stream_idle_timeout() -> None:
+                    nonlocal stream_idle_timed_out
+                    stream_idle_timed_out = True
+                    if current_task is not None:
+                        current_task.cancel()
+
+                def _reset_stream_idle_timer() -> None:
+                    nonlocal idle_timeout_handle
+                    if idle_timeout_handle is not None:
+                        idle_timeout_handle.cancel()
+                    idle_timeout_handle = loop.call_later(
+                        idle_timeout, _cancel_on_stream_idle_timeout
+                    )
+
+                _reset_stream_idle_timer()
+                try:
                     try:
-                        chunk = await asyncio.wait_for(
-                            aiter.__anext__(), timeout=stream_idle_timeout
-                        )
-                    except asyncio.TimeoutError:
-                        verbose_router_logger.warning(
-                            f"stream_idle_timeout={stream_idle_timeout}s exceeded for model={response.model}: "
-                            "provider stalled mid-stream"
-                        )
-                        raise litellm.Timeout(
-                            message=f"Router stream_idle_timeout={stream_idle_timeout}s exceeded: provider stalled mid-stream",
-                            model=response.model or "",
-                            llm_provider=response.custom_llm_provider or "",
-                        )
-                    except StopAsyncIteration:
-                        break
-                    chunks.append(chunk)
+                        async for chunk in aiter:
+                            chunks.append(chunk)
+                            _reset_stream_idle_timer()
+                    except asyncio.CancelledError:
+                        if stream_idle_timed_out:
+                            verbose_router_logger.warning(
+                                f"stream_idle_timeout={stream_idle_timeout}s exceeded for model={response.model}: "
+                                "provider stalled mid-stream"
+                            )
+                            raise litellm.Timeout(
+                                message=f"Router stream_idle_timeout={stream_idle_timeout}s exceeded: provider stalled mid-stream",
+                                model=response.model or "",
+                                llm_provider=response.custom_llm_provider or "",
+                            )
+                        raise
+                finally:
+                    if idle_timeout_handle is not None:
+                        idle_timeout_handle.cancel()
             else:
                 async for chunk in aiter:
                     chunks.append(chunk)

tests/test_litellm/test_router.py

@@ -5317,6 +5317,47 @@ async def test_router_stream_idle_timeout_completes_when_not_stalled():
     assert result is reconstructed
 
 
+@pytest.mark.asyncio
+async def test_router_stream_idle_timeout_avoids_per_chunk_wait_for():
+    """The idle timer should use one rescheduled watchdog, not wrap every streamed
+    chunk in asyncio.wait_for."""
+    from unittest.mock import patch
+
+    from litellm import ModelResponse
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "test-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
+            }
+        ],
+        stream_idle_timeout=5.0,
+    )
+
+    chunks = [
+        _make_chunk("Hello"),
+        _make_chunk(" world"),
+        _make_chunk("", finish_reason="stop"),
+    ]
+
+    fake_stream = _fake_stream(lambda: _async_chunks(*chunks))
+    reconstructed = MagicMock(spec=ModelResponse)
+
+    with (
+        patch("asyncio.wait_for", side_effect=AssertionError("per-chunk wait_for")),
+        patch("litellm.main.stream_chunk_builder", return_value=reconstructed),
+    ):
+        result = await router._collect_stream_with_ttft_timeout(
+            response=fake_stream,
+            messages=[{"role": "user", "content": "hi"}],
+            ttft_timeout=None,
+            stream_idle_timeout=5.0,
+        )
+
+    assert result is reconstructed
+
+
 @pytest.mark.asyncio
 async def test_router_stream_idle_timeout_does_not_fire_before_first_token():
     """When stream_idle_timeout is set without ttft_timeout, a slow first token must NOT be