From 018ccff23f1df9e4d223741e0f084447f08bcade Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:11:20 +0530
Subject: [PATCH] fix(vertex-ai): address greptile review feedback on batch
 cancel

- Add try/except httpx.HTTPStatusError blocks in _async_cancel_batch for
  both POST cancel and GET retrieve calls, with verbose_logger error logging
- Fix endpoint extraction inconsistency: compute endpoint from URL without
  :cancel suffix so it matches behaviour of create_batch/retrieve_batch
- Add explicit validation that api_base ends with ':cancel' before
  stripping it, raising a descriptive error for unsupported custom proxy
  URL rewriting scenarios
- Use string-based patch() in test instead of patch.object() for robustness
  against import order changes

Made-with: Cursor
---
 litellm/llms/vertex_ai/batches/handler.py     | 54 ++++++++++++++-----
 .../test_vertex_ai_batch_transformation.py    |  2 +-
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index c7b9287c08..36728499b9 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -399,10 +399,12 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_project=vertex_project or project_id,
         )
 
-        default_api_base = f"{default_api_base}/{batch_id}:cancel"
+        # Compute endpoint from the URL without :cancel for consistency with other methods
+        base_without_cancel = f"{default_api_base}/{batch_id}"
+        default_api_base = f"{base_without_cancel}:cancel"
 
-        if len(default_api_base.split(":")) > 1:
-            endpoint = default_api_base.split(":")[-1]
+        if len(base_without_cancel.split(":")) > 1:
+            endpoint = base_without_cancel.split(":")[-1]
         else:
             endpoint = ""
 
@@ -420,6 +422,14 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
+        if not api_base.endswith(":cancel"):
+            raise ValueError(
+                f"cancel_batch: expected api_base to end with ':cancel', got: {api_base!r}. "
+                "Custom proxy URL rewriting is not supported for this operation."
+            )
+
+        retrieve_api_base = api_base.rsplit(":cancel", 1)[0]
+
         headers = {
             "Content-Type": "application/json; charset=utf-8",
             "Authorization": f"Bearer {access_token}",
@@ -428,7 +438,7 @@ class VertexAIBatchPrediction(VertexLLM):
         if _is_async is True:
             return self._async_cancel_batch(
                 api_base=api_base,
-                retrieve_api_base=api_base.rsplit(":cancel", 1)[0],
+                retrieve_api_base=retrieve_api_base,
                 headers=headers,
             )
 
@@ -443,7 +453,7 @@ class VertexAIBatchPrediction(VertexLLM):
             raise Exception(f"Error: {response.status_code} {response.text}")
 
         retrieve_response = sync_handler.get(
-            url=api_base.rsplit(":cancel", 1)[0],
+            url=retrieve_api_base,
             headers=headers,
         )
         if retrieve_response.status_code != 200:
@@ -466,18 +476,34 @@ class VertexAIBatchPrediction(VertexLLM):
         client = get_async_httpx_client(
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
-        response = await client.post(
-            url=api_base,
-            headers=headers,
-            data=json.dumps({}),
-        )
+        try:
+            response = await client.post(
+                url=api_base,
+                headers=headers,
+                data=json.dumps({}),
+            )
+        except httpx.HTTPStatusError as e:
+            litellm.verbose_logger.error(
+                "Vertex AI batch cancel failed: status=%s, body=%s",
+                e.response.status_code,
+                e.response.text[:1000],
+            )
+            raise
         if response.status_code != 200:
             raise Exception(f"Error: {response.status_code} {response.text}")
 
-        retrieve_response = await client.get(
-            url=retrieve_api_base,
-            headers=headers,
-        )
+        try:
+            retrieve_response = await client.get(
+                url=retrieve_api_base,
+                headers=headers,
+            )
+        except httpx.HTTPStatusError as e:
+            litellm.verbose_logger.error(
+                "Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
+                e.response.status_code,
+                e.response.text[:1000],
+            )
+            raise
         if retrieve_response.status_code != 200:
             raise Exception(
                 f"Error: {retrieve_response.status_code} {retrieve_response.text}"
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 d27cd8ba8a..e03555a9c0 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
@@ -101,7 +101,7 @@ async def test_litellm_cancel_batch_vertex_ai():
     mock_response.id = "batch_123"
     mock_response.status = "cancelling"
     
-    with patch.object(litellm.batches.main, "vertex_ai_batches_instance") as mock_instance:
+    with patch("litellm.batches.main.vertex_ai_batches_instance") as mock_instance:
         mock_instance.cancel_batch.return_value = mock_response
         
         response = litellm.cancel_batch(