greptile comments

litellm/batches/main.py

@@ -1048,7 +1048,7 @@ def cancel_batch(
                 litellm_params=litellm_params,
             )
         elif custom_llm_provider == "vertex_ai":
-            api_base = optional_params.api_base or ""
+            api_base = optional_params.api_base or None
             vertex_ai_project = (
                 optional_params.vertex_project
                 or litellm.vertex_project

litellm/llms/vertex_ai/batches/handler.py

@@ -400,28 +400,41 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         retrieve_api_base_default = f"{default_api_base}/{batch_id}"
-        default_api_base = f"{retrieve_api_base_default}:cancel"
+        cancel_api_base_default = f"{retrieve_api_base_default}:cancel"
 
-        # The Vertex AI action suffix for this operation
-        endpoint = "cancel"
+        # Save the caller-supplied value before _check_custom_proxy overwrites api_base,
+        # so we can pass it unchanged to the second proxy-check for the retrieve URL.
+        caller_api_base = api_base
 
         _, api_base = self._check_custom_proxy(
-            api_base=api_base,
+            api_base=caller_api_base,
             custom_llm_provider="vertex_ai",
             gemini_api_key=None,
-            endpoint=endpoint,
+            endpoint="cancel",
             stream=None,
             auth_header=None,
-            url=default_api_base,
+            url=cancel_api_base_default,
             model=None,
             vertex_project=vertex_project or project_id,
             vertex_location=vertex_location or "us-central1",
             vertex_api_version="v1",
         )
 
-        # Use the canonical retrieve URL built from components rather than stripping
-        # ":cancel" from api_base, so custom proxy URL rewriting does not break retrieval.
-        retrieve_api_base = retrieve_api_base_default
+        # Route the retrieve GET through the same proxy as the cancel POST by running
+        # _check_custom_proxy a second time with the non-cancel default URL.
+        _, retrieve_api_base = self._check_custom_proxy(
+            api_base=caller_api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="",
+            stream=None,
+            auth_header=None,
+            url=retrieve_api_base_default,
+            model=None,
+            vertex_project=vertex_project or project_id,
+            vertex_location=vertex_location or "us-central1",
+            vertex_api_version="v1",
+        )
 
         headers = {
             "Content-Type": "application/json; charset=utf-8",

tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py

@@ -45,7 +45,7 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl()
 
 
 @pytest.mark.asyncio
-async def test_vertex_ai_cancel_batch():
+def test_vertex_ai_cancel_batch():
     """Test that vertex_ai cancel_batch calls the correct API endpoint"""
     handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")