feat(router): add stream_idle_timeout to catch stalled providers mid-stream

Extends the ttft_timeout feature with stream_idle_timeout: a per-chunk
inter-token deadline that fires litellm.Timeout when a provider accepts
a connection, sends some tokens, then goes silent. Both parameters are
independent; either or both can be set at router or per-deployment level.
This commit is contained in:
Nathan Price 2026-06-13 10:35:19 -05:00
parent 806db5663d
commit 9e7c1021f0
2 changed files with 271 additions and 103 deletions

View file

@ -275,6 +275,7 @@ class Router:
timeout: Optional[float] = None,
stream_timeout: Optional[float] = None,
ttft_timeout: Optional[float] = None,
stream_idle_timeout: Optional[float] = None,
default_litellm_params: Optional[
dict
] = None, # default params for Router.chat.completion.create
@ -433,9 +434,9 @@ class Router:
) # names of models under litellm_params. ex. azure/chatgpt-v-2
self.deployment_latency_map = {}
### CACHING ###
cache_type: Literal[
"local", "redis", "redis-semantic", "s3", "disk"
] = "local" # default to an in-memory cache
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = (
"local" # default to an in-memory cache
)
redis_cache = None
cache_config: Dict[str, Any] = {}
@ -483,9 +484,9 @@ class Router:
self.default_max_parallel_requests = default_max_parallel_requests
self.provider_default_deployment_ids: List[str] = []
self.pattern_router = PatternMatchRouter()
self.team_pattern_routers: Dict[
str, PatternMatchRouter
] = {} # {"TEAM_ID": PatternMatchRouter}
self.team_pattern_routers: Dict[str, PatternMatchRouter] = (
{}
) # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
@ -567,12 +568,13 @@ class Router:
self.timeout = timeout or litellm.request_timeout
self.stream_timeout = stream_timeout
self.ttft_timeout = ttft_timeout
self.stream_idle_timeout = stream_idle_timeout
self.retry_after = retry_after
self.routing_strategy = self._normalize_strategy(routing_strategy)
self._routing_groups_input: Optional[
List[Union[RoutingGroup, dict]]
] = routing_groups
self._routing_groups_input: Optional[List[Union[RoutingGroup, dict]]] = (
routing_groups
)
## SETTING FALLBACKS ##
### validate if it's set + in correct format
@ -699,12 +701,12 @@ class Router:
)
)
self.model_group_retry_policy: Optional[
Dict[str, RetryPolicy]
] = model_group_retry_policy
self.model_group_affinity_config: Optional[
Dict[str, List[str]]
] = model_group_affinity_config
self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
model_group_retry_policy
)
self.model_group_affinity_config: Optional[Dict[str, List[str]]] = (
model_group_affinity_config
)
self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None
if allowed_fails_policy is not None:
@ -2618,20 +2620,20 @@ class Router:
# _ageneric_api_call_with_fallbacks_helper.
# original_generic_function is preserved by the caller so
# the helper knows what underlying API to invoke per attempt.
initial_kwargs[
"original_function"
] = self._ageneric_api_call_with_fallbacks_helper
initial_kwargs["original_function"] = (
self._ageneric_api_call_with_fallbacks_helper
)
if e.is_pre_first_chunk or not e.generated_content:
# No content generated before the error — retry with the
# original input. Adding a continuation prompt would
# waste tokens and confuse the model.
pass
else:
initial_kwargs[
"input"
] = Router._build_responses_continuation_input(
initial_kwargs.get("input"),
e.generated_content,
initial_kwargs["input"] = (
Router._build_responses_continuation_input(
initial_kwargs.get("input"),
e.generated_content,
)
)
# The Responses-API path stores observability metadata
# under "litellm_metadata" (not the default "metadata") —
@ -2858,49 +2860,72 @@ class Router:
self,
response: CustomStreamWrapper,
messages: List[Dict[str, str]],
ttft_timeout: float,
ttft_timeout: Optional[float],
stream_idle_timeout: Optional[float] = None,
) -> ModelResponse:
from litellm.main import stream_chunk_builder
chunks: List = []
aiter = response.__aiter__()
loop = asyncio.get_running_loop()
deadline = loop.time() + ttft_timeout
first_token_received = False
while not first_token_received:
remaining = deadline - loop.time()
if remaining <= 0:
verbose_router_logger.warning(
f"ttft_timeout={ttft_timeout}s exceeded for model={response.model}: "
"provider accepted connection but sent no tokens"
)
raise litellm.Timeout(
message=f"Router ttft_timeout={ttft_timeout}s exceeded: provider accepted connection but sent no tokens",
model=response.model or "",
llm_provider=response.custom_llm_provider or "",
)
try:
chunk = await asyncio.wait_for(aiter.__anext__(), timeout=remaining)
except asyncio.TimeoutError:
verbose_router_logger.warning(
f"ttft_timeout={ttft_timeout}s exceeded for model={response.model}: "
"provider accepted connection but sent no tokens"
)
raise litellm.Timeout(
message=f"Router ttft_timeout={ttft_timeout}s exceeded: provider accepted connection but sent no tokens",
model=response.model or "",
llm_provider=response.custom_llm_provider or "",
)
except StopAsyncIteration:
break
chunks.append(chunk)
delta = chunk.choices[0].delta if chunk.choices else None
if delta and (delta.content or delta.tool_calls):
first_token_received = True
if ttft_timeout is not None:
loop = asyncio.get_running_loop()
deadline = loop.time() + ttft_timeout
first_token_received = False
async for chunk in aiter:
chunks.append(chunk)
while not first_token_received:
remaining = deadline - loop.time()
if remaining <= 0:
verbose_router_logger.warning(
f"ttft_timeout={ttft_timeout}s exceeded for model={response.model}: "
"provider accepted connection but sent no tokens"
)
raise litellm.Timeout(
message=f"Router ttft_timeout={ttft_timeout}s exceeded: provider accepted connection but sent no tokens",
model=response.model or "",
llm_provider=response.custom_llm_provider or "",
)
try:
chunk = await asyncio.wait_for(aiter.__anext__(), timeout=remaining)
except asyncio.TimeoutError:
verbose_router_logger.warning(
f"ttft_timeout={ttft_timeout}s exceeded for model={response.model}: "
"provider accepted connection but sent no tokens"
)
raise litellm.Timeout(
message=f"Router ttft_timeout={ttft_timeout}s exceeded: provider accepted connection but sent no tokens",
model=response.model or "",
llm_provider=response.custom_llm_provider or "",
)
except StopAsyncIteration:
break
chunks.append(chunk)
delta = chunk.choices[0].delta if chunk.choices else None
if delta and (delta.content or delta.tool_calls):
first_token_received = True
if stream_idle_timeout is not None:
while True:
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)
else:
async for chunk in aiter:
chunks.append(chunk)
result = stream_chunk_builder(chunks, messages=messages)
if result is None:
@ -2914,7 +2939,10 @@ class Router:
async def _acompletion( # noqa: PLR0915
self, model: str, messages: List[Dict[str, str]], **kwargs
) -> Union[ModelResponse, CustomStreamWrapper,]:
) -> Union[
ModelResponse,
CustomStreamWrapper,
]:
"""
- Get an available deployment
- call it with a semaphore over the call
@ -2996,9 +3024,12 @@ class Router:
input_kwargs.pop("silent_model", None)
_ttft_timeout = self._get_ttft_timeout(kwargs=kwargs, data=litellm_params)
_forced_stream_for_ttft = (
_ttft_timeout is not None and not input_kwargs.get("stream", False)
_stream_idle_timeout = self._get_stream_idle_timeout(
kwargs=kwargs, data=litellm_params
)
_forced_stream_for_ttft = (
_ttft_timeout is not None or _stream_idle_timeout is not None
) and not input_kwargs.get("stream", False)
if _forced_stream_for_ttft:
input_kwargs["stream"] = True
@ -3060,11 +3091,12 @@ class Router:
)
if isinstance(response, CustomStreamWrapper):
if _forced_stream_for_ttft and _ttft_timeout is not None:
if _forced_stream_for_ttft:
reconstructed = await self._collect_stream_with_ttft_timeout(
response=response,
messages=messages,
ttft_timeout=_ttft_timeout,
stream_idle_timeout=_stream_idle_timeout,
)
if self._should_raise_content_policy_error(
model=model, response=reconstructed, kwargs=kwargs
@ -3397,6 +3429,17 @@ class Router:
return source
return None
def _get_stream_idle_timeout(self, kwargs: dict, data: dict) -> Optional[float]:
for source in (
kwargs.get("stream_idle_timeout"),
data.get("stream_idle_timeout"),
self.stream_idle_timeout,
self.default_litellm_params.get("stream_idle_timeout"),
):
if source is not None:
return source
return None
def _get_timeout(self, kwargs: dict, data: dict) -> Optional[Union[float, int]]:
"""Helper to get timeout from kwargs or deployment params"""
timeout: Optional[Union[float, int]] = None
@ -5345,9 +5388,9 @@ class Router:
healthy_deployments=healthy_deployments, responses=responses
)
returned_response = cast(OpenAIFileObject, responses[0])
returned_response._hidden_params[
"model_file_id_mapping"
] = model_file_id_mapping
returned_response._hidden_params["model_file_id_mapping"] = (
model_file_id_mapping
)
return returned_response
except Exception as e:
verbose_router_logger.exception(
@ -6666,11 +6709,11 @@ class Router:
if isinstance(e, litellm.ContextWindowExceededError):
if context_window_fallbacks is not None:
context_window_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
context_window_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
)
)
if context_window_fallback_model_group is None:
raise original_exception
@ -6702,11 +6745,11 @@ class Router:
e.message += "\n{}".format(error_message)
elif isinstance(e, litellm.ContentPolicyViolationError):
if content_policy_fallbacks is not None:
content_policy_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
content_policy_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
)
)
if content_policy_fallback_model_group is None:
raise original_exception
@ -6928,9 +6971,9 @@ class Router:
)
## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking
_metadata["attempted_retries"] = 0
_metadata[
"max_retries"
] = num_retries # Updated after overrides in exception handler
_metadata["max_retries"] = (
num_retries # Updated after overrides in exception handler
)
try:
self._handle_mock_testing_rate_limit_error(
model_group=model_group, kwargs=kwargs
@ -8129,26 +8172,26 @@ class Router:
"""
from litellm.router_strategy.auto_router.auto_router import AutoRouter
auto_router_config_path: Optional[
str
] = deployment.litellm_params.auto_router_config_path
auto_router_config_path: Optional[str] = (
deployment.litellm_params.auto_router_config_path
)
auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config
if auto_router_config_path is None and auto_router_config is None:
raise ValueError(
"auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params"
)
default_model: Optional[
str
] = deployment.litellm_params.auto_router_default_model
default_model: Optional[str] = (
deployment.litellm_params.auto_router_default_model
)
if default_model is None:
raise ValueError(
"auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params"
)
embedding_model: Optional[
str
] = deployment.litellm_params.auto_router_embedding_model
embedding_model: Optional[str] = (
deployment.litellm_params.auto_router_embedding_model
)
if embedding_model is None:
raise ValueError(
"auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params"
@ -8191,13 +8234,13 @@ class Router:
ComplexityRouter,
)
complexity_router_config: Optional[
dict
] = deployment.litellm_params.complexity_router_config
complexity_router_config: Optional[dict] = (
deployment.litellm_params.complexity_router_config
)
default_model: Optional[
str
] = deployment.litellm_params.complexity_router_default_model
default_model: Optional[str] = (
deployment.litellm_params.complexity_router_default_model
)
# If no default model specified, try to get from config tiers
if default_model is None and complexity_router_config:
@ -8384,13 +8427,13 @@ class Router:
QualityRouter,
)
quality_router_config: Optional[
dict
] = deployment.litellm_params.quality_router_config
quality_router_config: Optional[dict] = (
deployment.litellm_params.quality_router_config
)
default_model: Optional[
str
] = deployment.litellm_params.quality_router_default_model
default_model: Optional[str] = (
deployment.litellm_params.quality_router_default_model
)
if default_model is None and quality_router_config:
default_model = quality_router_config.get("default_model")
@ -9161,9 +9204,9 @@ class Router:
# Add custom_llm_provider
if deployment.litellm_params.custom_llm_provider:
credentials[
"custom_llm_provider"
] = deployment.litellm_params.custom_llm_provider
credentials["custom_llm_provider"] = (
deployment.litellm_params.custom_llm_provider
)
elif "/" in deployment.litellm_params.model:
# Extract provider from "provider/model" format
credentials["custom_llm_provider"] = deployment.litellm_params.model.split(

View file

@ -4979,3 +4979,128 @@ async def test_router_ttft_timeout_acompletion_intercept():
assert result is reconstructed
assert mock_collect.called
# ---------------------------------------------------------------------------
# stream_idle_timeout tests
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_router_stream_idle_timeout_raises_on_stalled_provider():
"""After first token arrives, stream_idle_timeout fires if no subsequent chunk arrives in time."""
import asyncio
from unittest.mock import patch
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
}
],
stream_idle_timeout=0.05,
)
async def _stalled_after_first():
yield _make_chunk("Hello")
await asyncio.sleep(10) # stalls; will be killed by stream_idle_timeout
yield _make_chunk("", finish_reason="stop")
fake_stream = MagicMock()
fake_stream.model = "gpt-4o"
fake_stream.custom_llm_provider = "openai"
gen = _stalled_after_first()
fake_stream.__aiter__ = lambda self: gen
with pytest.raises(litellm.Timeout, match="stream_idle_timeout"):
await router._collect_stream_with_ttft_timeout(
response=fake_stream,
messages=[{"role": "user", "content": "hi"}],
ttft_timeout=None,
stream_idle_timeout=0.05,
)
@pytest.mark.asyncio
async def test_router_stream_idle_timeout_completes_when_not_stalled():
"""stream_idle_timeout does not fire when chunks arrive within the timeout."""
from unittest.mock import patch
from litellm import ModelResponse
import litellm
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 = MagicMock()
fake_stream.model = "gpt-4o"
fake_stream.custom_llm_provider = "openai"
fake_stream.__aiter__ = lambda self: _async_chunks(*chunks)
reconstructed = MagicMock(spec=ModelResponse)
with 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_ttft_and_idle_timeout_both_active():
"""When both ttft_timeout and stream_idle_timeout are set, both phases are enforced."""
import asyncio
from unittest.mock import patch
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
}
],
ttft_timeout=5.0,
stream_idle_timeout=0.05,
)
async def _stalled_after_first():
yield _make_chunk("Hello")
await asyncio.sleep(10)
yield _make_chunk("", finish_reason="stop")
fake_stream = MagicMock()
fake_stream.model = "gpt-4o"
fake_stream.custom_llm_provider = "openai"
gen = _stalled_after_first()
fake_stream.__aiter__ = lambda self: gen
with pytest.raises(litellm.Timeout, match="stream_idle_timeout"):
await router._collect_stream_with_ttft_timeout(
response=fake_stream,
messages=[{"role": "user", "content": "hi"}],
ttft_timeout=5.0,
stream_idle_timeout=0.05,
)