From dd6ef9e1bc22f4a07051e143af4e5fe266c40067 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 17 Sep 2026 10:37:08 -0700
Subject: [PATCH] fix(e2e): delete raw cloud-storage batch files with the
 master key

DELETE /v1/files/{id} only lets a proxy admin key delete a raw s3:// or
gs:// file id, because such ids skip the managed-file owner check. The
batch lifecycle cleanup deleted the vertex_ai raw ids with the test's
own virtual key and got a 403 at teardown on every build since #194

Raw cloud-storage ids now go through the master key; managed and
provider-native ids keep using the creating key
---
 tests/e2e/batches/batch_cleanup.py      | 11 +++++++++--
 tests/e2e/batches/batch_client.py       |  8 ++++++++
 tests/e2e/batches/capabilities.py       |  7 +++++++
 tests/e2e/batches/test_batch_cleanup.py |  4 ++++
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py
index 9284882ad82..86e47c0b1e1 100644
--- a/tests/e2e/batches/batch_cleanup.py
+++ b/tests/e2e/batches/batch_cleanup.py
@@ -5,7 +5,7 @@ from time import monotonic, sleep
 from typing import Final, Protocol
 
 from batch_client import BatchObject, FileDeleteResponse
-from capabilities import is_managed_id
+from capabilities import is_cloud_storage_id, is_managed_id
 from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
 from pydantic import BaseModel
 
@@ -19,6 +19,8 @@ BATCH_CANCEL_POLL_SECONDS: Final = 10.0
 class BatchCleanupClient(Protocol):
     def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ...
 
+    def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]: ...
+
     def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
 
     def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
@@ -49,7 +51,12 @@ def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) ->
 
 
 def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None:
-    result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider))
+    delete: Final[Callable[[], Result[FileDeleteResponse]]] = (
+        (lambda: client.delete_file_as_admin(file_id, provider=provider))
+        if is_cloud_storage_id(file_id)
+        else (lambda: client.delete_file(file_id, key=key, provider=provider))
+    )
+    result: Final = cleanup_result(delete)
     if isinstance(result, UnknownApiError) and result.status_code == 404:
         return
     deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}")
diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py
index c9c77e1f12e..8745140a818 100644
--- a/tests/e2e/batches/batch_client.py
+++ b/tests/e2e/batches/batch_client.py
@@ -233,6 +233,14 @@ class BatchClient:
             response_type=FileDeleteResponse,
         )
 
+    def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]:
+        return self.proxy.transport.delete(
+            f"{_files_path(provider)}/{file_id}",
+            headers=self.proxy.transport.master,
+            json=NoBody(),
+            response_type=FileDeleteResponse,
+        )
+
 
 def _files_path(provider: str | None) -> str:
     return f"/{provider}/v1/files" if provider else "/v1/files"
diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py
index 17749c2fb87..d510426dee2 100644
--- a/tests/e2e/batches/capabilities.py
+++ b/tests/e2e/batches/capabilities.py
@@ -222,6 +222,13 @@ def is_managed_id(id_str: str) -> bool:
     return _b64_decode(id_str).startswith("litellm_proxy")
 
 
+CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
+
+
+def is_cloud_storage_id(id_str: str) -> bool:
+    return id_str.startswith(CLOUD_STORAGE_SCHEMES)
+
+
 def is_model_encoded_id(id_str: str) -> bool:
     for prefix in ("file-", "batch_"):
         if id_str.startswith(prefix):
diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py
index d0038139dcf..a0932a80dfe 100644
--- a/tests/e2e/batches/test_batch_cleanup.py
+++ b/tests/e2e/batches/test_batch_cleanup.py
@@ -45,6 +45,10 @@ class CleanupClient:
         self.calls(f"delete {provider} {file_id}")
         return self.file_response()
 
+    def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]:
+        self.calls(f"admin delete {provider} {file_id}")
+        return self.file_response()
+
     def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
         self.calls(f"retrieve {provider} {batch_id}")
         return self.batch_response()