From 9db6e3b1265cb1a0d23b9eee4304c70ed600fee2 Mon Sep 17 00:00:00 2001 From: insuffer Date: Sat, 10 Oct 2026 14:16:09 +0800 Subject: [PATCH] fix(google_genai): retry a stream that drops before its first event The async generateContent handler returned its streaming iterator as soon as the response headers arrived, so the first body read happened in the proxy's data generator, long after the router had recorded the call as a success and spent its retry budget. A stream that dropped before its first event therefore reached the client as an in-stream error frame with attempted_retries unset, and router_settings.num_retries never applied. Read the first SSE event inside the handler instead, so a drop raises where the router's retry loop can see it and is mapped to a retriable status. Gemini sends no lifecycle frames before content, so one event proves the stream is alive, and the event is replayed to the consumer unchanged. A stream that ends before its first event stays a valid empty answer and ends the iterator with its logging intact. The abandoned response is closed when priming fails, so a dropped stream does not hold a pooled connection while the retry runs. The synchronous path has the same shape but is out of scope here: it has no caller that returns before the first read. Refs #45457 --- litellm/google_genai/streaming_iterator.py | 37 ++++++- litellm/llms/custom_httpx/llm_http_handler.py | 14 ++- .../google_genai/test_google_genai_main.py | 8 ++ .../test_google_genai_streaming_iterator.py | 80 +++++++++++++- .../custom_httpx/test_llm_http_handler.py | 102 ++++++++++++++++++ 5 files changed, 238 insertions(+), 3 deletions(-) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 37d08dfb59a..c9bee357e76 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -185,13 +185,48 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. self.stream_iterator = response.aiter_lines() + self._primed_chunk: bytes | None = None + self._primed_exhausted: bool = False + + async def _read_next_chunk(self) -> bytes: + """Read one SSE event off the line iterator. Callers record it, so a primed event is counted once, when it is handed over.""" + chunk: Final[bytes] = await _anext_google_genai_sse_chunk(self.stream_iterator) + return chunk + + async def prime_first_chunk(self) -> None: + """ + Read the first SSE event now, inside the caller's own retry window. + + The handler that builds this iterator returns as soon as the response headers + arrive, so a stream that drops before its first event would otherwise fail in + whoever consumes the iterator, long after the call that opened it was recorded + as a success and its retry budget spent. Gemini sends no lifecycle frames + before content, so one event is enough to prove the stream is alive. A stream + that ends before its first event is a valid empty answer and is replayed as an + immediate end of stream. + """ + if self._primed_chunk is not None or self._primed_exhausted: + return + try: + self._primed_chunk = await self._read_next_chunk() + except StopAsyncIteration: + self._primed_exhausted = True def __aiter__(self): return self async def __anext__(self): + if self._primed_chunk is not None: + primed: Final[bytes] = self._primed_chunk + self._primed_chunk = None + self.collected_chunks.append(primed) + return primed + if self._primed_exhausted: + self._primed_exhausted = False + await self._handle_async_streaming_logging() + raise StopAsyncIteration try: - chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator) + chunk: Final[bytes] = await self._read_next_chunk() self.collected_chunks.append(chunk) return chunk except StopAsyncIteration: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 48858dfa441..c39cbfaa4b3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12266,7 +12266,7 @@ class BaseLLMHTTPHandler: stream=True, ) # Return async streaming iterator - return AsyncGoogleGenAIGenerateContentStreamingIterator( + iterator: Final = AsyncGoogleGenAIGenerateContentStreamingIterator( response=response, model=model, logging_obj=logging_obj, @@ -12281,6 +12281,18 @@ class BaseLLMHTTPHandler: response_headers=response.headers, ), ) + # Read the first event before returning: this call is what the router + # wraps in its retry loop, so a stream that drops before its first event + # has to fail here, not in the consumer that reads the iterator later. + try: + await iterator.prime_first_chunk() + except Exception: + try: + await response.aclose() + except Exception as close_error: # noqa: BLE001 # closing must not mask the stream failure + verbose_logger.debug("Could not close the dropped google genai stream: %s", close_error) + raise + return iterator else: response = await async_httpx_client.post( url=api_base, diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 5cac61e727e..47340001f10 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -54,6 +54,14 @@ def _mock_gemini_post_response(): } ] } + + async def _aiter_lines(): + # The async handler reads the first SSE event before returning, so a stream + # stand-in has to end (a bare MagicMock yields mock lines forever). + yield 'data: {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}' + yield "" + + resp.aiter_lines = _aiter_lines return resp diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py index e8ec2848233..e87281c74ef 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -1,6 +1,7 @@ import json -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from litellm.google_genai.streaming_iterator import ( @@ -154,3 +155,80 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): chunk = await iterator.__anext__() assert chunk == b": keepalive\n\n" + + +def _async_response_with_lines(lines: list[str], drop_after: Exception | None = None): + """An httpx.Response stand-in whose line iterator yields `lines`, then raises `drop_after`.""" + mock_response = MagicMock() + + async def _aiter_lines(): + for line in lines: + yield line + if drop_after is not None: + raise drop_after + + mock_response.aiter_lines = _aiter_lines + return mock_response + + +def _async_iterator(mock_response) -> AsyncGoogleGenAIGenerateContentStreamingIterator: + return AsyncGoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-3.8-flash", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="vertex_ai", + ) + + +@pytest.mark.asyncio +async def test_priming_surfaces_a_stream_that_drops_before_its_first_event(): + """A drop on the first read fails the call that opened the stream, not the consumer that reads it later.""" + mock_response = _async_response_with_lines([], drop_after=httpx.ReadError("Connection closed.")) + + iterator = _async_iterator(mock_response) + + with pytest.raises(httpx.ReadError): + await iterator.prime_first_chunk() + + +@pytest.mark.asyncio +async def test_priming_replays_the_first_event_without_losing_or_duplicating_it(): + mock_response = _async_response_with_lines(['data: {"first": 1}', "", 'data: {"second": 2}', ""]) + + iterator = _async_iterator(mock_response) + await iterator.prime_first_chunk() + + assert await iterator.__anext__() == b'data: {"first": 1}\n\n' + assert await iterator.__anext__() == b'data: {"second": 2}\n\n' + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + + +@pytest.mark.asyncio +async def test_priming_an_empty_stream_ends_the_iterator_and_still_logs(): + """A stream that ends before its first event is a valid (empty) answer, not a failure.""" + mock_response = _async_response_with_lines([]) + + iterator = _async_iterator(mock_response) + iterator._handle_async_streaming_logging = AsyncMock() + + await iterator.prime_first_chunk() + + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + + iterator._handle_async_streaming_logging.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_priming_twice_does_not_consume_a_second_event(): + mock_response = _async_response_with_lines(['data: {"first": 1}', "", 'data: {"second": 2}', ""]) + + iterator = _async_iterator(mock_response) + await iterator.prime_first_chunk() + await iterator.prime_first_chunk() + + assert await iterator.__anext__() == b'data: {"first": 1}\n\n' + assert await iterator.__anext__() == b'data: {"second": 2}\n\n' diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 8b46597b603..0d511dffe89 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5391,3 +5391,105 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ sent: Final = received[0] assert "safeguards" not in json.loads(sent.content) assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) + + +class _DroppingByteStream(httpx.AsyncByteStream): + """A response body that fails on its first read, the way a dropped connection does.""" + + async def __aiter__(self): + raise httpx.ReadError("Connection closed.") + yield b"" # pragma: no cover - unreachable, makes this an async generator + + +def _google_genai_streaming_logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="agenerate_content_stream", + start_time=time.time(), + litellm_call_id="google-genai-stream-priming", + function_id="f", + ) + logging_obj.update_environment_variables( + model="gemini-3.8-flash", + user="", + optional_params={}, + litellm_params={}, + custom_llm_provider="gemini", + ) + return logging_obj + + +async def _call_google_genai_stream(upstream: httpx.Response): + """Drive async_generate_content_handler against `upstream` faked at the HTTP boundary. + + The provider config and the HTTP client are the real ones; only the upstream + response is a double, so the error mapping under test is the provider's own. + """ + from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig + + transport: Final = httpx.MockTransport(lambda request: upstream) + async with httpx.AsyncClient(transport=transport) as async_client: + client: Final = AsyncHTTPHandler() + await client.close() + client.client = async_client + return await BaseLLMHTTPHandler().async_generate_content_handler( + model="gemini-3.8-flash", + contents=[{"role": "user", "parts": [{"text": "hi"}]}], + generate_content_provider_config=GoogleGenAIConfig(), + generate_content_config_dict={}, + tools=None, + custom_llm_provider="gemini", + litellm_params=GenericLiteLLMParams( + api_key="test-key", + api_base="https://generativelanguage.googleapis.com", + ), + logging_obj=_google_genai_streaming_logging_obj(), + client=client, + stream=True, + ) + + +@pytest.mark.asyncio +async def test_google_genai_async_stream_priming_fails_retriably_when_the_stream_drops_before_its_first_event(): + """ + A stream that drops before its first event must fail inside the caller's own retry window. + + The handler returns as soon as the response headers arrive, so without priming the first + body read happens in the proxy's data generator, long after the router recorded the call + as a success and spent its retry budget. + """ + upstream: Final = httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + stream=_DroppingByteStream(), + ) + + with pytest.raises(BaseLLMException) as exc_info: + await _call_google_genai_stream(upstream) + + # The router retries on this status, so the drop now lands inside its retry loop. + assert litellm.should_retry(exc_info.value.status_code) + + +@pytest.mark.asyncio +async def test_google_genai_async_stream_priming_returns_an_iterator_that_still_yields_the_whole_stream(): + """Priming must not consume the first event: the consumer sees the entire stream.""" + from litellm.google_genai.streaming_iterator import ( + AsyncGoogleGenAIGenerateContentStreamingIterator, + ) + + upstream: Final = httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"first": 1}\n\ndata: {"second": 2}\n\n', + ) + + iterator = await _call_google_genai_stream(upstream) + + assert isinstance(iterator, AsyncGoogleGenAIGenerateContentStreamingIterator) + assert await iterator.__anext__() == b'data: {"first": 1}\n\n' + assert await iterator.__anext__() == b'data: {"second": 2}\n\n'