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
This commit is contained in:
Yuneng Jiang 2026-09-17 10:37:08 -07:00
parent 1d71564063
commit dd6ef9e1bc
No known key found for this signature in database
4 changed files with 28 additions and 2 deletions

View file

@ -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}")

View file

@ -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"

View file

@ -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):

View file

@ -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()