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.

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

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)