Merge 9db6e3b126 into d56ff281de

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:

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,

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
 
 

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'

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'