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'