From 74ae17d15305334d0614b107e866bb5f6f7ef0e9 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:41:46 +0530
Subject: [PATCH]  greptile comments

---
 litellm/batches/main.py                       |  2 +-
 litellm/llms/vertex_ai/batches/handler.py     | 31 +++++++++++++------
 .../test_vertex_ai_batch_transformation.py    |  2 +-
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index 36093d071ba..ae79469dd10 100644
--- a/litellm/batches/main.py
+++ b/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
diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index b30130688d2..f4ba0533c81 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/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",
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index e03555a9c09..1cf6fa32661 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/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")