From d7d186d1f7ed5ac1e4256af58b07ba4955ff8043 Mon Sep 17 00:00:00 2001 From: apex-mochen <2756823972@qq.com> Date: Sat, 10 Oct 2026 11:30:00 +0000 Subject: [PATCH] fix(proxy): harden Vertex pass-through route matching against substring false positives (#45787) A plain pass-through endpoint whose target URL contains the substring "predict", "search", "embedContent", "rawPredict", "streamGenerateContent" or "predictLongRunning" could be misrouted to the Vertex AI logging handler. Inside vertex_passthrough_handler the unguarded response["predictions"] access then raised KeyError in the logging worker and silently dropped the spend row (issue #45787). - vertex_passthrough_logging_handler.py: match Vertex method routes by the canonical ":method" suffix (mirroring PassThroughEndpointLogging.is_vertex_route) instead of bare substring matching; the batchPredictionJobs resource route is path-based and kept as-is. - vertex_embeddings/transformation.py: raise a clear ValueError when the embedding response lacks the "predictions" field instead of a raw KeyError surfacing deep in the logging worker. - test_success_handler.py: regression tests for the issue URLs (negative) and for authentic Vertex ":method" URLs (positive). AI-assisted disclosure: root cause located and patch written with the help of an AI assistant; all tests and lint verified locally. Signed-off-by: apex-mochen <2756823972@qq.com> --- .../vertex_embeddings/transformation.py | 6 ++ .../vertex_passthrough_logging_handler.py | 18 ++--- .../test_success_handler.py | 65 +++++++++++++++++++ 3 files changed, 80 insertions(+), 9 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index cab2bed8a60..dd3dbfa0319 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -229,6 +229,9 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_response(response=response, model=model, model_response=model_response) + if not isinstance(response, dict) or "predictions" not in response: + raise ValueError("Vertex embedding response missing 'predictions' field") + _predictions: Final = response["predictions"] embedding_response: Final = [] @@ -257,6 +260,9 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): """ Transforms a vertex fine-tuned model embedding response to an openai response format. """ + if not isinstance(response, dict) or "predictions" not in response: + raise ValueError("Vertex embedding response missing 'predictions' field") + _predictions: Final = response["predictions"] embedding_response: Final = [] diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index cfa569b5219..90b8d831e3f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -148,7 +148,7 @@ class VertexPassthroughLoggingHandler: custom_llm_provider="vertex_ai", vertex_location=vertex_location, ) - if "predictLongRunning" in url_route: + if ":predictLongRunning" in url_route: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) vertex_video_config: Final = VertexAIVideoConfig() @@ -188,7 +188,7 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } - elif "generateContent" in url_route: + elif ":generateContent" in url_route or ":streamGenerateContent" in url_route: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) instance_of_vertex_llm: Final = litellm.VertexGeminiConfig() @@ -220,7 +220,7 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } - elif "embedContent" in url_route or "batchEmbedContents" in url_route: + elif ":embedContent" in url_route or ":batchEmbedContents" in url_route: return VertexPassthroughLoggingHandler._handle_embed_content_response( httpx_response=httpx_response, logging_obj=logging_obj, @@ -228,14 +228,14 @@ class VertexPassthroughLoggingHandler: kwargs=kwargs, request_body=request_body, ) - elif "predict" in url_route: + elif ":predict" in url_route: return VertexPassthroughLoggingHandler._handle_predict_response( httpx_response=httpx_response, logging_obj=logging_obj, url_route=url_route, kwargs=kwargs, ) - elif "rawPredict" in url_route or "streamRawPredict" in url_route: + elif ":rawPredict" in url_route or ":streamRawPredict" in url_route: from litellm.llms.vertex_ai.vertex_ai_partner_models import ( get_vertex_ai_partner_model_config, ) @@ -287,7 +287,7 @@ class VertexPassthroughLoggingHandler: "result": litellm_prediction_response, "kwargs": kwargs, } - elif "search" in url_route: + elif ":search" in url_route: litellm_vs_response: Final = vertex_search_api_config.transform_search_vector_store_response( response=httpx_response, litellm_logging_obj=logging_obj, @@ -513,7 +513,7 @@ class VertexPassthroughLoggingHandler: model: Final = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) response_json: Final = httpx_response.json() - is_batch: Final = "batchEmbedContents" in url_route + is_batch: Final = ":batchEmbedContents" in url_route input_text: Final = VertexPassthroughLoggingHandler._extract_embed_content_input( request_body=request_body, batch=is_batch @@ -626,7 +626,7 @@ class VertexPassthroughLoggingHandler: url_route: str, ) -> ModelResponse | TextCompletionResponse | None: parsed_chunks = [] - if "generateContent" in url_route or "streamGenerateContent" in url_route: + if ":generateContent" in url_route or ":streamGenerateContent" in url_route: vertex_iterator: Any = VertexModelResponseIterator( streaming_response=None, sync_stream=False, @@ -634,7 +634,7 @@ class VertexPassthroughLoggingHandler: ) chunk_parsing_logic: Callable[..., ModelResponseStream | None] = vertex_iterator._common_chunk_parsing_logic parsed_chunks = [chunk_parsing_logic(chunk) for chunk in all_chunks] - elif "rawPredict" in url_route or "streamRawPredict" in url_route: + elif ":rawPredict" in url_route or ":streamRawPredict" in url_route: from litellm.llms.anthropic.chat.handler import ModelResponseIterator from litellm.llms.base_llm.base_model_iterator import ( BaseModelResponseIterator, diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index 9cd27a79fe0..ab0c37e037c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -584,3 +584,68 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response( assert normalized["standard_logging_response_object"] is None assert normalized["kwargs"] == {"model": _LIVE_MODEL} + + +# --- regression for #45787: substring route matching must not classify plain +# passthrough URLs (e.g. "http://example.com/api/predict/laya") as Vertex AI --- + +_NON_VERTEX_URLS_WITH_METHOD_SUBSTRINGS: Final = [ + "http://example.com/api/predict/laya", + "http://example.com/api/search/items", + "http://example.com/api/embedContent/proxy", + "http://example.com/v1/rawPredict", + "http://example.com/predictLongRunning/jobs/1", + "http://example.com/api/streamGenerateContent/relay", +] + + +@pytest.mark.parametrize("url", _NON_VERTEX_URLS_WITH_METHOD_SUBSTRINGS) +def test_is_vertex_route_ignores_non_vertex_urls_containing_method_substrings(url: str) -> None: + assert PassThroughEndpointLogging().is_vertex_route(url) is False + + +def test_is_vertex_route_matches_authentic_vertex_method_urls() -> None: + handler = PassThroughEndpointLogging() + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/" + "publishers/google/models/gemini-2.0-flash:predict" + ) + is True + ) + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/" + "publishers/google/models/gemini-2.0:search" + ) + is True + ) + assert ( + handler.is_vertex_route( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/" + "publishers/google/models/text-embedding-005:embedContent" + ) + is True + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_non_vertex_predict_url_normalizes_without_predictions_keyerror() -> None: + # Issue #45787: a plain passthrough URL containing the substring "predict" + # used to be misrouted to the Vertex handler, whose unguarded + # response["predictions"] access then raised KeyError inside the logging + # worker and silently dropped the spend row. + normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=httpx.Response(200, request=httpx.Request("GET", "http://example.com/api/predict/laya")), + response_body={}, # no "predictions" key + request_body={}, + logging_obj=_logging_obj(), + url_route="http://example.com/api/predict/laya", + result="", + start_time=_START, + end_time=_END, + cache_hit=False, + custom_llm_provider=None, + ) + + assert normalized["standard_logging_response_object"] is not None