From c908505e6a527e2328e755c9dd003b288e836013 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 2 Jun 2026 03:08:19 +0530 Subject: [PATCH] fix(proxy): omit OpenAI [DONE] on google-genai streamGenerateContent (#29426) * fix(proxy): omit OpenAI [DONE] on google-genai streamGenerateContent google-genai SDK uses ?alt=sse and cannot parse the proxy's trailing data: [DONE] chunk. Skip that terminator for agenerate_content_stream. Co-authored-by: Cursor * fix(proxy): address Greptile review on google-genai stream fix Always yield stream error_message; only gate data: [DONE] on the skip flag. Set _litellm_skip_openai_stream_done in google_endpoints instead of common_request_processing. Co-authored-by: Cursor --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/google_endpoints/endpoints.py | 2 + litellm/proxy/proxy_server.py | 7 +- tests/test_litellm/proxy/test_proxy_server.py | 104 ++++++++++++++++++ 3 files changed, 110 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 1f503247bf4..cc20f0cf3b3 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -105,6 +105,8 @@ async def google_stream_generate_content( if "model" not in data: data["model"] = model_name data["stream"] = True + # google-genai SDK (?alt=sse) must not receive OpenAI's data: [DONE] terminator. + data["_litellm_skip_openai_stream_done"] = True processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b296792cd09..e0f139dee57 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7072,11 +7072,12 @@ async def async_data_generator( # noqa: PLR0915 # still flush their post-stream logging. ProxyLogging._fire_deferred_stream_logging(request_data) - # Streaming is done, yield the [DONE] chunk if error_message is not None: yield error_message - done_message = "[DONE]" - yield f"data: {done_message}\n\n" + # OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not. + if not request_data.get("_litellm_skip_openai_stream_done"): + done_message = "[DONE]" + yield f"data: {done_message}\n\n" except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b912fec2479..8aa839cdfcb 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5165,6 +5165,110 @@ async def test_async_data_generator_passes_through_google_native_sse_bytes(): assert yielded_text[-1] == "data: [DONE]\n\n" +@pytest.mark.asyncio +async def test_async_data_generator_google_genai_stream_omits_openai_done(): + """ + google-genai SDK streamGenerateContent?alt=sse must not receive data: [DONE]. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_request_data = { + "model": "gemini-2.0-flash", + "_litellm_skip_openai_stream_done": True, + } + gemini_event = ( + b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n' + ) + + class MockStream: + def __aiter__(self): + return self._stream() + + async def _stream(self): + yield gemini_event + + async def aclose(self): + pass + + mock_response = MockStream() + mock_response.aclose = AsyncMock() + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.has_streaming_callbacks.return_value = False + mock_proxy_logging_obj.needs_iterator_wrap.return_value = False + mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock() + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): + with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): + yielded_data = [] + async for data in async_data_generator( + mock_response, mock_user_api_key_dict, mock_request_data + ): + yielded_data.append(data) + + yielded_text = [ + chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + for chunk in yielded_data + ] + assert yielded_text == [gemini_event.decode("utf-8")] + assert "[DONE]" not in "".join(yielded_text) + + +@pytest.mark.asyncio +async def test_async_data_generator_google_genai_stream_forwards_error_without_done(): + """Stream errors must still reach the client when OpenAI [DONE] is skipped.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + + error_sse = 'data: {"error": {"message": "stream failed"}}\n\n' + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_request_data = { + "model": "gemini-2.0-flash", + "_litellm_skip_openai_stream_done": True, + } + + class MockStream: + def __aiter__(self): + return self._stream() + + async def _stream(self): + yield error_sse + + async def aclose(self): + pass + + mock_response = MockStream() + mock_response.aclose = AsyncMock() + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.has_streaming_callbacks.return_value = False + mock_proxy_logging_obj.needs_iterator_wrap.return_value = False + mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock() + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): + with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): + yielded_data = [] + async for data in async_data_generator( + mock_response, mock_user_api_key_dict, mock_request_data + ): + yielded_data.append(data) + + yielded_text = [ + chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + for chunk in yielded_data + ] + assert yielded_text == [error_sse] + assert "[DONE]" not in "".join(yielded_text) + + @pytest.mark.asyncio async def test_async_data_generator_cleanup_on_normal_completion(): """