fix(vertex_ai): pin custom-endpoint batch jobs to the deployment's own endpoint

The endpoint id in a custom-endpoints/ file path is caller-controlled (raw
gs:// file ids are accepted at batch create), so it must match the endpoint
the routed deployment's api_base names, and container resolution only runs
for deployments marked custom_endpoint; otherwise an attacker-supplied file
id could run another endpoint's container with the deployment's project
credentials.

litellm/llms/vertex_ai/batches/handler.py

@@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.llms.vertex_ai.common_utils import (
     VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,
     VertexAIError,
+    get_custom_endpoint_id_from_api_base,
     get_vertex_base_url,
 )
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
@@ -156,14 +157,29 @@ class VertexAIBatchPrediction(VertexLLM):
                 vertex_location=vertex_location or "us-central1",
             )
         )
-        if custom_endpoint and "/custom-endpoints/" not in transformed_batch_request.get("model", ""):
+        # The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file
+        # ids are accepted), so it must never pick which container runs: batch jobs execute the
+        # container with the deployment's project credentials. The id has to match the endpoint
+        # the routed deployment's own api_base names, and container resolution only happens for
+        # deployments the admin marked custom_endpoint.
+        job_model: Final = transformed_batch_request.get("model", "")
+        deployment_endpoint_id: Final = get_custom_endpoint_id_from_api_base(api_base)
+        if custom_endpoint and not job_model.endswith(f"/custom-endpoints/{deployment_endpoint_id}"):
             raise VertexAIError(
                 status_code=400,
                 message=(
                     "Vertex AI batch prediction on a `custom_endpoint` deployment requires an input "
-                    "file uploaded through LiteLLM against that deployment (its file id carries a "
-                    "custom-endpoints/<endpoint id> path); this input file targets a publisher or "
-                    "fine-tuned Gemini model instead."
+                    "file uploaded through LiteLLM against that same deployment: the file id's "
+                    "custom-endpoints/<endpoint id> path must name the endpoint the deployment's "
+                    "`api_base` serves from."
+                ),
+            )
+        if not custom_endpoint and "/custom-endpoints/" in job_model:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "This input file was uploaded for a `custom_endpoint` deployment; create the "
+                    "batch against that deployment instead."
                 ),
             )
         gateway_api_base: Final = _gateway_api_base_or_none(api_base)
@@ -174,12 +190,16 @@ class VertexAIBatchPrediction(VertexLLM):
             api_base=gateway_api_base,
             vertex_location=vertex_location or "us-central1",
         )
-        vertex_batch_request: Final = self._resolve_custom_endpoint_container(
-            vertex_batch_request=resolved_batch_request,
-            headers=headers,
-            sync_handler=sync_handler,
-            api_base=gateway_api_base,
-            vertex_location=vertex_location or "us-central1",
+        vertex_batch_request: Final = (
+            self._resolve_custom_endpoint_container(
+                vertex_batch_request=resolved_batch_request,
+                headers=headers,
+                sync_handler=sync_handler,
+                api_base=gateway_api_base,
+                vertex_location=vertex_location or "us-central1",
+            )
+            if custom_endpoint
+            else resolved_batch_request
         )
         is_unmanaged_container_job: Final = "unmanagedContainerModel" in vertex_batch_request
 

litellm/llms/vertex_ai/common_utils.py

@@ -5,6 +5,7 @@ from enum import Enum
 from functools import lru_cache
 from types import MappingProxyType
 from typing import Any, Final, Literal, cast, get_type_hints
+from urllib.parse import urlparse
 
 import httpx
 from pydantic import TypeAdapter, ValidationError
@@ -373,6 +374,23 @@ def get_vertex_base_model_name(model: str) -> str:
 VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: Final = "litellm_custom_id"
 
 
+def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None:
+    """
+    The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its
+    api_base (`.../endpoints/<id>:rawPredict` or a dedicated-domain equivalent); batch jobs need
+    that id to read the endpoint's containerSpec, so extract it (verb suffix stripped).
+    """
+    if not api_base:
+        return None
+    path_segments: Final = urlparse(api_base).path.split("/")
+    after_endpoints: Final = tuple(
+        segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints"
+    )
+    if not after_endpoints:
+        return None
+    return after_endpoints[-1].split(":")[0] or None
+
+
 def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None:
     """
     Fine-tuned Gemini deployments are addressed by a numeric endpoint id,

litellm/llms/vertex_ai/files/transformation.py

@@ -8,7 +8,7 @@ import time
 from collections.abc import Callable, Iterable, Iterator, Mapping
 from types import MappingProxyType
 from typing import Any, Final, TypedDict
-from urllib.parse import quote, unquote, urlparse
+from urllib.parse import quote, unquote
 
 import httpx
 from httpx import Headers, Response
@@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import (
 from litellm.llms.vertex_ai.common_utils import (
     VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,
     _convert_vertex_datetime_to_openai_datetime,
+    get_custom_endpoint_id_from_api_base,
     get_vertex_ai_fine_tuned_endpoint_id,
 )
 from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
@@ -713,23 +714,6 @@ _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions"
 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
 
 
-def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None:
-    """
-    The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its
-    api_base (`.../endpoints/<id>:rawPredict` or a dedicated-domain equivalent); batch jobs need
-    that id to read the endpoint's containerSpec, so extract it (verb suffix stripped).
-    """
-    if not api_base:
-        return None
-    path_segments: Final = urlparse(api_base).path.split("/")
-    after_endpoints: Final = tuple(
-        segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints"
-    )
-    if not after_endpoints:
-        return None
-    return after_endpoints[-1].split(":")[0] or None
-
-
 def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: Mapping[str, object]) -> Mapping[str, object]:
     """
     One OpenAI batch JSONL line as the instance a vLLM-serving Vertex container consumes:

tests/test_litellm/llms/vertex_ai/batches/test_handler.py

@@ -269,6 +269,10 @@ CONTAINER_SPEC = {
     "healthRoute": "/ping",
 }
 MACHINE_SPEC = {"machineType": "g2-standard-12", "acceleratorType": "NVIDIA_L4", "acceleratorCount": 1}
+CUSTOM_ENDPOINT_API_BASE = (
+    f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}"
+    f"/locations/{LOCATION}/endpoints/{CUSTOM_ENDPOINT_ID}:rawPredict"
+)
 
 
 def _custom_endpoint_get_response() -> MagicMock:
@@ -317,7 +321,7 @@ def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job():
         out = h.create_batch(
             _is_async=False,
             create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA,
-            api_base=None,
+            api_base=CUSTOM_ENDPOINT_API_BASE,
             vertex_credentials=None,
             vertex_project=PROJECT,
             vertex_location=LOCATION,
@@ -360,7 +364,7 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400():
             h.create_batch(
                 _is_async=False,
                 create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA,
-                api_base=None,
+                api_base=CUSTOM_ENDPOINT_API_BASE,
                 vertex_credentials=None,
                 vertex_project=PROJECT,
                 vertex_location=LOCATION,
@@ -374,6 +378,63 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400():
     client.post.assert_not_called()
 
 
+def test_create_batch_sync_custom_endpoint_rejects_file_for_other_endpoint():
+    """The endpoint id in the file path is caller-controlled (raw gs:// ids are accepted), so a
+    file staged for a different endpoint must not make the deployment run that endpoint's
+    container with its own project credentials."""
+    h = _make_handler()
+    client = MagicMock()
+    foreign_file = {"input_file_id": "gs://bucket/litellm-vertex-files/custom-endpoints/999999/file-uuid"}
+
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get") as safe_get,
+    ):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=foreign_file,
+                api_base=CUSTOM_ENDPOINT_API_BASE,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+                custom_endpoint=True,
+            )
+
+    assert exc_info.value.status_code == 400
+    safe_get.assert_not_called()
+    client.post.assert_not_called()
+
+
+def test_create_batch_sync_non_custom_deployment_rejects_custom_endpoint_file():
+    """A custom-endpoints file id sent through an ordinary Vertex deployment must not trigger
+    container resolution at all."""
+    h = _make_handler()
+    client = MagicMock()
+
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get") as safe_get,
+    ):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA,
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+            )
+
+    assert exc_info.value.status_code == 400
+    safe_get.assert_not_called()
+    client.post.assert_not_called()
+
+
 def test_create_batch_sync_ignores_resource_shaped_api_base():
     """A deployment api_base like `.../endpoints/<id>:rawPredict` targets online inference, not
     the Vertex API root; grafting batch urls onto it yields guaranteed 404s, so batch operations