From 63c3b7a34eebf9050926f516e55efe834e9cab2c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 14 Oct 2026 18:51:07 -0700 Subject: [PATCH] fix(router): read the first stream chunk through __aiter__ so MCP auto-run streaming works again (#46578) * fix(router): read the first stream chunk through __aiter__ so wrappers that hand iteration off (MCPStreamWrapper) stream again * fix(router): mark the _RestT TypeVar Final for the type discipline gate --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 45 ++++++++----- tests/unit/test_router/test_router.py | 97 +++++++++++++++++++++------ 2 files changed, 105 insertions(+), 37 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 9066311451b..0696012c2a5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,7 +38,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, TypeVar, Union, cast import anyio import httpx @@ -694,42 +694,57 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: return error.original_exception if error.original_exception is not None else error +_RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) + + @dataclass(frozen=True, slots=True) -class _FirstChunk: - """The chunk Router._acompletion / _completion already read off a stream; None means the provider closed it empty.""" +class _FirstChunk(Generic[_RestT]): + """ + The chunk Router._acompletion / _completion already read off a stream (None: the + provider closed it empty) and the iterator it was read from, which carries on after it. + """ chunk: ModelResponseStream | None + rest: _RestT -async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk: +async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[AsyncIterator[object]]: """ Read the first chunk while Router._acompletion still holds the request, so a stream that fails before any token raises there, inside async_function_with_retries and the fallback chain, instead of surfacing to the caller mid-iteration. + + Iterates the way `async for` does, one __aiter__ call and then __anext__ on what it + returned: a CustomStreamWrapper subclass may hand iteration to another object and + leave completion_stream None (the MCP gateway's MCPStreamWrapper). """ from litellm.exceptions import MidStreamFallbackError + rest: Final = stream.__aiter__() try: - return _FirstChunk(chunk=await stream.__anext__()) + return _FirstChunk(chunk=await rest.__anext__(), rest=rest) except StopAsyncIteration: - return _FirstChunk(chunk=None) + return _FirstChunk(chunk=None, rest=rest) except MidStreamFallbackError as e: raise _mid_stream_retry_trigger(e) from e -def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk: +def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[Iterator[object]]: """Sync twin of _aread_first_chunk for Router._completion.""" from litellm.exceptions import MidStreamFallbackError + rest: Final = iter(stream) try: - return _FirstChunk(chunk=next(stream)) + return _FirstChunk(chunk=next(rest), rest=rest) except StopIteration: - return _FirstChunk(chunk=None) + return _FirstChunk(chunk=None, rest=rest) except MidStreamFallbackError as e: raise _mid_stream_retry_trigger(e) from e -def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk | None) -> Iterator[object]: +def _chunks_with_first( + model_response: Iterable[object], taken: _FirstChunk[Iterator[object]] | None +) -> Iterator[object]: """The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is).""" if taken is None: yield from model_response @@ -737,11 +752,11 @@ def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk | No if taken.chunk is None: return yield taken.chunk - yield from model_response + yield from taken.rest async def _achunks_with_first( - model_response: AsyncIterable[object], taken: _FirstChunk | None + model_response: AsyncIterable[object], taken: _FirstChunk[AsyncIterator[object]] | None ) -> AsyncIterator[object]: """Async twin of _chunks_with_first.""" if taken is None: @@ -751,7 +766,7 @@ async def _achunks_with_first( if taken.chunk is None: return yield taken.chunk - async for item in model_response: + async for item in taken.rest: yield item @@ -3117,7 +3132,7 @@ class Router: messages: list[dict[str, str]], initial_kwargs: dict, deployment_slot: contextlib.AsyncExitStack | None = None, - first_chunk: _FirstChunk | None = None, + first_chunk: _FirstChunk[AsyncIterator[object]] | None = None, ) -> CustomStreamWrapper: """ Helper to iterate over a streaming response. @@ -3747,7 +3762,7 @@ class Router: model_response: CustomStreamWrapper, messages: list[dict[str, str]], initial_kwargs: dict, - first_chunk: _FirstChunk | None = None, + first_chunk: _FirstChunk[Iterator[object]] | None = None, ) -> CustomStreamWrapper: """ Sync equivalent of _acompletion_streaming_iterator. diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 0a795ff2a0f..aa89a6ad488 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -26243,16 +26243,21 @@ class _FirstPullStream: raise self._first return self._first + def __iter__(self) -> "_FirstPullStream": + return self + def __next__(self) -> object: return self._pull(StopIteration) + def __aiter__(self) -> "_FirstPullStream": + return self + async def __anext__(self) -> object: return self._pull(StopAsyncIteration) def _first_chunk_cases(): from litellm.exceptions import MidStreamFallbackError - from litellm.router import _FirstChunk from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) @@ -26265,8 +26270,8 @@ def _first_chunk_cases(): is_pre_first_chunk=True, ) return [ - pytest.param(chunk, _FirstChunk(chunk=chunk), None, id="first-chunk"), - pytest.param(None, _FirstChunk(chunk=None), None, id="empty-stream"), + pytest.param(chunk, chunk, None, id="first-chunk"), + pytest.param(None, None, None, id="empty-stream"), pytest.param(wrapped, None, provider_error, id="pre-content-failure"), ] @@ -26275,11 +26280,13 @@ def _first_chunk_cases(): def test_read_first_chunk_direct_call(first, expected, raised): from litellm.router import _read_first_chunk + stream = _FirstPullStream(first) if raised is None: - assert _read_first_chunk(_FirstPullStream(first)) == expected + taken = _read_first_chunk(stream) + assert (taken.chunk, taken.rest) == (expected, stream) return with pytest.raises(type(raised)) as caught: - _read_first_chunk(_FirstPullStream(first)) + _read_first_chunk(stream) assert caught.value is raised @@ -26288,38 +26295,84 @@ def test_read_first_chunk_direct_call(first, expected, raised): async def test_aread_first_chunk_direct_call(first, expected, raised): from litellm.router import _aread_first_chunk + stream = _FirstPullStream(first) if raised is None: - assert await _aread_first_chunk(_FirstPullStream(first)) == expected + taken = await _aread_first_chunk(stream) + assert (taken.chunk, taken.rest) == (expected, stream) return with pytest.raises(type(raised)) as caught: - await _aread_first_chunk(_FirstPullStream(first)) + await _aread_first_chunk(stream) assert caught.value is raised def _chunks_with_first_cases(): - from litellm.router import _FirstChunk - return [ pytest.param(None, ["a", "b"], ["a", "b"], id="nothing-read-stream-as-is"), - pytest.param(_FirstChunk(chunk=None), ["a", "b"], [], id="empty-stream-not-touched"), - pytest.param(_FirstChunk(chunk="first"), ["a", "b"], ["first", "a", "b"], id="first-chunk-put-back"), + pytest.param((None, ["a", "b"]), ["a", "b"], [], id="empty-stream-not-touched"), + pytest.param(("first", ["a", "b"]), ["never-iterated"], ["first", "a", "b"], id="first-chunk-put-back"), ] -@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases()) -def test_chunks_with_first_direct_call(taken, rest, expected): - from litellm.router import _chunks_with_first +@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) +def test_chunks_with_first_direct_call(taken, stream, expected): + from litellm.router import _chunks_with_first, _FirstChunk - assert list(_chunks_with_first(iter(rest), taken)) == expected + first = None if taken is None else _FirstChunk(chunk=taken[0], rest=iter(taken[1])) + assert list(_chunks_with_first(iter(stream), first)) == expected @pytest.mark.asyncio -@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases()) -async def test_achunks_with_first_direct_call(taken, rest, expected): - from litellm.router import _achunks_with_first +@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) +async def test_achunks_with_first_direct_call(taken, stream, expected): + from litellm.router import _achunks_with_first, _FirstChunk - async def stream(): - for item in rest: - yield item + first = None if taken is None else _FirstChunk(chunk=taken[0], rest=_AsyncList(taken[1])) + assert [item async for item in _achunks_with_first(_AsyncList(stream), first)] == expected - assert [item async for item in _achunks_with_first(stream(), taken)] == expected + +class _IterationHandedOffStream(CustomStreamWrapper): + """A CustomStreamWrapper that, like the MCP gateway's MCPStreamWrapper, has no completion_stream of its + own and hands iteration to another object from __aiter__ / __iter__, a fresh one per call so a second + __aiter__ would replay the first chunk.""" + + def __init__(self, chunks: list[litellm.ModelResponseStream]) -> None: + super().__init__(completion_stream=None, model="gpt-5.4", custom_llm_provider="openai", logging_obj=MagicMock()) + self._chunks: Final = chunks + + def __aiter__(self): + return _AsyncList(self._chunks) + + def __iter__(self): + return iter(self._chunks) + + +def _handed_off_chunks() -> list[litellm.ModelResponseStream]: + return [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in ("tool ", "answer") + ] + + +@pytest.mark.asyncio +async def test_acompletion_streams_a_wrapper_that_hands_iteration_to_another_iterator(): + """The eager first-chunk read goes through __aiter__ like `async for` does, so a wrapper with + completion_stream=None (MCPStreamWrapper) streams instead of failing with a NoneType iterator.""" + router = _mid_stream_opt_out_router() + chunks = _handed_off_chunks() + + with patch("litellm.acompletion", new=AsyncMock(return_value=_IterationHandedOffStream(chunks))): + response = await router.acompletion(model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True) + streamed = [chunk async for chunk in response] + + assert streamed == chunks + + +def test_completion_streams_a_wrapper_that_hands_iteration_to_another_iterator(): + """Sync twin: Router._completion's first-chunk read goes through __iter__ like a for loop does.""" + router = _mid_stream_opt_out_router() + chunks = _handed_off_chunks() + + with patch("litellm.completion", return_value=_IterationHandedOffStream(chunks)): + response = router.completion(model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True) + streamed = list(response) + + assert streamed == chunks