From 39eddafb315ee74d595e7504646b64f320ab13d1 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Wed, 13 May 2026 16:49:10 +0800 Subject: [PATCH] fix(vertex_ai): raise error on mid-stream 429/error chunks instead of silently swallowing (#23711) Squash-merged by litellm-agent from krisxia0506's PR. --- .../vertex_and_google_ai_studio_gemini.py | 30 ++ ...test_vertex_and_google_ai_studio_gemini.py | 441 ++++++++++++++++++ 2 files changed, 471 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 6278de662f8..f9899197854 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3139,6 +3139,31 @@ class ModelResponseIterator: self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False + @staticmethod + def _check_streaming_error(chunk: dict) -> None: + """Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError.""" + if "error" not in chunk: + return + error_data = chunk["error"] + if not isinstance(error_data, dict): + raise VertexAIError( + status_code=500, + message=f"Unexpected error format in mid-stream chunk: {error_data}", + ) + raw_code = error_data.get("code", 500) + if raw_code is None: + raw_code = 500 + try: + error_code = int(raw_code) + except (TypeError, ValueError): + error_code = 500 + error_message = error_data.get("message", "Unknown error") + error_status = error_data.get("status", "UNKNOWN") + raise VertexAIError( + status_code=error_code, + message=f"{error_status} - {error_message}", + ) + def _apply_stream_candidates( self, _candidates: List[Candidates], @@ -3256,6 +3281,11 @@ class ModelResponseIterator: def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]: try: verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}") + + # Detect mid-stream error chunks (e.g. 429 RESOURCE_EXHAUSTED). + # Vertex AI can return errors as HTTP 200 but with an "error" field in the SSE body. + self._check_streaming_error(chunk) + from litellm.types.utils import ModelResponseStream processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 353d19b0198..1e0ad04c3c2 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -4290,3 +4290,444 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + + +def test_chunk_parser_raises_on_429_error_chunk(): + """Test chunk_parser raises VertexAIError on 429 RESOURCE_EXHAUSTED error chunk""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 429, + "message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.", + "status": "RESOURCE_EXHAUSTED", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + assert "RESOURCE_EXHAUSTED" in exc_info.value.message + assert "Resource exhausted" in exc_info.value.message + + +def test_chunk_parser_raises_on_500_error_chunk(): + """Test chunk_parser raises VertexAIError on 500 INTERNAL error chunk""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 500, + "message": "Internal error encountered.", + "status": "INTERNAL", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "INTERNAL" in exc_info.value.message + + +def test_chunk_parser_raises_on_error_chunk_with_minimal_fields(): + """Test chunk_parser handles error chunks with missing optional fields""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": 429, + "message": "Resource exhausted.", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + + +def test_chunk_parser_normal_chunk_unaffected_by_error_check(): + """Test that normal streaming chunks still work correctly after error check addition""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + normal_chunk = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello"}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 1, + "totalTokenCount": 6, + }, + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + result = streaming_obj.chunk_parser(normal_chunk) + assert result is not None + assert len(result.choices) > 0 + assert result.choices[0].delta.content == "Hello" + + +def test_chunk_parser_raises_on_non_dict_error(): + """Test chunk_parser raises VertexAIError when chunk['error'] is not a dict""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": "something went wrong"} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + + +def test_chunk_parser_raises_on_string_error_code(): + """Test chunk_parser correctly converts string error code to int""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # code field is a string "429" rather than an int + error_chunk = { + "error": { + "code": "429", + "message": "Resource exhausted.", + "status": "RESOURCE_EXHAUSTED", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.status_code, int) + + +def test_chunk_parser_error_chunk_explicit_null_code_uses_500(): + """JSON null for code must not call int(None); status defaults to 500.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": None, + "message": "Something went wrong.", + "status": "UNKNOWN", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Something went wrong" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_numeric_code_defaults_to_500(): + """Non-numeric code must not become ValueError -> RuntimeError in __next__.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = { + "error": { + "code": "NOT_A_NUMBER", + "message": "Malformed.", + "status": "INVALID", + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Malformed" in exc_info.value.message + + +def test_chunk_parser_error_chunk_empty_dict_defaults_to_500(): + """Empty error object {} uses default code 500 and default message/status strings.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": {}} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "UNKNOWN" in exc_info.value.message + assert "Unknown error" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_dict_int_value(): + """Non-dict error payloads (e.g. bare JSON number) must raise with status 500, not TypeError.""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": 503} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + assert "503" in exc_info.value.message + + +def test_chunk_parser_error_chunk_non_dict_null_value(): + """JSON null for error must hit the non-dict branch (same as int/string).""" + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + error_chunk = {"error": None} + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + with pytest.raises(VertexAIError) as exc_info: + streaming_obj.chunk_parser(error_chunk) + + assert exc_info.value.status_code == 500 + assert "Unexpected error format" in exc_info.value.message + + +def test_mid_stream_429_error_raises_during_iteration(): + """ + Simulate a full streaming scenario: normal thinking chunks arrive first, + then a 429 RESOURCE_EXHAUSTED error chunk arrives mid-stream. + Verify that ModelResponseIterator raises VertexAIError during iteration. + """ + import json + from unittest.mock import Mock + + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Simulate Vertex AI SSE stream: normal chunks followed by a 429 error chunk + normal_chunk_1 = json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Let me think about this...", "thought": True}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + "modelVersion": "gemini-3.1-flash-image-preview", + } + ) + + normal_chunk_2 = json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "I'll generate the image now.", "thought": True}], + }, + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 12, + "totalTokenCount": 22, + }, + } + ) + + error_chunk = json.dumps( + { + "error": { + "code": 429, + "message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.", + "status": "RESOURCE_EXHAUSTED", + } + } + ) + + # Build a mock SSE stream (lines returned by iter_lines) + sse_lines = iter([normal_chunk_1, normal_chunk_2, error_chunk]) + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=sse_lines, + sync_stream=True, + logging_obj=logging_obj, + ) + + # Iterate the stream: first chunks should succeed, then 429 error should be raised + results = [] + with pytest.raises(VertexAIError) as exc_info: + for chunk in streaming_obj: + if chunk is not None: + results.append(chunk) + + # Verify: received normal chunks before the error + assert ( + len(results) >= 1 + ), "Should have received at least 1 normal chunk before the error" + + # Verify: 429 error is properly raised + assert exc_info.value.status_code == 429 + assert "RESOURCE_EXHAUSTED" in str(exc_info.value.message)