From 5e00a49e7f66190bf06e3c4d6528becfce36c3ce Mon Sep 17 00:00:00 2001
From: Kris Xia <xiajiayi0506@gmail.com>
Date: Wed, 7 Jan 2026 02:11:23 +0800
Subject: [PATCH] fix(streaming): normalize status code extraction to prevent
 4xx errors from triggering mid-stream fallback (#18698)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

在流式处理错误时，添加状态码标准化逻辑，确保 4xx 客户端错误直接抛出而不是被包装成 MidStreamFallbackError。

- 新增 _normalize_status_code 函数用于从异常对象提取状态码
- 优先从异常的 status_code 属性获取，其次从 response.status_code 获取
- 当映射异常或原始异常的状态码在 400-499 范围内时，直接抛出映射异常
- 添加单元测试验证 Vertex AI 400 错误正确抛出为 BadRequestError
- 确保流式处理中的客户端错误能够正确传播，而不会触发回退机制
---
 .../litellm_core_utils/streaming_handler.py   |  54 ++-
 .../test_streaming_handler.py                 |  23 ++
 .../test_function_call_args_serialization.py  | 355 ++++++++++++++++++
 ...test_vertex_and_google_ai_studio_gemini.py |  35 ++
 4 files changed, 456 insertions(+), 11 deletions(-)
 create mode 100644 tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index d92af41717..6baaae7ae3 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -2000,24 +2000,56 @@ class CustomStreamWrapper:
                 )
             ## Map to OpenAI Exception
             try:
-                raise exception_type(
+                mapped_exception = exception_type(
                     model=self.model,
                     custom_llm_provider=self.custom_llm_provider,
                     original_exception=e,
                     completion_kwargs={},
                     extra_kwargs={},
                 )
-            except Exception as e:
-                from litellm.exceptions import MidStreamFallbackError
+            except Exception as mapping_error:
+                mapped_exception = mapping_error
 
-                raise MidStreamFallbackError(
-                    message=str(e),
-                    model=self.model,
-                    llm_provider=self.custom_llm_provider or "anthropic",
-                    original_exception=e,
-                    generated_content=self.response_uptil_now,
-                    is_pre_first_chunk=not self.sent_first_chunk,
-                )
+            def _normalize_status_code(exc: Exception) -> Optional[int]:
+                """
+                Best-effort status_code extraction.
+                Uses status_code on the exception, then falls back to the response.
+                """
+                try:
+                    code = getattr(exc, "status_code", None)
+                    if code is not None:
+                        return int(code)
+                except Exception:
+                    pass
+
+                response = getattr(exc, "response", None)
+                if response is not None:
+                    try:
+                        status_code = getattr(response, "status_code", None)
+                        if status_code is not None:
+                            return int(status_code)
+                    except Exception:
+                        pass
+                return None
+
+            mapped_status_code = _normalize_status_code(mapped_exception)
+            original_status_code = _normalize_status_code(e)
+
+            if mapped_status_code is not None and 400 <= mapped_status_code < 500:
+                raise mapped_exception
+            if original_status_code is not None and 400 <= original_status_code < 500:
+                raise mapped_exception
+
+            from litellm.exceptions import MidStreamFallbackError
+
+            raise MidStreamFallbackError(
+                message=str(mapped_exception),
+                model=self.model,
+                llm_provider=self.custom_llm_provider or "anthropic",
+                original_exception=mapped_exception,
+                generated_content=self.response_uptil_now,
+                is_pre_first_chunk=not self.sent_first_chunk,
+            )
 
     @staticmethod
     def _strip_sse_data_from_chunk(chunk: Optional[str]) -> Optional[str]:
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
index 6a528fef8f..ec2f528a35 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -691,6 +691,29 @@ async def test_streaming_completion_start_time(logging_obj: Logging):
     )
 
 
+@pytest.mark.asyncio
+async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging):
+    """Ensure Vertex bad request errors surface as 400, not mid-stream fallbacks."""
+    from litellm.llms.vertex_ai.common_utils import VertexAIError
+
+    async def _raise_bad_request(**kwargs):
+        raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None)
+
+    response = CustomStreamWrapper(
+        completion_stream=None,
+        model="gemini-3-pro-preview",
+        logging_obj=logging_obj,
+        custom_llm_provider="vertex_ai_beta",
+        make_call=_raise_bad_request,
+    )
+
+    with pytest.raises(litellm.BadRequestError) as excinfo:
+        await response.__anext__()
+
+    assert getattr(excinfo.value, "status_code", None) == 400
+    assert "invalid maxOutputTokens" in str(excinfo.value)
+
+
 def test_streaming_handler_with_created_time_propagation(
     initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging
 ):
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py b/tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py
new file mode 100644
index 0000000000..0f369fbb8b
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py
@@ -0,0 +1,355 @@
+"""
+Test cases for functionCall args serialization in Vertex AI Gemini.
+
+This test file specifically tests the edge cases where Vertex AI might return
+functionCall args in unexpected formats that could lead to invalid JSON strings
+like: {"x":"x"}{"a":"a"}
+"""
+import json
+from typing import List, Optional
+
+import pytest
+
+from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+    VertexGeminiConfig,
+)
+from litellm.types.llms.vertex_ai import HttpxPartType
+
+
+class TestFunctionCallArgsSerialization:
+    """Test cases for functionCall args serialization edge cases."""
+
+    def test_normal_dict_args(self):
+        """Test normal case: args is a dict."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": {"location": "Boston", "unit": "celsius"},
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        assert tools[0]["function"]["name"] == "get_weather"
+        
+        # Verify arguments is a valid JSON string
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        # Should be valid JSON
+        parsed = json.loads(arguments)
+        assert parsed == {"location": "Boston", "unit": "celsius"}
+
+    def test_none_args(self):
+        """Test case: args is None."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": None,
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        # Should serialize None to "null" or empty dict
+        assert isinstance(arguments, str)
+        parsed = json.loads(arguments)
+        # json.dumps(None) returns "null"
+        assert parsed is None or parsed == {}
+
+    def test_args_as_string_valid_json(self):
+        """Test case: args is already a valid JSON string."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": '{"location": "Boston"}',  # String, not dict
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        # If args is a string, json.dumps will double-encode it
+        # This would result in: "{\"location\": \"Boston\"}"
+        assert isinstance(arguments, str)
+        # This is the problematic case - string gets double-encoded
+        # The result would be a JSON string containing a JSON string
+        parsed = json.loads(arguments)
+        # If it's double-encoded, parsed would be a string, not a dict
+        if isinstance(parsed, str):
+            # Double-encoded case
+            inner_parsed = json.loads(parsed)
+            assert inner_parsed == {"location": "Boston"}
+        else:
+            # Normal case (shouldn't happen if args is string)
+            assert parsed == {"location": "Boston"}
+
+    def test_args_as_string_invalid_json_concatenated(self):
+        """Test case: args is a string with concatenated JSON objects (the bug case).
+        
+        When args is a string like '{"x":"x"}{"a":"a"}', json.dumps() will serialize it
+        as a JSON string, resulting in: "{\"x\":\"x\"}{\"a\":\"a\"}"
+        This is a valid JSON string (the outer quotes), but the content inside is invalid JSON.
+        When you try to parse the inner content, it fails.
+        """
+        # This simulates the case where Vertex might return something like:
+        # args = '{"x":"x"}{"a":"a"}'  # Two JSON objects concatenated
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": '{"x":"x"}{"a":"a"}',  # Invalid concatenated JSON
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        
+        # json.dumps() on a string will escape it, so we get:
+        # arguments = '"{\\"x\\":\\"x\\"}{\\"a\\":\\"a\\"}"'
+        # This is a valid JSON string (the outer quotes), but the inner content is invalid
+        parsed_outer = json.loads(arguments)
+        assert isinstance(parsed_outer, str)
+        
+        # The inner string is invalid JSON (two objects concatenated)
+        # This is the bug: the inner content cannot be parsed as valid JSON
+        with pytest.raises(json.JSONDecodeError):
+            json.loads(parsed_outer)
+        
+        # The arguments string would be: "{\"x\":\"x\"}{\"a\":\"a\"}"
+        # Which when parsed gives: '{"x":"x"}{"a":"a"}' (invalid JSON)
+
+    def test_args_as_array(self):
+        """Test case: args is an array (unexpected but possible)."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": [{"x": "x"}, {"a": "a"}],  # Array of objects
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        # Should serialize array correctly
+        parsed = json.loads(arguments)
+        assert parsed == [{"x": "x"}, {"a": "a"}]
+
+    def test_args_missing_key(self):
+        """Test case: args key is missing from functionCall.
+        
+        This will raise a KeyError because the code directly accesses part["functionCall"]["args"]
+        without checking if the key exists. This is a bug that should be fixed.
+        """
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    # args key missing
+                }
+            }
+        ]
+
+        # This should raise KeyError because args key is missing
+        with pytest.raises(KeyError):
+            VertexGeminiConfig._transform_parts(
+                parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+            )
+
+    def test_multiple_function_calls(self):
+        """Test case: multiple function calls in parts."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": {"location": "Boston"},
+                }
+            },
+            {
+                "functionCall": {
+                    "name": "get_time",
+                    "args": {"timezone": "EST"},
+                }
+            },
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 2
+        assert tools[0]["function"]["name"] == "get_weather"
+        assert tools[1]["function"]["name"] == "get_time"
+        
+        # Both should have valid JSON arguments
+        args1 = json.loads(tools[0]["function"]["arguments"])
+        args2 = json.loads(tools[1]["function"]["arguments"])
+        assert args1 == {"location": "Boston"}
+        assert args2 == {"timezone": "EST"}
+
+    def test_args_with_vertex_protobuf_format(self):
+        """Test case: args in Vertex protobuf format with string_value, etc."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": {
+                        "location": {"string_value": "Boston, MA"},
+                        "unit": {"string_value": "celsius"},
+                    },
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        # Should serialize the nested structure correctly
+        parsed = json.loads(arguments)
+        assert "location" in parsed
+        assert "unit" in parsed
+
+    def test_args_as_empty_dict(self):
+        """Test case: args is an empty dict."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": {},
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        parsed = json.loads(arguments)
+        assert parsed == {}
+
+    def test_args_with_special_characters(self):
+        """Test case: args contains special characters that need escaping."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": {
+                        "location": 'Boston, MA "downtown"',
+                        "note": "Line 1\nLine 2",
+                    },
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        # Should handle special characters correctly
+        parsed = json.loads(arguments)
+        assert parsed["location"] == 'Boston, MA "downtown"'
+        assert parsed["note"] == "Line 1\nLine 2"
+
+    def test_args_as_list_of_strings_that_look_like_json(self):
+        """Test case: args is a list containing strings that look like JSON objects."""
+        # This could potentially cause issues if not handled correctly
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "get_weather",
+                    "args": ['{"x":"x"}', '{"a":"a"}'],  # List of JSON strings
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        # Should serialize list correctly
+        parsed = json.loads(arguments)
+        assert isinstance(parsed, list)
+        assert parsed == ['{"x":"x"}', '{"a":"a"}']
+
+    def test_args_as_dict_with_nested_structures(self):
+        """Test case: args contains nested dicts and lists."""
+        parts: List[HttpxPartType] = [
+            {
+                "functionCall": {
+                    "name": "complex_function",
+                    "args": {
+                        "nested": {"key": "value"},
+                        "list": [1, 2, 3],
+                        "mixed": [{"a": 1}, {"b": 2}],
+                    },
+                }
+            }
+        ]
+
+        function, tools, idx = VertexGeminiConfig._transform_parts(
+            parts=parts, cumulative_tool_call_idx=0, is_function_call=False
+        )
+
+        assert tools is not None
+        assert len(tools) == 1
+        arguments = tools[0]["function"]["arguments"]
+        assert isinstance(arguments, str)
+        parsed = json.loads(arguments)
+        assert parsed["nested"] == {"key": "value"}
+        assert parsed["list"] == [1, 2, 3]
+        assert parsed["mixed"] == [{"a": 1}, {"b": 2}]
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
+
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 91a28ee6ec..d09de3a0f2 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
@@ -10,11 +10,13 @@ from pydantic import BaseModel
 
 import litellm
 from litellm import ModelResponse, completion
+from litellm.llms.vertex_ai.common_utils import VertexAIError
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     VertexGeminiConfig,
 )
 from litellm.types.llms.vertex_ai import UsageMetadata
 from litellm.types.utils import ChoiceLogprobs, Usage
+from litellm.utils import CustomStreamWrapper
 
 
 def test_top_logprobs():
@@ -1605,6 +1607,39 @@ def test_vertex_ai_annotation_streaming_events():
     assert "Weather information" in annotation["url_citation"]["title"]
 
 
+@pytest.mark.asyncio
+async def test_vertex_ai_streaming_bad_request_is_not_wrapped():
+    class DummyLogging:
+        def __init__(self):
+            self.model_call_details = {"litellm_params": {}}
+            self.optional_params = {}
+            self.messages = []
+            self.completion_start_time = None
+            self.stream_options = None
+
+        def failure_handler(self, *args, **kwargs):
+            return None
+
+        async def async_failure_handler(self, *args, **kwargs):
+            return None
+
+    async def failing_make_call(client=None, **kwargs):
+        raise VertexAIError(status_code=400, message="bad input", headers={})
+
+    stream = CustomStreamWrapper(
+        completion_stream=None,
+        make_call=failing_make_call,
+        model="gemini-3-pro-preview",
+        logging_obj=DummyLogging(),
+        custom_llm_provider="vertex_ai_beta",
+    )
+
+    with pytest.raises(litellm.BadRequestError) as exc_info:
+        await stream.__anext__()
+
+    assert getattr(exc_info.value, "status_code", None) == 400
+
+
 def test_vertex_ai_annotation_conversion():
     """
     Test the conversion of Vertex AI grounding metadata to OpenAI annotations.