fix(router): wrap aresponses streaming iterator for mid-stream fallbacks - #28214
fix(router): wrap aresponses streaming iterator for mid-stream fallbacks#28214cwang-otto wants to merge 1 commit into
Conversation
MidStreamFallbackError raised during aresponses streaming was bypassing the Router's fallback chain because only the chat completions path wrapped its CustomStreamWrapper. The aresponses dispatch goes through _ageneric_api_call_with_fallbacks which returns the streaming iterator unwrapped, so any MidStreamFallbackError propagated past the Router and the configured cross-provider fallback (e.g. anthropic -> vertex_ai) never fired. Add _aresponses_streaming_iterator (parity with _acompletion_streaming_iterator) plus a thin _aresponses_with_streaming_fallbacks dispatch wrapper. On MidStreamFallbackError, re-enter the chain via async_function_with_fallbacks_common_utils with original_function set to the per-attempt helper and original_generic_function preserved, so the helper invokes litellm.aresponses on each fallback attempt. Scope: pre-first-chunk retry only. Responses-API input shape differs from chat completions, so partial-content continuation is intentionally out of scope.
|
Re-opening against shin_agent_oss_staging_05_19_2026 per the Guard main branch policy (see #28201 for the pattern). |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR fixes mid-stream fallback handling for
Confidence Score: 4/5The happy path (no MidStreamFallbackError) is unchanged; the new wrapper is only activated on streaming Responses API calls, making the blast radius of any regression small. The change is narrowly scoped to the Responses API streaming path and does not touch the chat-completions or synchronous fallback paths. Two maintainability concerns are present: FallbackResponsesStreamWrapper leaves several BaseResponsesAPIStreamingIterator attributes uninitialised, and the shallow kwargs.copy() lets litellm_metadata be mutated by the primary attempt so fallback attempts carry a stale model-group in that dict. litellm/router.py — specifically the FallbackResponsesStreamWrapper.init and the fallback_kwargs = kwargs.copy() snapshot in _aresponses_with_streaming_fallbacks.
|
| Filename | Overview |
|---|---|
| litellm/router.py | Adds _aresponses_streaming_iterator and _aresponses_with_streaming_fallbacks to bring mid-stream fallback parity to the Responses API path; the FallbackResponsesStreamWrapper bypasses super().__init__() leaving several base-class attributes uninitialised (safe today but fragile), and the shallow kwargs.copy() snapshot can propagate stale litellm_metadata to fallback attempts. |
| tests/test_litellm/test_router.py | Adds a well-structured mock-only unit test that verifies chunk forwarding, MidStreamFallbackError interception, correct kwargs passed to the fallback utils, and isinstance preservation; no real network calls. |
Reviews (1): Last reviewed commit: "fix(router): wrap aresponses streaming i..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| Subclasses BaseResponsesAPIStreamingIterator only for isinstance | ||
| compatibility (proxy + interactions code paths check the type). | ||
| Bypasses the parent constructor and delegates iteration to an | ||
| async generator. | ||
| """ | ||
|
|
||
| def __init__(self, async_generator: AsyncGenerator): | ||
| self._async_generator = async_generator | ||
| self.finished = False | ||
| # Preserve hidden params and identity attributes so response | ||
| # headers (model_id, api_base, additional_headers) still flow. | ||
| self._hidden_params = dict( | ||
| getattr(source_iterator, "_hidden_params", {}) or {} | ||
| ) | ||
| self.model = getattr(source_iterator, "model", "") | ||
| self.custom_llm_provider = getattr( | ||
| source_iterator, "custom_llm_provider", None | ||
| ) | ||
| self.logging_obj = getattr(source_iterator, "logging_obj", None) | ||
| self.litellm_metadata = getattr( | ||
| source_iterator, "litellm_metadata", None | ||
| ) | ||
| self.responses_api_provider_config = getattr( | ||
| source_iterator, "responses_api_provider_config", None | ||
| ) | ||
|
|
||
| def __aiter__(self): | ||
| return self | ||
|
|
||
| async def __anext__(self): | ||
| return await self._async_generator.__anext__() | ||
|
|
||
| async def aclose(self): | ||
| close = getattr(self._async_generator, "aclose", None) | ||
| if close is not None: |
There was a problem hiding this comment.
FallbackResponsesStreamWrapper intentionally bypasses super().__init__() but leaves many BaseResponsesAPIStreamingIterator instance attributes unset: completed_response, _stream_created_time, start_time, _failure_handled, response, request_data, call_type, _completed_response_cached, _completed_response_logged, _completed_response_cache_hit, and _persist_completed_response_before_logging.
Currently this is safe because FallbackResponsesStreamWrapper.__anext__ delegates entirely to the async generator (so base-class methods like _check_max_streaming_duration and _log_success are never called on the wrapper itself). However, any code path that receives a BaseResponsesAPIStreamingIterator and accesses instance.completed_response or similar (e.g., new post-stream logging hooks) will get AttributeError on this wrapper. Initializing these with safe sentinels (e.g., self.completed_response = None, self._failure_handled = False, self._stream_created_time = time.time()) would prevent that class of failure.
| fallback_kwargs = kwargs.copy() | ||
| fallback_kwargs["original_generic_function"] = original_function |
There was a problem hiding this comment.
fallback_kwargs = kwargs.copy() is a shallow copy. When _ageneric_api_call_with_fallbacks(**kwargs) runs, it calls _update_kwargs_before_fallbacks which does kwargs.setdefault("litellm_metadata", {}).update({"model_group": model, ...}). If litellm_metadata already exists in kwargs, the in-place .update() mutates the same dict object referenced by fallback_kwargs["litellm_metadata"]. As a result, by the time a MidStreamFallbackError triggers and stream_with_fallbacks calls _update_kwargs_before_fallbacks again (with the default metadata_variable_name="metadata", not "litellm_metadata"), fallback_kwargs["litellm_metadata"]["model_group"] still carries the primary deployment's model group rather than the fallback model group. This means the fallback attempt's litellm_metadata is stale, which can produce incorrect routing metadata in logs.
1 participant
Summary
MidStreamFallbackErrorraised mid-stream duringRouter.aresponses(stream=True)bypasses the Router's fallback chain, so configured cross-provider fallbacks (e.g.anthropic → vertex_ai) never fire when the primary provider's stream fails mid-flight.The chat completions path wraps its
CustomStreamWrappervia_acompletion_streaming_iterator, which catchesMidStreamFallbackErrorand re-entersasync_function_with_fallbacks_common_utils. The Responses API path doesn't have an equivalent: it dispatches through_ageneric_api_call_with_fallbacks→_ageneric_api_call_with_fallbacks_helper, which awaitslitellm.aresponses(**response_kwargs)and returns the streaming iterator unwrapped. AnyMidStreamFallbackErrorraised during iteration (e.g. by the underlyingCustomStreamWrapperwhenLiteLLMCompletionStreamingIteratorbridges through completion) propagates past the Router.Observed in production: Anthropic socket timed out before the first chunk on
Router.aresponses(stream=True). Configuredanthropic → vertex_aifallback was never invoked; the error surfaced to the caller.Changes
Router._aresponses_streaming_iteratormirroring_acompletion_streaming_iterator. WrapsBaseResponsesAPIStreamingIterator, catchesMidStreamFallbackError, re-entersasync_function_with_fallbacks_common_utilswithoriginal_function=_ageneric_api_call_with_fallbacks_helperandoriginal_generic_functionpreserved so the helper invokeslitellm.aresponseson each fallback attempt. The wrapper subclassesBaseResponsesAPIStreamingIterator(bypassing the parent constructor) to preserveisinstancecompatibility with downstream consumers (proxy cursor endpoint atlitellm/proxy/response_api_endpoints/endpoints.pyandlitellm/interactions/litellm_responses_transformation/handler.py).Router._aresponses_with_streaming_fallbacksdispatcher that snapshots kwargs (preservingoriginal_generic_function), calls_ageneric_api_call_with_fallbacks, and wraps the result whenstream=Trueand the response is aBaseResponsesAPIStreamingIterator."aresponses"out of the genericfactory_functiondispatch bucket so it flows through the new wrapper.Scope
Pre-first-chunk retry only — covers the observed timeout-before-first-chunk case. The Responses-API input shape differs from chat completions, so the partial-content
prefixassistant-message continuation used in_acompletion_streaming_iteratordoesn't translate cleanly and is intentionally out of scope.Test plan
tests/test_litellm/test_router.py::test_aresponses_streaming_iterator_fallbackmirrorstest_acompletion_streaming_iteratorand asserts:MidStreamFallbackErrortriggersasync_function_with_fallbacks_common_utilskwargs["original_function"]is_ageneric_api_call_with_fallbacks_helperkwargs["original_generic_function"]islitellm.aresponsesmodel_groupis the original model anddisable_fallbacks=False_hidden_paramsandisinstance(BaseResponsesAPIStreamingIterator)_acompletion_streaming_iteratortests still pass.uv run black .anduv run ruff check litellm/router.pyclean on touched files.