scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(router): read the first stream chunk through __aiter__ so MCP auto-run streaming works again (#46578)

Browse source 
* 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>
This commit is contained in:
Kerry Lu 2026-10-14 18:51:07 -07:00 • committed by GitHub
parent c70dd4ece6
commit 63c3b7a34e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 105 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/router.py
View file

@ -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.

97
tests/unit/test_router/test_router.py
View file

@ -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