scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 71

fix(google_genai): retry a stream that drops before its first event

Browse source 
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
This commit is contained in:
Cabeza 2026-10-12 15:48:43 +08:00
parent 15154a40a8
commit 6879090caa
5 changed files with 277 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/google_genai/streaming_iterator.py
View file

@ -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 = await self._read_next_chunk()
self.collected_chunks.append(chunk)
return chunk
except StopAsyncIteration:

14
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -12055,7 +12055,7 @@ class BaseLLMHTTPHandler:
stream=True,
)
# Return async streaming iterator
return AsyncGoogleGenAIGenerateContentStreamingIterator(
iterator: Final = AsyncGoogleGenAIGenerateContentStreamingIterator(
response=response,
model=model,
logging_obj=logging_obj,
@ -12070,6 +12070,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,

8
tests/unit/google_genai/test_google_genai_main.py
View file

@ -52,6 +52,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

80
tests/unit/google_genai/test_google_genai_streaming_iterator.py
View file

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

170
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -410,9 +410,7 @@ def test_response_api_handler_sends_no_stream_options_when_web_search_forces_a_n
):
sent: Final = _SentRequestRecorder()
recorder: Final = _AgenticLoopKwargsRecorder()
monkeypatch.setattr(
litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"]), recorder]
)
monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"]), recorder])
logging_obj: Final = Mock()
logging_obj.dynamic_success_callbacks = []
@ -2233,7 +2231,9 @@ async def test_async_anthropic_messages_handler_passes_deployment_api_base_to_ag
super().__init__()
self.hook_kwargs: dict | None = None
async def async_should_run_agentic_loop(self, response, model, messages, tools, stream, custom_llm_provider, kwargs):
async def async_should_run_agentic_loop(
self, response, model, messages, tools, stream, custom_llm_provider, kwargs
):
self.hook_kwargs = dict(kwargs)
return False, {}
@ -2593,9 +2593,7 @@ def _transform_subtitle_response(payload):
def test_subtitle_synthesis_fallback_without_timings_drops_words():
response = _transform_subtitle_response(
{"text": "hello world", "words": [{"word": "hello"}, {"word": "world"}]}
)
response = _transform_subtitle_response({"text": "hello world", "words": [{"word": "hello"}, {"word": "world"}]})
assert response.text == "hello world"
assert "words" not in response
@ -3531,7 +3529,13 @@ def _capture_video_create_request(captured):
captured["body"] = request.content
return httpx.Response(
200,
json={"id": "video_123", "object": "video", "status": "queued", "created_at": 1712697600, "model": "sora-2"},
json={
"id": "video_123",
"object": "video",
"status": "queued",
"created_at": 1712697600,
"model": "sora-2",
},
)
return respond
@ -3554,7 +3558,9 @@ def test_video_generation_without_file_sends_multipart_form_data():
captured = {}
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured))))
result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(OpenAIVideoConfig()))
result = BaseLLMHTTPHandler().video_generation_handler(
client=client, **_video_create_call_kwargs(OpenAIVideoConfig())
)
assert captured["content_type"].startswith("multipart/form-data")
assert _multipart_text_fields(captured["content_type"], captured["body"]) == {
@ -3594,7 +3600,9 @@ def test_azure_video_generation_without_file_sends_multipart_form_data():
captured = {}
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured))))
result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(AzureVideoConfig()))
result = BaseLLMHTTPHandler().video_generation_handler(
client=client, **_video_create_call_kwargs(AzureVideoConfig())
)
assert captured["content_type"].startswith("multipart/form-data")
assert _multipart_text_fields(captured["content_type"], captured["body"]) == {
@ -3609,7 +3617,9 @@ def test_video_generation_json_provider_keeps_json_body():
captured = {}
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured))))
result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(_JSONBodyVideoConfig()))
result = BaseLLMHTTPHandler().video_generation_handler(
client=client, **_video_create_call_kwargs(_JSONBodyVideoConfig())
)
assert captured["content_type"] == "application/json"
assert json.loads(captured["body"]) == {"model": "sora-2", "prompt": "a cat surfing", "seconds": "4"}
@ -3639,6 +3649,7 @@ def test_video_generation_with_input_reference_keeps_file_multipart():
AZURE_AI_BASE = "https://myfoundry.services.ai.azure.com"
AZURE_AI_CHAT_COMPLETIONS_URL = f"{AZURE_AI_BASE}/models/chat/completions"
def _a_tool_with_an_unsupported_field() -> dict:
return {
"type": "function",
@ -3646,14 +3657,13 @@ def _a_tool_with_an_unsupported_field() -> dict:
"strict": True,
}
A_COMPLETION = {
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1,
"model": "grok-3",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "sent"}, "finish_reason": "stop"}
],
"choices": [{"index": 0, "message": {"role": "assistant", "content": "sent"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
@ -3697,9 +3707,7 @@ def _call_azure_ai(recorder: _RecordedAzureAI, **overrides):
def test_a_tool_field_the_provider_rejects_is_dropped_and_the_call_retried():
recorder = _RecordedAzureAI(
[_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)]
)
recorder = _RecordedAzureAI([_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)])
response = _call_azure_ai(recorder)
@ -3710,9 +3718,7 @@ def test_a_tool_field_the_provider_rejects_is_dropped_and_the_call_retried():
def test_the_retry_changes_only_the_field_the_provider_named():
recorder = _RecordedAzureAI(
[_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)]
)
recorder = _RecordedAzureAI([_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)])
_call_azure_ai(recorder)
@ -3751,9 +3757,7 @@ def test_an_extra_input_outside_a_tool_is_not_retried_unless_dropping_params_was
def test_an_extra_input_outside_a_tool_is_retried_when_dropping_params_was_asked_for():
recorder = _RecordedAzureAI(
[_rejection(UNRELATED_REJECTION), httpx.Response(200, json=A_COMPLETION)]
)
recorder = _RecordedAzureAI([_rejection(UNRELATED_REJECTION), httpx.Response(200, json=A_COMPLETION)])
response = _call_azure_ai(recorder, drop_params=True)
@ -3767,9 +3771,7 @@ async def test_a_tool_field_the_provider_rejects_is_dropped_and_retried_on_the_a
):
import respx
recorder = _RecordedAzureAI(
[_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)]
)
recorder = _RecordedAzureAI([_rejection(TOOL_LEVEL_REJECTION), httpx.Response(200, json=A_COMPLETION)])
with respx.mock(assert_all_called=True) as router:
router.post(AZURE_AI_CHAT_COMPLETIONS_URL).mock(side_effect=recorder)
@ -3910,7 +3912,11 @@ async def test_create_batch_says_which_setting_is_missing_when_the_provider_reso
with pytest.raises(ValueError, match="api_base is required for create_batch"):
await BaseLLMHTTPHandler().create_batch(
create_batch_data={"input_file_id": "file_1", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
create_batch_data={
"input_file_id": "file_1",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
litellm_params={},
provider_config=provider_config,
headers={},
@ -4091,7 +4097,9 @@ def _start_async_completion(config, logging_obj=None):
custom_llm_provider="openai",
model_response=ModelResponse(),
encoding=None,
logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}),
logging_obj=logging_obj
if logging_obj is not None
else Mock(dynamic_success_callbacks=None, model_call_details={}),
optional_params={},
timeout=10.0,
litellm_params={},
@ -4431,7 +4439,9 @@ async def test_async_realtime_bridges_a_transcription_session_through_the_provid
logging_obj.dispatch_failure_handlers = AsyncMock()
handler = BaseLLMHTTPHandler()
with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(side_effect=AssertionError("dialed a websocket"))) as dial:
with patch.object(
handler, "_open_realtime_backend_ws", AsyncMock(side_effect=AssertionError("dialed a websocket"))
) as dial:
await handler.async_realtime(
model="chirp_3",
websocket=client_ws,
@ -5144,3 +5154,105 @@ async def test_retrieve_file_rejects_unverified_bedrock_range_error(
)
finally:
sync_client.close()
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'