From 970ea2949eec6c40091b64098a186560747d301d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 25 Jul 2026 17:13:23 -0700 Subject: [PATCH] fix(vertex): decide rawPredict passthrough streaming from the request body Vertex passthrough classified any target URL containing "stream" as a streaming request. `:streamRawPredict` carries that substring, so a unary Claude-on-Vertex call whose body omits `stream` was routed through the streaming logging path. That path never consults the response content-type, so a complete `"type": "message"` JSON body was handed to the Anthropic SSE chunk parser, which recognises none of it; the spend log recorded 0 prompt tokens, 0 completion tokens and zero cost Streaming for the rawPredict family now comes from the request body, which is what the Anthropic Messages contract uses for those endpoints. The generateContent family keeps its URL signal because the Gemini REST body has no `stream` field, and `?alt=sse` is still appended for every request that is classified as streaming, so Gemini framing and its usage parsing are unchanged Both passthrough streaming predicates read `.get("stream")` off a body that is only annotated as a dict; `_read_request_body` returns whatever the JSON parser produced, so an array body raised AttributeError. The two predicates are now one owner that answers False for any non-object body, which covers the vertex, mistral, anthropic, vllm and azure passthrough routes --- .../llm_passthrough_endpoints.py | 22 ++- .../test_llm_pass_through_endpoints.py | 169 ++++++++++++++++++ 2 files changed, 183 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7e573de261b..28d2c62f1f1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -86,11 +86,16 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: Op return False -def is_passthrough_request_streaming(request_body: dict) -> bool: +def is_passthrough_request_streaming(request_body: object) -> bool: """ - Returns True if the request is streaming + Returns True if the request is streaming. + + A JSON body need not be an object, so a list or scalar can reach here; it + carries no streaming flag. """ - return request_body.get("stream", False) + if not isinstance(request_body, dict): + return False + return bool(request_body.get("stream", False)) async def llm_passthrough_factory_proxy_route( @@ -551,8 +556,7 @@ async def is_streaming_request_fn(request: Request) -> bool: _request_body = await get_form_data(request) else: _request_body = await _read_request_body(request) - if _request_body.get("stream"): - return True + return is_passthrough_request_streaming(_request_body) return False @@ -1755,9 +1759,11 @@ async def _base_vertex_proxy_route( ## check for streaming target = str(updated_url) - is_streaming_request = False - if "stream" in str(updated_url): - is_streaming_request = True + if ":rawPredict" in target or ":streamRawPredict" in target: + is_streaming_request = await is_streaming_request_fn(request) + else: + is_streaming_request = "stream" in target + if is_streaming_request: target += "?alt=sse" ## CREATE PASS-THROUGH diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index cf3351c4ff8..181846fe289 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -3022,3 +3022,172 @@ class TestCursorProxyRoute: assert call_args["target"] == "https://api.cursor.com/v0/agents" assert result["id"] == "bc_abc123" assert result["status"] == "CREATING" + + +class TestVertexRawPredictStreamingClassification: + """ + Regression coverage for LIT-4761. + + `_base_vertex_proxy_route` classified any target URL containing "stream" as a + streaming request. `:streamRawPredict` carries that substring, so a unary + Anthropic-on-Vertex call (no `stream` field in the body) was sent with + `?alt=sse` and logged through the streaming chunk collector, which parses + Anthropic SSE deltas and finds no usage in a complete `"type": "message"` + body; the spend log recorded 0 tokens and $0 cost. + + Streaming for the rawPredict family is decided by the request body, per the + Anthropic Messages contract. The Gemini generateContent family stays + URL-signalled because the Gemini REST body has no `stream` field. + """ + + RAW_PREDICT_ENDPOINT = ( + "v1/projects/test-project/locations/us-east5/publishers/anthropic/models/" + "claude-sonnet-4-6:streamRawPredict" + ) + GENERATE_CONTENT_ENDPOINT = ( + "v1/projects/test-project/locations/us-east5/publishers/google/models/" + "gemini-2.5-flash:streamGenerateContent" + ) + + async def _capture_passthrough_kwargs(self, endpoint: str, body: object) -> dict: + raw_body = json.dumps(body).encode("utf-8") + + async def receive(): + return {"type": "http.request", "body": raw_body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": f"/vertex_ai/{endpoint}", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + }, + receive=receive, + ) + + captured: dict = {} + + def fake_create_pass_through_route(**kwargs): + captured.update(kwargs) + return AsyncMock(return_value={"status": "success"}) + + mock_credentials = Mock() + mock_credentials.token = "test-token" + + base_url = "https://us-east5-aiplatform.googleapis.com/" + mock_handler = Mock() + mock_handler.get_default_base_target_url.return_value = base_url + mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url) + + module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + with ( + mock.patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth", + return_value=(mock_credentials, "test-project"), + ), + mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), + mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"), + mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value={"api_key": "test-key"})), + mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), + ): + await vertex_proxy_route( + endpoint=endpoint, + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(token="test-key"), + ) + + assert captured, "create_pass_through_route was never called" + return captured + + @pytest.mark.asyncio + async def test_raw_predict_without_stream_field_is_not_streaming(self): + captured = await self._capture_passthrough_kwargs( + endpoint=self.RAW_PREDICT_ENDPOINT, + body={ + "anthropic_version": "vertex-2023-10-16", + "messages": [{"role": "user", "content": "Explain MLOps"}], + "max_tokens": 5000, + }, + ) + + assert captured["is_streaming_request"] is False + assert "alt=sse" not in captured["target"] + + @pytest.mark.asyncio + async def test_raw_predict_with_stream_false_is_not_streaming(self): + captured = await self._capture_passthrough_kwargs( + endpoint=self.RAW_PREDICT_ENDPOINT, + body={ + "anthropic_version": "vertex-2023-10-16", + "stream": False, + "messages": [{"role": "user", "content": "Explain MLOps"}], + "max_tokens": 5000, + }, + ) + + assert captured["is_streaming_request"] is False + assert "alt=sse" not in captured["target"] + + @pytest.mark.asyncio + async def test_raw_predict_with_stream_true_still_streams(self): + captured = await self._capture_passthrough_kwargs( + endpoint=self.RAW_PREDICT_ENDPOINT, + body={ + "anthropic_version": "vertex-2023-10-16", + "stream": True, + "messages": [{"role": "user", "content": "Explain MLOps"}], + "max_tokens": 5000, + }, + ) + + assert captured["is_streaming_request"] is True + assert captured["target"].endswith("?alt=sse") + + @pytest.mark.asyncio + async def test_gemini_stream_generate_content_stays_url_signalled(self): + captured = await self._capture_passthrough_kwargs( + endpoint=self.GENERATE_CONTENT_ENDPOINT, + body={"contents": [{"role": "user", "parts": [{"text": "Explain MLOps"}]}]}, + ) + + assert captured["is_streaming_request"] is True + assert captured["target"].endswith("?alt=sse") + + @pytest.mark.asyncio + async def test_raw_predict_with_non_object_body_is_not_streaming(self): + captured = await self._capture_passthrough_kwargs( + endpoint=self.RAW_PREDICT_ENDPOINT, + body=[{"role": "user", "content": "Explain MLOps"}], + ) + + assert captured["is_streaming_request"] is False + assert "alt=sse" not in captured["target"] + + +@pytest.mark.parametrize( + "request_body, expected", + [ + ({"stream": True}, True), + ({"stream": "true"}, True), + ({"stream": False}, False), + ({}, False), + ([{"role": "user"}], False), + ([], False), + ("stream", False), + (7, False), + (None, False), + ], +) +def test_is_passthrough_request_streaming_tolerates_non_object_bodies(request_body, expected): + """ + A JSON request body is not required to be an object. Every passthrough + streaming decision funnels through this predicate, so a list or scalar body + must answer False instead of raising AttributeError. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + is_passthrough_request_streaming, + ) + + assert is_passthrough_request_streaming(request_body) is expected