From 6c21be619441dbd24879e1f8a4b897c6dbd667fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:15:06 -0700 Subject: [PATCH 01/13] fix(e2e): clean up batch files reliably and expire Azure inputs --- tests/e2e/batches/COVERAGE.md | 18 ++ tests/e2e/batches/batch_cleanup.py | 90 +++++++++ tests/e2e/batches/batch_client.py | 16 +- tests/e2e/batches/capabilities.py | 4 + tests/e2e/batches/conftest.py | 10 +- tests/e2e/batches/test_batch_cleanup.py | 187 ++++++++++++++++++ tests/e2e/batches/test_batches_e2e.py | 80 ++++---- .../test_managed_files_enforcement_e2e.py | 3 +- tests/e2e/lifecycle.py | 23 ++- 9 files changed, 381 insertions(+), 50 deletions(-) create mode 100644 tests/e2e/batches/batch_cleanup.py create mode 100644 tests/e2e/batches/test_batch_cleanup.py diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 8a7b68511ec..899530dde2b 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -120,6 +120,24 @@ create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: nested managed ids round-trip retrieve. This self-chaining only needs the proxy to reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. +## Cleanup + +Batch teardown cancels active batches before deleting their input files and keys. +Raw file IDs from both `model_param` and `provider_fallback` uploads use the upload +provider when deleted. Model-encoded and managed file IDs route themselves + +File deletion and batch cancellation check their responses and retry transient +failures up to three times. Teardown attempts every registered cleanup before +reporting failures as test errors. Already deleted files and batches that are +terminal are safe to clean up again. Cancellation polls for up to ten minutes +before input deletion, because accepting cancellation does not finish it + +Azure input uploads request `expires_after` anchored to `created_at` with +`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a +fallback for interrupted runs: immediate deletion remains the normal cleanup. +Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot +be requested through its Files API + ## Terminal state + cost write-back (cross-run marker baton) The 24h completion window rules out submit-and-wait inside one run, so diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py new file mode 100644 index 00000000000..a5b5e9bba37 --- /dev/null +++ b/tests/e2e/batches/batch_cleanup.py @@ -0,0 +1,90 @@ +from collections.abc import Callable +from time import monotonic, sleep +from typing import Final, Protocol + +from pydantic import BaseModel + +from batch_client import BatchObject, FileDeleteResponse +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError + +CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) +BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 600.0 +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 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]: ... + + +def cleanup_result[R: BaseModel]( + action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep +) -> Result[R]: + for delay, result in ((delay, action()) for delay in CLEANUP_DELAYS): + match result: + case NetworkError() | RateLimitedError(): + wait(delay) + case UnknownApiError(status_code=code) if code in {408, 429, 500, 502, 503, 504}: + wait(delay) + case _: + return result + return action() + + +def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) -> R: + match result: + case Success(data=data): + return data + case UnknownApiError(status_code=code): + raise AssertionError(f"{operation} failed: HTTP {code}") + case _: + raise AssertionError(f"{operation} failed: {result.kind}") + + +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)) + if isinstance(result, UnknownApiError) and result.status_code == 404: + return + deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") + assert deleted.deleted, f"Delete file {file_id} did not confirm deletion" + + +def cleanup_batch( + client: BatchCleanupClient, + batch_id: str, + *, + key: str, + provider: str | None = None, + wait: Callable[[float], None] = sleep, + clock: Callable[[], float] = monotonic, +) -> None: + fetched: Final = _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} for cleanup", + ) + if fetched.status in BATCH_TERMINAL_STATUSES: + return + if fetched.status != "cancelling": + result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) + if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): + cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") + assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( + f"Cancel batch {batch_id} left status {cancelled.status}" + ) + deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS + while True: + current = _require_cleanup_success( + cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), + f"Retrieve batch {batch_id} after cancellation", + ) + if current.status in BATCH_TERMINAL_STATUSES: + return + assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" + assert clock() < deadline, ( + f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" + ) + wait(BATCH_CANCEL_POLL_SECONDS) diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 31e49f22450..84a902b0b11 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -13,8 +13,9 @@ co-located here because only this suite uses them. from __future__ import annotations from dataclasses import dataclass +from typing import Final, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field from proxy_client import ProxyClient from e2e_http import ( @@ -27,6 +28,18 @@ from e2e_http import ( from models import LiteLLMParamsBody UPLOAD_FILENAME = "batch_input.jsonl" +AZURE_FILE_EXPIRY_SECONDS: Final = 14 * 24 * 60 * 60 + + +class ExpiringFileUploadForm(FileUploadForm): + expires_after_anchor: Literal["created_at"] = Field(default="created_at", alias="expires_after[anchor]") + expires_after_seconds: int = Field(default=AZURE_FILE_EXPIRY_SECONDS, alias="expires_after[seconds]") + + +def batch_upload_form(provider: str, *, target_model_names: str | None = None) -> FileUploadForm: + if provider == "azure": + return ExpiringFileUploadForm(target_model_names=target_model_names) + return FileUploadForm(target_model_names=target_model_names) class FileObject(BaseModel): @@ -37,6 +50,7 @@ class FileObject(BaseModel): bytes: int | None = None status: str | None = None created_at: int | None = None + expires_at: int | None = None class FileList(BaseModel): diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 1bcea0a61ee..17749c2fb87 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -108,6 +108,10 @@ class Capability: def id(self) -> str: return f"{self.provider}-{self.scenario}" + @property + def file_provider(self) -> str | None: + return self.provider if self.scenario in {"model_param", "provider_fallback"} else None + @property def jsonl_model(self) -> str: # Always the provider deployment name. Unified routes via diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 3b133fab680..91a365b6b92 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -13,7 +13,7 @@ the proxy config. from __future__ import annotations import os -from typing import Iterator +from typing import Final, Iterator import pytest @@ -21,6 +21,7 @@ from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_config import MANAGED_FILES_OPT_IN_ENV from e2e_http import NoBody +from lifecycle import ResourceManager from proxy_client import ProxyClient @@ -52,6 +53,13 @@ def client(proxy: ProxyClient) -> BatchClient: return build_client(proxy) +@pytest.fixture +def resources(client: BatchClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + yield manager + manager.teardown() + + @pytest.fixture(scope="session") def batch_deployments(client: BatchClient) -> Iterator[None]: probe = client.proxy.probe("/health/liveliness", params=NoBody()) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py new file mode 100644 index 00000000000..8ebfd750360 --- /dev/null +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -0,0 +1,187 @@ +from builtins import ExceptionGroup +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Final + +import pytest + +from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result +from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form +from capabilities import CAPABILITIES, Capability +from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from lifecycle import ResourceManager +from models import KeyGenerateBody + + +@dataclass +class CleanupClient: + files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) + batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + calls: list[str] = field(default_factory=list) + + def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: + self.calls.append(f"delete {provider} {file_id}") + return next(self.files) + + def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls.append(f"retrieve {provider} {batch_id}") + return next(self.batches) + + def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: + self.calls.append(f"cancel {provider} {batch_id}") + return next(self.cancellations) + + def generate_key(self, body: KeyGenerateBody) -> str: + return "test-key" + + def delete_key(self, key: str) -> None: + self.calls.append(f"delete key {key}") + + def delete_customers(self, user_ids: list[str]) -> None: + self.calls.append(f"delete customers {user_ids}") + + +def batch(status: str) -> Success[BatchObject]: + return Success(status_code=200, data=BatchObject(id="batch-1", status=status)) + + +def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]: + return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted)) + + +class TestFileCleanup: + @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) + def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: + client: Final = CleanupClient(files=iter((deleted_file(),))) + cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) + expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None + assert client.calls == [f"delete {expected_provider} file-1"] + + def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="secret response"),))) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure")) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + assert client.calls == ["delete azure file-1", "delete key test-key"] + assert len(caught.value.exceptions) == 1 + assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403" + + def test_success_response_must_confirm_deletion(self) -> None: + client: Final = CleanupClient(files=iter((deleted_file(deleted=False),))) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, "file-1", key="test-key") + + def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=404, body="missing"),))) + cleanup_file(client, "file-1", key="test-key", provider="azure") + assert client.calls == ["delete azure file-1"] + + def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: + client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="forbidden"),))) + manager: Final = ResourceManager(client=client) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.teardown() + assert client.calls == ["delete None file-1", "delete key test-key"] + + +class TestCleanupRetries: + @pytest.mark.parametrize( + "failure", + [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")], + ) + def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: + outcomes: Final = iter((failure, deleted_file())) + delays: Final[list[float]] = [] + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + assert isinstance(result, Success) and result.data.deleted + assert delays == [1.0] + + def test_persistent_error_has_bounded_retries(self) -> None: + failure: Final = UnknownApiError(status_code=503, body="unavailable") + outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) + delays: Final[list[float]] = [] + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + assert result is failure + assert tuple(delays) == CLEANUP_DELAYS + assert next(outcomes, None) is None + + def test_permanent_error_is_not_retried(self) -> None: + failure: Final = UnknownApiError(status_code=403, body="forbidden") + outcomes: Final = iter((failure, deleted_file())) + delays: Final[list[float]] = [] + assert cleanup_result(lambda: next(outcomes), wait=delays.append) is failure + assert delays == [] + assert isinstance(next(outcomes), Success) + + +class TestBatchCancellation: + def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: + client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) + delays: Final[list[float]] = [] + cleanup_batch(client, "batch-1", key="test-key", wait=delays.append) + assert client.calls == ["retrieve None batch-1"] * 3 + assert delays == [10.0] + + def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: + client: Final = CleanupClient( + batches=iter((batch("cancelling"), batch("cancelling"))), files=iter((deleted_file(),)) + ) + ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.defer(lambda: cleanup_batch(client, "batch-1", key=key, clock=lambda: next(ticks))) + with pytest.raises(ExceptionGroup) as caught: + manager.teardown() + assert "cancellation did not finish" in str(caught.value.exceptions[0]) + assert client.calls == [ + "retrieve None batch-1", + "retrieve None batch-1", + "delete None file-1", + "delete key test-key", + ] + + @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) + def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: + client: Final = CleanupClient(batches=iter((batch(status),))) + cleanup_batch(client, "batch-1", key="test-key") + assert client.calls == ["retrieve None batch-1"] + + def test_active_batch_is_cancelled_through_its_provider(self) -> None: + client: Final = CleanupClient( + batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) + ) + cleanup_batch(client, "batch-1", key="test-key", provider="azure") + assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1", "retrieve azure batch-1"] + + @pytest.mark.parametrize("status", ["completed", "in_progress"]) + def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: + client: Final = CleanupClient( + batches=iter((batch("in_progress"), batch(status))), + cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), + ) + if status == "completed": + cleanup_batch(client, "batch-1", key="test-key") + else: + with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"): + cleanup_batch(client, "batch-1", key="test-key") + assert client.calls == ["retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"] + + +class TestAzureFileExpiry: + def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None: + form: Final = batch_upload_form("azure", target_model_names="azure-test") + assert form.model_dump(by_alias=True, exclude_none=True) == { + "purpose": "batch", + "target_model_names": "azure-test", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS, + } + + @pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"]) + def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None: + assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"} diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index ed7cf656d01..1b4a6ed266f 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -21,14 +21,16 @@ import os import re import time from datetime import datetime, timedelta, timezone -from typing import Callable import pytest from pydantic import BaseModel from e2e_config import PROXY_BASE_URL, unique_marker +from batch_cleanup import cleanup_batch, cleanup_file from batch_client import ( + AZURE_FILE_EXPIRY_SECONDS, + batch_upload_form, UPLOAD_FILENAME, BatchClient, BatchCreateBody, @@ -155,19 +157,19 @@ def upload_for_scenario( if cap.scenario == "encoded": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), model=cap.model, key=key, ) if cap.scenario == "unified": return client.upload_file( content=content, - form=FileUploadForm(purpose="batch", target_model_names=cap.model), + form=batch_upload_form(cap.provider, target_model_names=cap.model), key=key, ) return client.upload_file( content=content, - form=FileUploadForm(purpose="batch"), + form=batch_upload_form(cap.provider), key=key, provider=cap.provider, ) @@ -188,20 +190,11 @@ def create_for_scenario( def op_provider(cap: Capability) -> str | None: - """provider_fallback ids are raw, so retrieve/cancel/list/delete need the provider + """provider_fallback batch ids are raw, so retrieve/cancel/list need the provider hint; the other scenarios encode it into the id and route automatically.""" return cap.provider if cap.scenario == "provider_fallback" else None -def quietly(action: Callable[[], object]) -> Callable[[], None]: - """Adapt a value-returning call into a best-effort cleanup the teardown can run.""" - - def run() -> None: - action() - - return run - - def assert_file_object(file: FileObject, *, provider: str) -> None: assert file.object == "file", f"file.object={file.object!r}" assert file.purpose == "batch", f"file.purpose={file.purpose!r}" @@ -209,6 +202,10 @@ def assert_file_object(file: FileObject, *, provider: str) -> None: if provider != "bedrock": assert file.bytes > 0, f"file.bytes={file.bytes!r}" assert file.status, "file.status missing" + if provider == "azure": + assert file.expires_at is not None, "Azure batch input has no automatic expiry" + assert file.created_at is not None + assert file.expires_at - file.created_at == AZURE_FILE_EXPIRY_SECONDS assert ( file.created_at is not None and file.created_at > 0 ), "file.created_at missing" @@ -249,7 +246,7 @@ def test_batch_lifecycle( file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key)) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider=provider)) + lambda: cleanup_file(client, file.id, key=key, provider=cap.file_provider) ) assert_file_object(file, provider=cap.provider) assert matches_id_shape( @@ -260,7 +257,7 @@ def test_batch_lifecycle( require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( - quietly(lambda: client.cancel_batch(batch.id, key=key, provider=provider)) + lambda: cleanup_batch(client, batch.id, key=key, provider=provider) ) assert batch.id, f"create returned no batch id (body={created.body[:200]})" @@ -339,7 +336,7 @@ def test_batch_key_model_access_denied( denied_upload = client.upload_file( content=render_jsonl(AZURE_BATCH_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) @@ -356,7 +353,7 @@ def test_batch_key_model_access_denied( ) ).id resources.defer( - quietly(lambda: client.delete_file(raw_file, key=key, provider="openai")) + lambda: cleanup_file(client, raw_file, key=key, provider="openai") ) denied_create = client.create_batch( @@ -383,6 +380,7 @@ def test_file_upload_and_delete_outputs( key=key, ) ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="openai") deleted = unwrap(client.delete_file(file.id, key=key)) @@ -458,12 +456,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) _ = client.proxy.poll_logs_for_key(key, min_rows=1) @@ -517,7 +515,7 @@ class TestBatchFileContent: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert file.id downloaded = client.proxy.transport.download( @@ -559,11 +557,11 @@ class TestBatchFileContent: file = unwrap( client.upload_file( content=payload, - form=FileUploadForm(purpose="batch", target_model_names=provider.model), + form=batch_upload_form(provider.name, target_model_names=provider.model), key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider=provider.name) assert is_managed_id(file.id), ( f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" @@ -626,7 +624,7 @@ class TestOpenAIFiles: ) ) resources.defer( - quietly(lambda: client.delete_file(file.id, key=key, provider="openai")) + lambda: cleanup_file(client, file.id, key=key, provider="openai") ) listed = unwrap(client.list_files(key=key)) @@ -690,7 +688,7 @@ class TestOpenAIFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) fetched = unwrap(client.retrieve_file(file.id, key=key)) assert fetched.id == file.id, "retrieve must echo the uploaded file id" @@ -760,7 +758,7 @@ class TestBatchRateLimitErrorMapping: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -813,7 +811,7 @@ class TestBatchEnqueuedTokenLimit: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) return file def _generate_enqueued_key( @@ -861,7 +859,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) @pytest.mark.covers( "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", @@ -904,7 +902,7 @@ class TestBatchEnqueuedTokenLimit: first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(first) first_batch = BatchObject.model_validate_json(first.body) - resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, first_batch.id, key=key)) blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) assert blocked.status_code == 429, ( @@ -928,7 +926,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(retried) retry_batch = BatchObject.model_validate_json(retried.body) - resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, retry_batch.id, key=key)) ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -984,13 +982,13 @@ class TestBedrockBatchAssumeRole: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="bedrock") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"assume-role create returned no batch id: {created.body[:200]}" assert is_managed_id(batch.id), ( @@ -1044,7 +1042,7 @@ class TestGeminiFiles: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="gemini") assert file.id, "gemini file upload returned no id" @@ -1099,13 +1097,13 @@ class TestHostedVllmBatch: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert_file_object(file, provider="hosted_vllm") created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}" assert batch.status in CREATED_BATCH_STATUSES, ( @@ -1192,7 +1190,7 @@ class TestBatchFailurePaths: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) @@ -1243,12 +1241,12 @@ class TestBatchFailurePaths: file = unwrap( client.upload_file( content=render_jsonl(AZURE_BATCH_RAW_MODEL), - form=FileUploadForm(purpose="batch"), + form=batch_upload_form("azure"), model=AZURE_BATCH_MODEL, key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( f"upload did not encode the azure deployment into the file id: {file.id!r}" ) @@ -1258,7 +1256,7 @@ class TestBatchFailurePaths: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( "create with a foreign encoded file id must route by the file's embedded model, " @@ -1307,7 +1305,7 @@ class TestBatchSecondHop: key=key, ) ) - resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) assert is_managed_id(file.id), ( f"second-hop unified upload must return a managed file id, got {file.id!r}" ) @@ -1315,7 +1313,7 @@ class TestBatchSecondHop: created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) assert is_managed_id(batch.id), ( f"second-hop create must return a managed batch id, got {batch.id!r}" diff --git a/tests/e2e/batches/test_managed_files_enforcement_e2e.py b/tests/e2e/batches/test_managed_files_enforcement_e2e.py index 7ad0b16adc3..4f703cf0fdc 100644 --- a/tests/e2e/batches/test_managed_files_enforcement_e2e.py +++ b/tests/e2e/batches/test_managed_files_enforcement_e2e.py @@ -21,6 +21,7 @@ from typing import Iterator import pytest from batch_client import BatchClient, FileObject +from batch_cleanup import cleanup_file from capabilities import batch_model_name, is_managed_id, openai_batch_params from e2e_config import unique_marker from e2e_http import FileUploadForm, Result, UnknownApiError, unwrap @@ -108,7 +109,7 @@ def test_cross_user_managed_id_denied_owner_allowed( key=owner_key, ) ) - resources.defer(lambda: client.delete_file(uploaded.id, key=owner_key)) + resources.defer(lambda: cleanup_file(client, uploaded.id, key=owner_key)) assert is_managed_id(uploaded.id), f"expected a managed unified file id, got {uploaded.id}" denied = client.retrieve_file(uploaded.id, key=other_key) diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index c9a67ebdb8c..eb9704d4dcb 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -8,8 +8,9 @@ ResourceManager; the test registers a cleanup for every resource it creates, and the fixture's teardown releases them all even when the test body raises. """ +from builtins import ExceptionGroup from dataclasses import dataclass, field -from typing import Callable, List, Protocol, runtime_checkable +from typing import Callable, Final, List, Protocol, runtime_checkable from proxy_client import ProxyClient from models import KeyGenerateBody @@ -52,6 +53,7 @@ class ResourceManager: """ client: ResourceClient + strict_cleanup: bool = False _cleanups: List[Callable[[], object]] = field( default_factory=list ) # mutable-ok: append-only teardown registry @@ -82,8 +84,17 @@ class ResourceManager: return customer_id def teardown(self) -> None: - for cleanup in reversed(self._cleanups): - try: - cleanup() - except Exception: - pass # best-effort: a failed cleanup must not block the rest + failures: Final = tuple( + failure for cleanup in reversed(self._cleanups) + if (failure := _run_cleanup(cleanup)) is not None + ) + if failures and self.strict_cleanup: + raise ExceptionGroup("Resource cleanup failed", failures) + + +def _run_cleanup(cleanup: Callable[[], object]) -> Exception | None: + try: + cleanup() + except Exception as exc: + return exc + return None From a56c60e8924f6f956200c816142fb0cb994fa3ed Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:27:51 -0700 Subject: [PATCH 02/13] fix(e2e): wait for managed batch cancellation before deleting inputs --- tests/e2e/batches/COVERAGE.md | 5 +++-- tests/e2e/batches/batch_cleanup.py | 10 +++++++++- tests/e2e/batches/test_batch_cleanup.py | 14 ++++++++------ 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 899530dde2b..69ba9d781ec 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -129,8 +129,9 @@ provider when deleted. Model-encoded and managed file IDs route themselves File deletion and batch cancellation check their responses and retry transient failures up to three times. Teardown attempts every registered cleanup before reporting failures as test errors. Already deleted files and batches that are -terminal are safe to clean up again. Cancellation polls for up to ten minutes -before input deletion, because accepting cancellation does not finish it +terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes +before input deletion: the ten-minute provider window plus a propagation margin. +Raw and model-encoded inputs can be deleted after cancellation is accepted Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index a5b5e9bba37..dd79c776758 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -5,11 +5,12 @@ from typing import Final, Protocol from pydantic import BaseModel from batch_client import BatchObject, FileDeleteResponse +from capabilities import is_managed_id from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) -BATCH_CANCEL_TIMEOUT_SECONDS: Final = 600.0 +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 BATCH_CANCEL_POLL_SECONDS: Final = 10.0 @@ -62,12 +63,15 @@ def cleanup_batch( wait: Callable[[float], None] = sleep, clock: Callable[[], float] = monotonic, ) -> None: + needs_terminal_state: Final = is_managed_id(batch_id) fetched: Final = _require_cleanup_success( cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), f"Retrieve batch {batch_id} for cleanup", ) if fetched.status in BATCH_TERMINAL_STATUSES: return + if fetched.status == "cancelling" and not needs_terminal_state: + return if fetched.status != "cancelling": result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): @@ -75,6 +79,8 @@ def cleanup_batch( assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( f"Cancel batch {batch_id} left status {cancelled.status}" ) + if not needs_terminal_state: + return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS while True: current = _require_cleanup_success( @@ -84,6 +90,8 @@ def cleanup_batch( if current.status in BATCH_TERMINAL_STATUSES: return assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" + if not needs_terminal_state: + return assert clock() < deadline, ( f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" ) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 8ebfd750360..9715370d9b7 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -12,6 +12,8 @@ from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApi from lifecycle import ResourceManager from models import KeyGenerateBody +MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" + @dataclass class CleanupClient: @@ -122,8 +124,8 @@ class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) delays: Final[list[float]] = [] - cleanup_batch(client, "batch-1", key="test-key", wait=delays.append) - assert client.calls == ["retrieve None batch-1"] * 3 + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays.append) + assert client.calls == [f"retrieve None {MANAGED_BATCH_ID}"] * 3 assert delays == [10.0] def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: @@ -134,13 +136,13 @@ class TestBatchCancellation: manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) - manager.defer(lambda: cleanup_batch(client, "batch-1", key=key, clock=lambda: next(ticks))) + manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=lambda: next(ticks))) with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) assert client.calls == [ - "retrieve None batch-1", - "retrieve None batch-1", + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", "delete None file-1", "delete key test-key", ] @@ -156,7 +158,7 @@ class TestBatchCancellation: batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") - assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1", "retrieve azure batch-1"] + assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1"] @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: From 6bdf206cc5228e85b47c908a2f4977374656289b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:48:07 -0700 Subject: [PATCH 03/13] fix(e2e): accept managed file deletion responses --- tests/e2e/batches/COVERAGE.md | 3 +++ tests/e2e/batches/batch_cleanup.py | 4 +++- tests/e2e/batches/batch_client.py | 2 +- tests/e2e/batches/test_batch_cleanup.py | 15 +++++++++++++++ 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 69ba9d781ec..f95ea1f2649 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -139,6 +139,9 @@ fallback for interrupted runs: immediate deletion remains the normal cleanup. Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot be requested through its Files API +The Azure entry in `files_settings` must use `api_version: 2025-04-01-preview` +for raw uploads to honor expiry, matching the batch deployment's API version + ## Terminal state + cost write-back (cross-run marker baton) The 24h completion window rules out submit-and-wait inside one run, so diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index dd79c776758..3fd6802f696 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -51,7 +51,9 @@ def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider if isinstance(result, UnknownApiError) and result.status_code == 404: return deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") - assert deleted.deleted, f"Delete file {file_id} did not confirm deletion" + assert deleted.deleted is True or ( + deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file" + ), f"Delete file {file_id} did not confirm deletion" def cleanup_batch( diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 84a902b0b11..c9c77e1f12e 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -99,7 +99,7 @@ class BatchList(BaseModel): class FileDeleteResponse(BaseModel): id: str object: str | None = None - deleted: bool + deleted: bool | None = None class BatchCreateBody(BaseModel): diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 9715370d9b7..15dead6d36d 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -12,6 +12,7 @@ from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApi from lifecycle import ResourceManager from models import KeyGenerateBody +MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" @@ -53,6 +54,20 @@ def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]: class TestFileCleanup: + def test_managed_delete_accepts_the_deleted_file_object(self) -> None: + response: Final = Success( + status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) + ) + client: Final = CleanupClient(files=iter((response,))) + cleanup_file(client, MANAGED_FILE_ID, key="test-key") + assert client.calls == [f"delete None {MANAGED_FILE_ID}"] + + @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) + def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: + client: Final = CleanupClient(files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),))) + with pytest.raises(AssertionError, match="did not confirm deletion"): + cleanup_file(client, file_id, key="test-key") + @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: client: Final = CleanupClient(files=iter((deleted_file(),))) From a096dd615c71e40be9473338ffeb8d1c17284f51 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 14:39:25 -0700 Subject: [PATCH 04/13] refactor(e2e): use immutable batch cleanup test expectations --- tests/e2e/batches/batch_cleanup.py | 7 +- tests/e2e/batches/test_batch_cleanup.py | 137 ++++++++++++++++-------- 2 files changed, 96 insertions(+), 48 deletions(-) diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 3fd6802f696..722df0c29bc 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -1,4 +1,5 @@ from collections.abc import Callable +from itertools import count from time import monotonic, sleep from typing import Final, Protocol @@ -84,11 +85,13 @@ def cleanup_batch( if not needs_terminal_state: return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS - while True: - current = _require_cleanup_success( + for current in ( + _require_cleanup_success( cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)), f"Retrieve batch {batch_id} after cancellation", ) + for _ in count() + ): if current.status in BATCH_TERMINAL_STATUSES: return assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 15dead6d36d..66b7079ccc2 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -16,33 +16,44 @@ MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" -@dataclass +@dataclass(frozen=True, slots=True) +class ExpectedCalls[T]: + values: Iterator[T] + + def __call__(self, value: T) -> None: + assert next(self.values, None) == value + + def assert_done(self) -> None: + assert tuple(self.values) == () + + +@dataclass(frozen=True, slots=True) class CleanupClient: + calls: ExpectedCalls[str] files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) - calls: list[str] = field(default_factory=list) def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: - self.calls.append(f"delete {provider} {file_id}") + self.calls(f"delete {provider} {file_id}") return next(self.files) def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: - self.calls.append(f"retrieve {provider} {batch_id}") + self.calls(f"retrieve {provider} {batch_id}") return next(self.batches) def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: - self.calls.append(f"cancel {provider} {batch_id}") + self.calls(f"cancel {provider} {batch_id}") return next(self.cancellations) def generate_key(self, body: KeyGenerateBody) -> str: return "test-key" def delete_key(self, key: str) -> None: - self.calls.append(f"delete key {key}") + self.calls(f"delete key {key}") def delete_customers(self, user_ids: list[str]) -> None: - self.calls.append(f"delete customers {user_ids}") + self.calls(f"delete customers {user_ids}") def batch(status: str) -> Success[BatchObject]: @@ -58,51 +69,71 @@ class TestFileCleanup: response: Final = Success( status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) ) - client: Final = CleanupClient(files=iter((response,))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete None {MANAGED_FILE_ID}",))), files=iter((response,)) + ) cleanup_file(client, MANAGED_FILE_ID, key="test-key") - assert client.calls == [f"delete None {MANAGED_FILE_ID}"] + client.calls.assert_done() @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: - client: Final = CleanupClient(files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete None {file_id}",))), + files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),)), + ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, file_id, key="test-key") + client.calls.assert_done() @pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES]) def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: - client: Final = CleanupClient(files=iter((deleted_file(),))) - cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None - assert client.calls == [f"delete {expected_provider} file-1"] + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"delete {expected_provider} file-1",))), files=iter((deleted_file(),)) + ) + cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) + client.calls.assert_done() def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="secret response"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete azure file-1", "delete key test-key"))), + files=iter((UnknownApiError(status_code=403, body="secret response"),)), + ) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure")) with pytest.raises(ExceptionGroup) as caught: manager.teardown() - assert client.calls == ["delete azure file-1", "delete key test-key"] + client.calls.assert_done() assert len(caught.value.exceptions) == 1 assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403" def test_success_response_must_confirm_deletion(self) -> None: - client: Final = CleanupClient(files=iter((deleted_file(deleted=False),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete None file-1",))), files=iter((deleted_file(deleted=False),)) + ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, "file-1", key="test-key") + client.calls.assert_done() def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=404, body="missing"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete azure file-1",))), + files=iter((UnknownApiError(status_code=404, body="missing"),)), + ) cleanup_file(client, "file-1", key="test-key", provider="azure") - assert client.calls == ["delete azure file-1"] + client.calls.assert_done() def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: - client: Final = CleanupClient(files=iter((UnknownApiError(status_code=403, body="forbidden"),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("delete None file-1", "delete key test-key"))), + files=iter((UnknownApiError(status_code=403, body="forbidden"),)), + ) manager: Final = ResourceManager(client=client) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) manager.teardown() - assert client.calls == ["delete None file-1", "delete key test-key"] + client.calls.assert_done() class TestCleanupRetries: @@ -112,40 +143,54 @@ class TestCleanupRetries: ) def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: outcomes: Final = iter((failure, deleted_file())) - delays: Final[list[float]] = [] - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + delays: Final = ExpectedCalls(iter((1.0,))) + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) assert isinstance(result, Success) and result.data.deleted - assert delays == [1.0] + delays.assert_done() def test_persistent_error_has_bounded_retries(self) -> None: failure: Final = UnknownApiError(status_code=503, body="unavailable") outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) - delays: Final[list[float]] = [] - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays.append) + delays: Final = ExpectedCalls(iter(CLEANUP_DELAYS)) + result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) assert result is failure - assert tuple(delays) == CLEANUP_DELAYS + delays.assert_done() assert next(outcomes, None) is None def test_permanent_error_is_not_retried(self) -> None: failure: Final = UnknownApiError(status_code=403, body="forbidden") outcomes: Final = iter((failure, deleted_file())) - delays: Final[list[float]] = [] - assert cleanup_result(lambda: next(outcomes), wait=delays.append) is failure - assert delays == [] + delays: Final = ExpectedCalls[float](iter(())) + assert cleanup_result(lambda: next(outcomes), wait=delays) is failure + delays.assert_done() assert isinstance(next(outcomes), Success) class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: - client: Final = CleanupClient(batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled")))) - delays: Final[list[float]] = [] - cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays.append) - assert client.calls == [f"retrieve None {MANAGED_BATCH_ID}"] * 3 - assert delays == [10.0] + client: Final = CleanupClient( + calls=ExpectedCalls(iter((f"retrieve None {MANAGED_BATCH_ID}",) * 3)), + batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled"))), + ) + delays: Final = ExpectedCalls(iter((10.0,))) + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays) + client.calls.assert_done() + delays.assert_done() def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: client: Final = CleanupClient( - batches=iter((batch("cancelling"), batch("cancelling"))), files=iter((deleted_file(),)) + calls=ExpectedCalls( + iter( + ( + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", + "delete None file-1", + "delete key test-key", + ) + ) + ), + batches=iter((batch("cancelling"), batch("cancelling"))), + files=iter((deleted_file(),)), ) ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) manager: Final = ResourceManager(client=client, strict_cleanup=True) @@ -155,29 +200,29 @@ class TestBatchCancellation: with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) - assert client.calls == [ - f"retrieve None {MANAGED_BATCH_ID}", - f"retrieve None {MANAGED_BATCH_ID}", - "delete None file-1", - "delete key test-key", - ] + client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: - client: Final = CleanupClient(batches=iter((batch(status),))) + client: Final = CleanupClient( + calls=ExpectedCalls(iter(("retrieve None batch-1",))), batches=iter((batch(status),)) + ) cleanup_batch(client, "batch-1", key="test-key") - assert client.calls == ["retrieve None batch-1"] + client.calls.assert_done() def test_active_batch_is_cancelled_through_its_provider(self) -> None: client: Final = CleanupClient( - batches=iter((batch("in_progress"), batch("cancelled"))), cancellations=iter((batch("cancelling"),)) + calls=ExpectedCalls(iter(("retrieve azure batch-1", "cancel azure batch-1"))), + batches=iter((batch("in_progress"), batch("cancelled"))), + cancellations=iter((batch("cancelling"),)), ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") - assert client.calls == ["retrieve azure batch-1", "cancel azure batch-1"] + client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( + calls=ExpectedCalls(iter(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"))), batches=iter((batch("in_progress"), batch(status))), cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), ) @@ -186,7 +231,7 @@ class TestBatchCancellation: else: with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"): cleanup_batch(client, "batch-1", key="test-key") - assert client.calls == ["retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"] + client.calls.assert_done() class TestAzureFileExpiry: From 7bff9bf9a2594d3cd5dd6a80a44e4e3b1cd38f61 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:22:32 -0700 Subject: [PATCH 05/13] fix(batches): handle provider cancellation and file cleanup gaps --- litellm/llms/bedrock/files/transformation.py | 70 ++++++++------ tests/e2e/batches/COVERAGE.md | 7 +- tests/e2e/batches/batch_cleanup.py | 63 +++++++++--- tests/e2e/batches/test_batch_cleanup.py | 69 +++++++++++++- tests/e2e/batches/test_batches_e2e.py | 14 +-- .../test_bedrock_files_transformation.py | 95 ++++++++++++++++--- 6 files changed, 258 insertions(+), 60 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 33b27943ad8..90b539ff37c 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -7,7 +7,7 @@ from contextlib import suppress from functools import cache from itertools import chain from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict +from typing import Any, Final, Literal, TypeAlias, TypedDict from urllib.parse import unquote import httpx @@ -60,11 +60,8 @@ from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id -# litellm_params key used to hand the SigV4-signed GET headers from -# `transform_file_content_request` to `validate_environment` (the only hook -# the shared file-content HTTP handler exposes for setting request headers). -# Same pattern as the `upload_url` handoff in `transform_create_file_request`. -S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" +S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" +S3_DELETE_FILE_ID_PARAM: Final = "_s3_delete_file_id" # litellm_params key carrying the size of the body uploaded to S3, handed from # `transform_create_file_request` to `transform_create_file_response`. @@ -291,7 +288,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) -> dict: result: Final[dict[str, object]] = {} result.update(headers) - signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None) + signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None) if isinstance(signed_headers, Mapping): result.update(signed_headers) # any-ok: untyped handoff headers # otherwise no extra headers - AWS credentials are handled by BaseAWSLLM @@ -1187,18 +1184,31 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + request: Final = self._transform_s3_file_request( + file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params + ) + litellm_params[S3_DELETE_FILE_ID_PARAM] = file_id + return request def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: - raise NotImplementedError("BedrockFilesConfig does not support file deletion") + if raw_response.status_code != 204: + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}", + headers=raw_response.headers, + ) + file_id: Final = litellm_params.get(S3_DELETE_FILE_ID_PARAM) + if not isinstance(file_id, str) or not file_id: + raise ValueError("Missing file id for Bedrock file deletion response") + return FileDeleted(id=file_id, deleted=True, object="file") def transform_list_files_request( self, @@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if not file_id: raise ValueError("file_id is required for Bedrock file content retrieval") + return self._transform_s3_file_request( + file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params + ) + + def _transform_s3_file_request( + self, + *, + file_id: str, + method: Literal["GET", "DELETE"], + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: s3_uri: Final = extract_s3_uri_from_file_id(file_id) bucket_name, object_key = _validate_file_id_against_configured_buckets( s3_uri=s3_uri, @@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) - # The shared file-content handler passes optional_params={}, so AWS - # credentials/region arrive via litellm_params here (unlike the upload - # path). s3_region_name wins over aws_region_name, same priority as - # get_complete_file_url above. - merged_params: Final[dict[str, object]] = {} - merged_params.update(litellm_params) - merged_params.update(optional_params) - request_params: Final = _BedrockS3RequestParams.model_validate(merged_params) + request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params}) region_preference: Final = request_params.s3_region_name or request_params.aws_region_name region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = ( + s3_endpoint_url: Final = ( request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" - litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body( api_base=url, aws_region_name=aws_region_name, request_params=request_params, + method=method, ) return url, {} - def _sign_s3_get_request( + def _sign_s3_request_without_body( self, api_base: str, aws_region_name: str, request_params: _BedrockS3RequestParams, + method: Literal["GET", "DELETE"] = "GET", ) -> dict[str, str]: - """ - SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT). - """ try: import hashlib @@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): empty_body_hash: Final = hashlib.sha256(b"").hexdigest() aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped - method="GET", + method=method, url=api_base, headers={"x-amz-content-sha256": empty_body_hash}, ) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index f95ea1f2649..ca44fc95e25 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -131,7 +131,12 @@ failures up to three times. Teardown attempts every registered cleanup before reporting failures as test errors. Already deleted files and batches that are terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes before input deletion: the ten-minute provider window plus a propagation margin. -Raw and model-encoded inputs can be deleted after cancellation is accepted +Accepted cancellation may still report validating or in_progress while the provider +updates its state. Raw and model-encoded batches are polled until cancelling or +terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes +output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE +restricted to the configured storage buckets and managed file prefixes. The low-RPM +test submits with its restricted key and cleans up with the test administrator key Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 722df0c29bc..9284882ad82 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -1,16 +1,17 @@ +from builtins import ExceptionGroup from collections.abc import Callable from itertools import count from time import monotonic, sleep from typing import Final, Protocol -from pydantic import BaseModel - from batch_client import BatchObject, FileDeleteResponse from capabilities import is_managed_id from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError +from pydantic import BaseModel CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) +BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"}) BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 BATCH_CANCEL_POLL_SECONDS: Final = 10.0 @@ -63,6 +64,7 @@ def cleanup_batch( *, key: str, provider: str | None = None, + delete_output_files: bool = False, wait: Callable[[float], None] = sleep, clock: Callable[[], float] = monotonic, ) -> None: @@ -72,18 +74,28 @@ def cleanup_batch( f"Retrieve batch {batch_id} for cleanup", ) if fetched.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, fetched, key=key, provider=provider) return if fetched.status == "cancelling" and not needs_terminal_state: return - if fetched.status != "cancelling": - result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) - if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}): - cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") - assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, ( - f"Cancel batch {batch_id} left status {cancelled.status}" - ) - if not needs_terminal_state: - return + result: Final = ( + Success(status_code=200, data=fetched) + if fetched.status == "cancelling" + else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider)) + ) + conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409} + if not conflicted: + cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}") + assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, ( + f"Cancel batch {batch_id} left status {cancelled.status}" + ) + if cancelled.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, cancelled, key=key, provider=provider) + return + if cancelled.status == "cancelling" and not needs_terminal_state: + return deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS for current in ( _require_cleanup_success( @@ -93,11 +105,36 @@ def cleanup_batch( for _ in count() ): if current.status in BATCH_TERMINAL_STATUSES: + if delete_output_files: + _cleanup_batch_outputs(client, current, key=key, provider=provider) return - assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}" - if not needs_terminal_state: + assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), ( + f"Cancel batch {batch_id} left status {current.status}" + ) + if current.status == "cancelling" and not needs_terminal_state: return assert clock() < deadline, ( f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s" ) wait(BATCH_CANCEL_POLL_SECONDS) + + +def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None: + errors: Final = tuple( + error + for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id)) + if file_id is not None and file_id != batch.input_file_id + if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None + ) + if errors: + raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors) + + +def _output_cleanup_error( + client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None +) -> Exception | None: + try: + cleanup_file(client, file_id, key=key, provider=provider) + except Exception as error: + return error + return None diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 66b7079ccc2..a875aee719b 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -4,7 +4,6 @@ from dataclasses import dataclass, field from typing import Final import pytest - from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form from capabilities import CAPABILITIES, Capability @@ -219,6 +218,74 @@ class TestBatchCancellation: cleanup_batch(client, "batch-1", key="test-key", provider="azure") client.calls.assert_done() + @pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID]) + @pytest.mark.parametrize("pending_status", ["validating", "in_progress"]) + def test_accepted_cancellation_waits_through_stale_provider_status( + self, batch_id: str, pending_status: str + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + iter( + ( + f"retrieve vertex_ai {batch_id}", + f"cancel vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + "delete vertex_ai file-1", + "delete key test-key", + ) + ) + ), + batches=iter((batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled"))), + cancellations=iter((batch(pending_status),)), + files=iter((deleted_file(),)), + ) + delays: Final = ExpectedCalls(iter((10.0, 10.0))) + manager: Final = ResourceManager(client=client, strict_cleanup=True) + key: Final = manager.key() + manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai")) + manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays)) + manager.teardown() + client.calls.assert_done() + delays.assert_done() + + @pytest.mark.parametrize("output_delete_fails", [False, True]) + def test_batch_that_completed_before_cleanup_deletes_output_and_error_files( + self, output_delete_fails: bool + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls( + iter(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")) + ), + batches=iter( + ( + Success( + status_code=200, + data=BatchObject( + id="batch-1", + status="completed", + input_file_id="file-input", + output_file_id="file-output", + error_file_id="file-error", + ), + ), + ) + ), + files=iter( + ( + UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), + deleted_file(), + ) + ), + ) + if output_delete_fails: + with pytest.raises(ExceptionGroup, match="output cleanup failed"): + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + else: + cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True) + client.calls.assert_done() + @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1b4a6ed266f..c4b699190b8 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -25,7 +25,7 @@ from datetime import datetime, timedelta, timezone import pytest from pydantic import BaseModel -from e2e_config import PROXY_BASE_URL, unique_marker +from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker from batch_cleanup import cleanup_batch, cleanup_file from batch_client import ( @@ -257,7 +257,9 @@ def test_batch_lifecycle( require_successful_call(created) batch = BatchObject.model_validate_json(created.body) resources.defer( - lambda: cleanup_batch(client, batch.id, key=key, provider=provider) + lambda: cleanup_batch( + client, batch.id, key=key, provider=provider, delete_output_files=cap.provider in {"openai", "azure"} + ) ) assert batch.id, f"create returned no batch id (body={created.body[:200]})" @@ -801,7 +803,7 @@ class TestBatchEnqueuedTokenLimit: """ def _upload_batch_file( - self, client: BatchClient, resources: ResourceManager, key: str + self, client: BatchClient, resources: ResourceManager, key: str, *, cleanup_key: str | None = None ) -> FileObject: file = unwrap( client.upload_file( @@ -811,7 +813,7 @@ class TestBatchEnqueuedTokenLimit: key=key, ) ) - resources.defer(lambda: cleanup_file(client, file.id, key=key)) + resources.defer(lambda: cleanup_file(client, file.id, key=cleanup_key or key)) return file def _generate_enqueued_key( @@ -848,7 +850,7 @@ class TestBatchEnqueuedTokenLimit: marker="rpm", rpm_limit=BATCH_RL_RPM_LIMIT, ) - file = self._upload_batch_file(client, resources, key) + file = self._upload_batch_file(client, resources, key, cleanup_key=MASTER_KEY) created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) @@ -859,7 +861,7 @@ class TestBatchEnqueuedTokenLimit: ) require_successful_call(created) batch = BatchObject.model_validate_json(created.body) - resources.defer(lambda: cleanup_batch(client, batch.id, key=key)) + resources.defer(lambda: cleanup_batch(client, batch.id, key=MASTER_KEY, delete_output_files=True)) @pytest.mark.covers( "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 541c0db15d8..2c02a58663e 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -5,6 +5,8 @@ Test bedrock files transformation functionality import json import os from collections.abc import Mapping +from contextlib import AsyncExitStack, closing +from typing import Final from unittest.mock import MagicMock from urllib.parse import unquote, urlparse @@ -1855,6 +1857,77 @@ class TestBedrockBatchNonChatEndpointRecords: ] +class TestBedrockFileDeletion: + S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl" + URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl" + + def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with respx.mock, closing(HTTPHandler()) as client: + route: Final = respx.delete(self.URL).mock(return_value=httpx.Response(204)) + deleted: Final = litellm.file_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + request: Final = route.calls[0].request + assert request.content == b"" + signed: Final = AWSRequest(method="DELETE", url=self.URL, headers={ + "X-Amz-Date": request.headers["X-Amz-Date"], + "X-Amz-Content-SHA256": request.headers["X-Amz-Content-SHA256"], + }) + signed.context["timestamp"] = request.headers["X-Amz-Date"] + auth: Final = S3SigV4Auth(Credentials("AKIAEXAMPLE", "test-secret"), "s3", "us-west-2") + signature: Final = auth.signature(auth.string_to_sign(signed, auth.canonical_request(signed)), signed) + assert request.headers["Authorization"].endswith(f"Signature={signature}") + assert deleted.id == self.S3_URI and deleted.deleted is True + + @pytest.mark.asyncio + async def test_adelete_file_propagates_s3_errors(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + import respx + + import litellm + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + async with AsyncExitStack() as stack: + client: Final = AsyncHTTPHandler() + stack.push_async_callback(client.close) + with respx.mock: + route: Final = respx.delete(self.URL).mock( + return_value=httpx.Response(403, content=b"AccessDenied") + ) + from litellm.llms.bedrock.common_utils import BedrockError + + with pytest.raises(BedrockError, match="AccessDenied"): + await litellm.afile_delete( + file_id=self.S3_URI, custom_llm_provider="bedrock", client=client, + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2", + ) + assert route.call_count == 1 + + @pytest.mark.parametrize("file_id, message", [ + ("s3://other-bucket/litellm-bedrock-files-model-abc.jsonl", "configured storage bucket"), + ("s3://my-bucket/private/data.jsonl", "LiteLLM-managed"), + ]) + def test_delete_rejects_untrusted_objects_before_signing( + self, file_id: str, message: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + with pytest.raises(ValueError, match=message): + BedrockFilesConfig().transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params={}) + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" @@ -1873,7 +1946,7 @@ class TestBedrockFileContentTransformation: import hashlib from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -1889,7 +1962,7 @@ class TestBedrockFileContentTransformation: assert url == self.EXPECTED_URL assert params == {} - signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM] + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] content_hashes = { value for name, value in signed_headers.items() @@ -2139,7 +2212,7 @@ class TestBedrockFileContentTransformation: def test_s3_region_name_wins_for_content_signing(self, monkeypatch): """s3_region_name must override aws_region_name for both the URL and the signature.""" from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2154,17 +2227,17 @@ class TestBedrockFileContentTransformation: ) assert url.startswith("https://s3.eu-west-1.amazonaws.com/") - authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"] + authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"] assert "/eu-west-1/s3/aws4_request" in authorization def test_validate_environment_merges_and_pops_signed_get_headers(self): from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) litellm_params = { - S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} + S3_SIGNED_REQUEST_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} } headers = BedrockFilesConfig().validate_environment( @@ -2179,7 +2252,7 @@ class TestBedrockFileContentTransformation: "x-custom": "kept", "Authorization": "AWS4-HMAC-SHA256 test", } - assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params + assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params def test_transform_file_content_response_wraps_binary_content(self): import httpx @@ -2379,7 +2452,7 @@ class TestBedrockFilesS3SignatureEncoding: self, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.llms.bedrock.files.transformation import ( - S3_SIGNED_GET_HEADERS_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig, ) @@ -2402,7 +2475,7 @@ class TestBedrockFilesS3SignatureEncoding: method="GET", url=url, body=None, - headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM], + headers=litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM], ) @@ -2457,7 +2530,7 @@ def test_sign_s3_request_assumes_role_with_external_id(monkeypatch): assert "ASIAFILESPUTROLE" in authorization -def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): +def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch): """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request.""" import datetime from unittest.mock import patch @@ -2504,7 +2577,7 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): assert request_params.aws_external_id == "external-id-files-get" with patch.object(boto3, "client", return_value=FakeSTSClient()): - signed_headers = BedrockFilesConfig()._sign_s3_get_request( + signed_headers = BedrockFilesConfig()._sign_s3_request_without_body( api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", aws_region_name="us-east-1", request_params=request_params, From 4ab5719ff9e0770ecb9f2d1b53c4caf58f19e5db Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:36:27 -0700 Subject: [PATCH 06/13] test(batches): use immutable expectations with explicit test doubles --- litellm/files/main.py | 2 +- tests/e2e/batches/test_batch_cleanup.py | 187 ++++++++++++------------ 2 files changed, 93 insertions(+), 96 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 19da77b7364..218518eb3cd 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -31,7 +31,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] import litellm from litellm import get_secret_str diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index a875aee719b..d0038139dcf 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -1,7 +1,7 @@ from builtins import ExceptionGroup -from collections.abc import Iterator -from dataclasses import dataclass, field +from collections.abc import Callable from typing import Final +from unittest.mock import Mock, call import pytest from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result @@ -15,35 +15,43 @@ MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" -@dataclass(frozen=True, slots=True) class ExpectedCalls[T]: - values: Iterator[T] + def __init__(self, values: tuple[T, ...]) -> None: + self.values: Final = values + self.recorder: Final = Mock() def __call__(self, value: T) -> None: - assert next(self.values, None) == value + self.recorder(value) def assert_done(self) -> None: - assert tuple(self.values) == () + assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values) -@dataclass(frozen=True, slots=True) class CleanupClient: - calls: ExpectedCalls[str] - files: Iterator[Result[FileDeleteResponse]] = field(default_factory=lambda: iter(())) - batches: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) - cancellations: Iterator[Result[BatchObject]] = field(default_factory=lambda: iter(())) + def __init__( + self, + *, + calls: ExpectedCalls[str], + files: tuple[Result[FileDeleteResponse], ...] = (), + batches: tuple[Result[BatchObject], ...] = (), + cancellations: tuple[Result[BatchObject], ...] = (), + ) -> None: + self.calls: Final = calls + self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files) + self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches) + self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations) def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: self.calls(f"delete {provider} {file_id}") - return next(self.files) + 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 next(self.batches) + return self.batch_response() def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: self.calls(f"cancel {provider} {batch_id}") - return next(self.cancellations) + return self.cancel_response() def generate_key(self, body: KeyGenerateBody) -> str: return "test-key" @@ -68,17 +76,15 @@ class TestFileCleanup: response: Final = Success( status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"}) ) - client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete None {MANAGED_FILE_ID}",))), files=iter((response,)) - ) + client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,)) cleanup_file(client, MANAGED_FILE_ID, key="test-key") client.calls.assert_done() @pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID]) def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete None {file_id}",))), - files=iter((Success(status_code=200, data=FileDeleteResponse(id=file_id)),)), + calls=ExpectedCalls((f"delete None {file_id}",)), + files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),), ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, file_id, key="test-key") @@ -88,15 +94,15 @@ class TestFileCleanup: def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None: expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"delete {expected_provider} file-1",))), files=iter((deleted_file(),)) + calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),) ) cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider) client.calls.assert_done() def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete azure file-1", "delete key test-key"))), - files=iter((UnknownApiError(status_code=403, body="secret response"),)), + calls=ExpectedCalls(("delete azure file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="secret response"),), ) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() @@ -109,7 +115,7 @@ class TestFileCleanup: def test_success_response_must_confirm_deletion(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete None file-1",))), files=iter((deleted_file(deleted=False),)) + calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),) ) with pytest.raises(AssertionError, match="did not confirm deletion"): cleanup_file(client, "file-1", key="test-key") @@ -117,16 +123,16 @@ class TestFileCleanup: def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete azure file-1",))), - files=iter((UnknownApiError(status_code=404, body="missing"),)), + calls=ExpectedCalls(("delete azure file-1",)), + files=(UnknownApiError(status_code=404, body="missing"),), ) cleanup_file(client, "file-1", key="test-key", provider="azure") client.calls.assert_done() def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("delete None file-1", "delete key test-key"))), - files=iter((UnknownApiError(status_code=403, body="forbidden"),)), + calls=ExpectedCalls(("delete None file-1", "delete key test-key")), + files=(UnknownApiError(status_code=403, body="forbidden"),), ) manager: Final = ResourceManager(client=client) key: Final = manager.key() @@ -141,37 +147,39 @@ class TestCleanupRetries: [NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")], ) def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None: - outcomes: Final = iter((failure, deleted_file())) - delays: Final = ExpectedCalls(iter((1.0,))) - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls((1.0,)) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) assert isinstance(result, Success) and result.data.deleted delays.assert_done() def test_persistent_error_has_bounded_retries(self) -> None: failure: Final = UnknownApiError(status_code=503, body="unavailable") - outcomes: Final[Iterator[Result[FileDeleteResponse]]] = iter((failure,) * (len(CLEANUP_DELAYS) + 1)) - delays: Final = ExpectedCalls(iter(CLEANUP_DELAYS)) - result: Final[Result[FileDeleteResponse]] = cleanup_result(lambda: next(outcomes), wait=delays) + outcomes: Final = Mock(return_value=failure) + delays: Final = ExpectedCalls(CLEANUP_DELAYS) + result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays) assert result is failure delays.assert_done() - assert next(outcomes, None) is None + assert outcomes.call_count == len(CLEANUP_DELAYS) + 1 def test_permanent_error_is_not_retried(self) -> None: failure: Final = UnknownApiError(status_code=403, body="forbidden") - outcomes: Final = iter((failure, deleted_file())) - delays: Final = ExpectedCalls[float](iter(())) - assert cleanup_result(lambda: next(outcomes), wait=delays) is failure + responses: Final = (failure, deleted_file()) + outcomes: Final = Mock(side_effect=responses) + delays: Final = ExpectedCalls[float](()) + assert cleanup_result(outcomes, wait=delays) is failure delays.assert_done() - assert isinstance(next(outcomes), Success) + assert outcomes.call_count == 1 class TestBatchCancellation: def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter((f"retrieve None {MANAGED_BATCH_ID}",) * 3)), - batches=iter((batch("cancelling"), batch("cancelling"), batch("cancelled"))), + calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3), + batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")), ) - delays: Final = ExpectedCalls(iter((10.0,))) + delays: Final = ExpectedCalls((10.0,)) cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays) client.calls.assert_done() delays.assert_done() @@ -179,23 +187,22 @@ class TestBatchCancellation: def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: client: Final = CleanupClient( calls=ExpectedCalls( - iter( - ( - f"retrieve None {MANAGED_BATCH_ID}", - f"retrieve None {MANAGED_BATCH_ID}", - "delete None file-1", - "delete key test-key", - ) + ( + f"retrieve None {MANAGED_BATCH_ID}", + f"retrieve None {MANAGED_BATCH_ID}", + "delete None file-1", + "delete key test-key", ) ), - batches=iter((batch("cancelling"), batch("cancelling"))), - files=iter((deleted_file(),)), + batches=(batch("cancelling"), batch("cancelling")), + files=(deleted_file(),), ) - ticks: Final = iter((0.0, BATCH_CANCEL_TIMEOUT_SECONDS)) + times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS) + ticks: Final[Callable[[], float]] = Mock(side_effect=times) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key)) - manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=lambda: next(ticks))) + manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks)) with pytest.raises(ExceptionGroup) as caught: manager.teardown() assert "cancellation did not finish" in str(caught.value.exceptions[0]) @@ -203,17 +210,15 @@ class TestBatchCancellation: @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) def test_inactive_batch_needs_no_cancellation(self, status: str) -> None: - client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve None batch-1",))), batches=iter((batch(status),)) - ) + client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),)) cleanup_batch(client, "batch-1", key="test-key") client.calls.assert_done() def test_active_batch_is_cancelled_through_its_provider(self) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve azure batch-1", "cancel azure batch-1"))), - batches=iter((batch("in_progress"), batch("cancelled"))), - cancellations=iter((batch("cancelling"),)), + calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")), + batches=(batch("in_progress"), batch("cancelled")), + cancellations=(batch("cancelling"),), ) cleanup_batch(client, "batch-1", key="test-key", provider="azure") client.calls.assert_done() @@ -225,23 +230,21 @@ class TestBatchCancellation: ) -> None: client: Final = CleanupClient( calls=ExpectedCalls( - iter( - ( - f"retrieve vertex_ai {batch_id}", - f"cancel vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - f"retrieve vertex_ai {batch_id}", - "delete vertex_ai file-1", - "delete key test-key", - ) + ( + f"retrieve vertex_ai {batch_id}", + f"cancel vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + f"retrieve vertex_ai {batch_id}", + "delete vertex_ai file-1", + "delete key test-key", ) ), - batches=iter((batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled"))), - cancellations=iter((batch(pending_status),)), - files=iter((deleted_file(),)), + batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")), + cancellations=(batch(pending_status),), + files=(deleted_file(),), ) - delays: Final = ExpectedCalls(iter((10.0, 10.0))) + delays: Final = ExpectedCalls((10.0, 10.0)) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai")) @@ -255,28 +258,22 @@ class TestBatchCancellation: self, output_delete_fails: bool ) -> None: client: Final = CleanupClient( - calls=ExpectedCalls( - iter(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")) - ), - batches=iter( - ( - Success( - status_code=200, - data=BatchObject( - id="batch-1", - status="completed", - input_file_id="file-input", - output_file_id="file-output", - error_file_id="file-error", - ), + calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")), + batches=( + Success( + status_code=200, + data=BatchObject( + id="batch-1", + status="completed", + input_file_id="file-input", + output_file_id="file-output", + error_file_id="file-error", ), - ) + ), ), - files=iter( - ( - UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), - deleted_file(), - ) + files=( + UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(), + deleted_file(), ), ) if output_delete_fails: @@ -289,9 +286,9 @@ class TestBatchCancellation: @pytest.mark.parametrize("status", ["completed", "in_progress"]) def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None: client: Final = CleanupClient( - calls=ExpectedCalls(iter(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1"))), - batches=iter((batch("in_progress"), batch(status))), - cancellations=iter((UnknownApiError(status_code=409, body="conflict"),)), + calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")), + batches=(batch("in_progress"), batch(status)), + cancellations=(UnknownApiError(status_code=409, body="conflict"),), ) if status == "completed": cleanup_batch(client, "batch-1", key="test-key") From 7d3b68fea5a6343781401e52f40b02f6c581541c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 16:59:53 -0700 Subject: [PATCH 07/13] fix(files): preserve managed deletion routing and response identity --- .../proxy/hooks/managed_files.py | 13 ++- tests/e2e/batches/COVERAGE.md | 3 + .../proxy/test_managed_files_hook.py | 104 ++++++++++++++++++ 3 files changed, 118 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc1eb6cebc2..6e0bb0da3f4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1779,7 +1779,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Remove conflicting keys from data to avoid duplicate keyword arguments filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): - delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + delete_data = { + **{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"}, + **( + {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))} + if credentials is not None + else {} + ), + } + delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data) stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) @@ -1790,7 +1799,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): prom_logger.record_managed_file_deleted(result="success") if stored_file_object: - return stored_file_object + return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id}) elif delete_response: delete_response.id = file_id return delete_response diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca44fc95e25..919c39f21a2 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -138,6 +138,9 @@ output and error files returned by terminal batches. Bedrock deletion uses a sig restricted to the configured storage buckets and managed file prefixes. The low-RPM test submits with its restricted key and cleans up with the test administrator key +Managed deletion forwards the deployment's trusted bucket configuration and returns +the requested managed file ID even when stored output metadata carries a provider ID + Azure input uploads request `expires_after` anchored to `created_at` with `seconds=1209600`, and the lifecycle tests check the returned expiry. This is a fallback for interrupted runs: immediate deletion remains the normal cleanup. diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 091b958d7c3..48fceb50403 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -1095,6 +1095,110 @@ async def test_afile_content_passes_trusted_model_credentials_to_router(): assert trusted_credentials["s3_bucket_name"] == "my-bucket" +def _managed_deletion_file_id(provider_file_id): + from litellm.types.utils import SpecialEnums + + value = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "test-file", "batch-model", provider_file_id, "model-123" + ) + return base64.urlsafe_b64encode(value.encode()).decode().rstrip("=") + + +def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object): + from litellm.caching import DualCache + from litellm.models.managed_files import LiteLLM_ManagedFileTable + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + row = LiteLLM_ManagedFileTable( + unified_file_id=unified_file_id, + model_mappings={"model-123": provider_file_id}, + flat_model_file_ids=[provider_file_id], + file_object=file_object, + ) + table = MagicMock( + find_first=AsyncMock(return_value=row), + delete=AsyncMock(), + ) + return _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)), + ), table + + +@pytest.mark.asyncio +async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete(monkeypatch): + import httpx + import respx + + from litellm import Router + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + router = Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + "model_info": {"id": "model-123"}, + } + ], + num_retries=0, + ) + s3_uri = "s3://my-bucket/litellm-bedrock-files/input.jsonl" + unified_file_id = _managed_deletion_file_id(s3_uri) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, s3_uri, None) + with respx.mock: + route = respx.delete( + "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl" + ).mock(return_value=httpx.Response(204)) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert len(route.calls) == 1 + assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert response.id == unified_file_id + assert response.deleted is True + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + +@pytest.mark.asyncio +async def test_afile_delete_returns_managed_id_for_stored_provider_output(): + from openai.types import FileDeleted + + provider_file_id = "file-error-output" + unified_file_id = _managed_deletion_file_id(provider_file_id) + stored_file = _make_file_object(provider_file_id) + managed_files, table = _managed_files_with_deletion_row(unified_file_id, provider_file_id, stored_file) + router = MagicMock( + get_deployment_credentials_with_provider=MagicMock(return_value=None), + afile_delete=AsyncMock(return_value=FileDeleted(id=provider_file_id, object="file", deleted=True)), + ) + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + _litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"}, + ) + + assert response.id == unified_file_id + assert response.object == "file" + assert response.filename == stored_file.filename + assert stored_file.id == provider_file_id + router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id) + table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) + + @pytest.mark.asyncio async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): """ From 00381ef03baedc9c9f7bec55f8c6139ecdb75f21 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 15:07:18 -0700 Subject: [PATCH 08/13] test: validate opaque stream IDs and hide log-reader credentials --- .../test_responses_bridge_streaming_e2e.py | 24 ++++++++++++------- tests/e2e/logging/datadog_reader.py | 10 ++++---- tests/e2e/logging/test_datadog_reader.py | 20 ++++++++++++++++ 3 files changed, 40 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/logging/test_datadog_reader.py diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py index 9a45743a0cd..75817340876 100644 --- a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p from __future__ import annotations +from typing import Final, Literal + import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import StreamingResponse @@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel): class _BridgeChunk(BaseModel): id: str - choices: list[_BridgeChoice] = [] + object: Literal["chat.completion.chunk"] + choices: list[_BridgeChoice] = Field(default_factory=list) class _WeatherArgs(BaseModel): @@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming: resources.key(), ChatBody( model=bridged_model, - messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], max_tokens=64, stream=True, ), ) - chunks = _bridge_chunks(result) - ids = {chunk.id for chunk in chunks} + chunks: Final = _bridge_chunks(result) + assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk" + ids: Final = frozenset(chunk.id for chunk in chunks) assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" - assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id" @pytest.mark.covers( "llm.chat_completions.openai.basic.stream.bridge_streams_sse", @@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming: chunks = _bridge_chunks(result) content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" - assert any( - choice.finish_reason for chunk in chunks for choice in chunk.choices - ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), ( + f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + ) assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" @pytest.mark.covers( diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index d0f478185c2..7b4372ef9ba 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -13,7 +13,7 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations import time -from dataclasses import dataclass +from dataclasses import dataclass, field import pytest from pydantic import BaseModel, ConfigDict, Field @@ -36,8 +36,8 @@ _RATE_LIMIT_RETRIES = 5 class _DdAuthHeaders(Headers): - api_key: str = Field(serialization_alias="DD-API-KEY") - app_key: str = Field(serialization_alias="DD-APPLICATION-KEY") + api_key: str = Field(serialization_alias="DD-API-KEY", repr=False) + app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False) class _SearchFilter(BaseModel): @@ -88,8 +88,8 @@ class _SearchResponse(BaseModel): @dataclass(frozen=True, slots=True) class DdLogsReader: site: str - api_key: str - app_key: str + api_key: str = field(repr=False) + app_key: str = field(repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py new file mode 100644 index 00000000000..00624bd3c84 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,20 @@ +from typing import Final + +from datadog_reader import DdLogsReader +from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization + + +def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: + api_key: Final = "test-datadog-api-secret" + app_key: Final = "test-datadog-app-secret" + reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key) + headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key) + + for value in (reader, headers): + assert api_key not in repr(value) + assert app_key not in repr(value) + + assert headers.model_dump(by_alias=True) == { + "DD-API-KEY": api_key, + "DD-APPLICATION-KEY": app_key, + } From 64afa9d6eca8d238da02458a2769294eff93d08d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 15:36:34 -0700 Subject: [PATCH 09/13] test: isolate auto-router scenarios and clean partial setup --- .../test_auto_router_regressions_e2e.py | 239 +++++++++--------- 1 file changed, 120 insertions(+), 119 deletions(-) diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 188db2a8eb5..374badcf5fc 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form. import json import os from collections.abc import Iterator +from contextlib import ExitStack from dataclasses import dataclass from typing import Final @@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel): @dataclass(frozen=True, slots=True) -class TagSplitDeployments: - """Scenario A mirrors the customer-shaped config from GitHub issue #36619: - plain deployment registered first, tier deployment and marker both tagged. - Scenario B flips both axes for GitHub issue #36621: marker registered first - and its tier deployment left untagged, so routing depends neither on - registration order nor on tier deployments carrying tags.""" - - tag_a: str - shared_a: str - tier_a: str - tag_b: str - shared_b: str - tier_b: str +class TagSplitDeployment: + tag: str + shared: str + tier: str @dataclass(frozen=True, slots=True) @@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]: } -def _key_for( - proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False -) -> str: +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str: key: Final = proxy.generate_key( KeyGenerateBody( models=models, @@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con ) -@pytest.fixture(scope="module") -def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _register_models( + proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...] +) -> None: + for name, params in registrations: + stack.callback(proxy.delete_model, proxy.create_model(name, params)) + + +def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment: marker: Final = unique_marker() - deployments: Final = TagSplitDeployments( - tag_a=f"e2e-split-a-{marker}", - shared_a=f"e2e-autoroute-a-{marker}", - tier_a=f"e2e-tier-a-{marker}", - tag_b=f"e2e-split-b-{marker}", - shared_b=f"e2e-autoroute-b-{marker}", - tier_b=f"e2e-tier-b-{marker}", + named: Final = TagSplitDeployment( + tag=f"e2e-split-{marker}", + shared=f"e2e-autoroute-{marker}", + tier=f"e2e-tier-{marker}", ) anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") - marker_params_a: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_a), - tags=[deployments.tag_a], + marker_registration: Final = ( + named.shared, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + tags=[named.tag], + ), ) - marker_params_b: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_b), - tags=[deployments.tag_b], + tier_registration: Final = ( + named.tier, + LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]), ) - registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( - (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), - (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), - (deployments.shared_a, marker_params_a), - (deployments.shared_b, marker_params_b), - (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), - (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)) + registrations: Final = ( + (marker_registration, tier_registration, plain_registration) + if marker_first + else (plain_registration, tier_registration, marker_registration) ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield deployments - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, stack, registrations) + return named -@pytest.fixture(scope="module") -def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: +@pytest.fixture(scope="class") +def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=False) + + +@pytest.fixture(scope="class") +def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=True) + + +@pytest.fixture(scope="class") +def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias: marker: Final = unique_marker() named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: +@pytest.fixture(scope="class") +def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit: marker: Final = unique_marker() named: Final = HeuristicSplit( alias=f"e2e-heuristic-router-{marker}", @@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: +@pytest.fixture(scope="class") +def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter: marker: Final = unique_marker() named: Final = SemanticAutoRouter( marker=f"e2e-semantic-router-{marker}", @@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.marker, marker_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: +@pytest.fixture(scope="class") +def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias: marker: Final = unique_marker() named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named class TestTagSplitRouting: @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") def test_body_tagged_chat_routes_through_the_marker_to_its_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36619: with tag filtering on, a chat request whose body metadata tags match the tagged marker under a shared model name is answered by the marker's tier deployment, not by the plain deployment that was registered under the name first.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag]))) assert chat.choices, "tagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_chat_is_always_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620: untagged chat requests to the shared name succeed on every call and are all served by the plain deployment; the tagged marker never captures them, so no intermittent auto-router errors and no tier hijacking.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) for _ in range(5): - chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared))) assert chat.choices, "untagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=5) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + _assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_messages_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620 on the /v1/messages surface: an untagged Anthropic-native request to the shared name is served by the plain deployment, not captured by the tagged marker.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared))) assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name" + ) class TestUntaggedTierDeployments: @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36621: a /v1/messages request tagged only via the x-litellm-tags header selects the tagged marker, and the rewrite still lands on the tier deployment even though that deployment carries no tags, because the marker consumed the routing tags.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag) answer: Final = unwrap( proxy.transport.post( "/v1/messages", headers=headers, - json=_hello_messages_body(split.shared_b), + json=_hello_messages_body(marker_first_split.shared), response_type=AnthropicMessagesResponse, ) ) assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + _assert_served_only_by( + rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name" + ) @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins the tag-consumption half of GitHub issue #36621: after the tagged marker rewrites the request to its tier model, the consumed routing tags no longer constrain deployment selection, so the untagged tier deployment serves the request instead of a strict-tag denial.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + chat: Final = unwrap( + proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag])) + ) assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + _assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier") @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """The tag-consumption fix must not loosen strict tag semantics: a tagged request aimed directly at an untagged deployment (no marker involved) is still rejected with the 401 tags-configuration error.""" - key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) - result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag])) assert isinstance(result, UnauthorizedError), ( f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" ) @@ -451,37 +450,39 @@ class TestUntaggedTierDeployments: class TestResponsesApiTagRouting: @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_header_tagged_responses_with_string_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the /v1/responses surface of the tag split (GitHub issues #36620/#36621): a /v1/responses request with string input, tagged via the x-litellm-tags header, succeeds and routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) ) assert answer.id, "header-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + _assert_served_only_by( + rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input" + ) @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_body_tagged_responses_with_list_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the body-tag and list-input combination of the same split: /v1/responses with litellm_metadata.tags and structured input items routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, + model=plain_first_split.shared, input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], max_output_tokens=64, - litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]), ) answer: Final = unwrap( proxy.transport.post( @@ -493,18 +494,18 @@ class TestResponsesApiTagRouting: ) assert answer.id, "body-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_responses_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the untagged half of the /v1/responses tag split: an untagged request to the shared name is served by the plain deployment, matching the chat and messages surfaces.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post( @@ -516,7 +517,9 @@ class TestResponsesApiTagRouting: ) assert answer.id, "untagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name" + ) class TestStrategyAliasPricing: @@ -551,9 +554,7 @@ class TestComplexityHeuristicScope: while the accompanying ~2KB agent system prompt is packed with enough reasoning and complexity keywords that scoring the combined text lands in REASONING; only ask-only scoring keeps this on the cheap tier.""" - key: Final = _key_for( - proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] - ) + key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]) body: Final = ChatBody( model=heuristic_split.alias, messages=[ From 253600fc6144747c1ffd954c3ebdfadd686adc2e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 16:52:56 -0700 Subject: [PATCH 10/13] test: wait for requested guardrail propagation --- tests/e2e/guardrails/guardrails_client.py | 22 ++++++- .../e2e/guardrails/test_guardrails_client.py | 63 +++++++++++++++++++ .../test_tool_permission_guardrail_e2e.py | 24 ++++--- 3 files changed, 98 insertions(+), 11 deletions(-) create mode 100644 tests/e2e/guardrails/test_guardrails_client.py diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 1f55a0f9a56..3770de0c6d5 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -7,7 +7,7 @@ from __future__ import annotations import time from collections.abc import Callable from dataclasses import dataclass -from typing import Literal +from typing import Final, Literal from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap @@ -405,6 +405,26 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient: return GuardrailsClient(proxy=proxy) +def poll_until_guardrail_applied( + call: Callable[[], StreamingResponse], + guardrail_name: str, + *, + timeout: float = POLL_TIMEOUT, + interval: float = POLL_INTERVAL, + now: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> StreamingResponse: + deadline: Final = now() + timeout + while ( + (result := call()).ok + and guardrail_name + not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(",")) + and (remaining := deadline - now()) > 0 + ): + sleep(min(interval, remaining)) + return result + + def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]: """Retry a call that a guardrail should reject until it is, returning the last result. diff --git a/tests/e2e/guardrails/test_guardrails_client.py b/tests/e2e/guardrails/test_guardrails_client.py new file mode 100644 index 00000000000..aab163d86d2 --- /dev/null +++ b/tests/e2e/guardrails/test_guardrails_client.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass +from itertools import chain, repeat +from typing import Final + +import pytest + +from e2e_http import StreamingResponse +from guardrails_client import poll_until_guardrail_applied + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +def _response(applied: str, status: int = 200) -> StreamingResponse: + return StreamingResponse(status_code=status, body="{}", headers={"x-litellm-applied-guardrails": applied}) + + +def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> None: + clock: Final = Clock() + expected: Final = _response("global-filter, tool-permission") + responses: Final = iter((_response("global-filter"), expected)) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is expected + assert clock.elapsed == 2 + + +@pytest.mark.parametrize("applied", ("", "global-filter", "tool-permission-sibling")) +def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None: + clock: Final = Clock() + missing: Final = _response(applied) + + result: Final = poll_until_guardrail_applied( + lambda: missing, "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is missing + assert clock.elapsed == 5 + + +@pytest.mark.parametrize("status", (400, 401, 429, 500)) +def test_http_failure_is_not_hidden_by_a_later_success(status: int) -> None: + clock: Final = Clock() + failed: Final = _response("", status) + responses: Final = iter(chain((failed,), repeat(_response("tool-permission")))) + + result: Final = poll_until_guardrail_applied( + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + ) + + assert result is failed + assert clock.elapsed == 0 diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py index 9ef3650625c..8d1047e53c7 100644 --- a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -30,6 +30,7 @@ from guardrails_client import ( ToolPermissionParamsBody, ToolPermissionRuleBody, poll_until_blocked, + poll_until_guardrail_applied, ) from lifecycle import ResourceManager from models import ChatResponse, ChatTool, ChatToolFunction @@ -84,8 +85,8 @@ def _register_tool_permission(client: GuardrailsClient, resources: ResourceManag resources.defer(lambda: client.delete_guardrail(guardrail_id)) -def _applied_guardrails(outcome: StreamingResponse) -> str: - return outcome.headers.get("x-litellm-applied-guardrails", "") +def _applied_guardrails(outcome: StreamingResponse) -> tuple[str, ...]: + return tuple(name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(",")) def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: @@ -144,14 +145,17 @@ class TestToolPermissionPreCall: name = f"e2e-toolperm-allow-{unique_marker()}" _register_tool_permission(client, resources, name=name) - outcome = client.chat_raw( - scoped_key, - MODEL, - TOOL_PROMPT, - guardrails=[name], - max_tokens=128, - tools=[ALLOWED_TOOL], - tool_choice="required", + outcome = poll_until_guardrail_applied( + lambda: client.chat_raw( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[ALLOWED_TOOL], + tool_choice="required", + ), + name, ) assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" From b8be30219c5e5dea4343b3a924b89494e7350303 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 17:02:00 -0700 Subject: [PATCH 11/13] test: stop guardrail retries at the polling deadline --- tests/e2e/guardrails/guardrails_client.py | 7 +++++-- tests/e2e/guardrails/test_guardrails_client.py | 5 ++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 3770de0c6d5..ed112a79b9b 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -415,13 +415,16 @@ def poll_until_guardrail_applied( sleep: Callable[[float], None] = time.sleep, ) -> StreamingResponse: deadline: Final = now() + timeout + if not (result := call()).ok: + return result while ( - (result := call()).ok - and guardrail_name + guardrail_name not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(",")) and (remaining := deadline - now()) > 0 ): sleep(min(interval, remaining)) + if now() >= deadline or not (result := call()).ok: + break return result diff --git a/tests/e2e/guardrails/test_guardrails_client.py b/tests/e2e/guardrails/test_guardrails_client.py index aab163d86d2..423c2ede599 100644 --- a/tests/e2e/guardrails/test_guardrails_client.py +++ b/tests/e2e/guardrails/test_guardrails_client.py @@ -40,13 +40,16 @@ def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None: clock: Final = Clock() missing: Final = _response(applied) + responses: Final = iter((missing, missing, missing)) result: Final = poll_until_guardrail_applied( - lambda: missing, "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep + lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep ) assert result is missing assert clock.elapsed == 5 + with pytest.raises(StopIteration): + next(responses) @pytest.mark.parametrize("status", (400, 401, 429, 500)) From 24ee66328a7bf1a791dc0f7d81051165611877c9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 17:55:21 -0700 Subject: [PATCH 12/13] test: bound Datadog read-back retries using reset headers --- tests/e2e/logging/datadog_reader.py | 114 ++++++++----- tests/e2e/logging/test_datadog_reader.py | 203 +++++++++++++++++++++++ 2 files changed, 278 insertions(+), 39 deletions(-) diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index 7b4372ef9ba..368c20cb6aa 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations +import math +import random import time +from collections.abc import Callable, Mapping from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,12 +31,28 @@ from e2e_config import ( DD_SITE, POLL_TIMEOUT, ) -from e2e_http import URL, Headers, RateLimitedError, Success, post +from e2e_http import URL, Headers, StreamingResponse, send -#: How many rate-limited responses in a row one search tolerates before the -#: hard fail; each retry sleeps a full search interval, so this rides out a -#: burst from a concurrent consumer of the org-wide search budget. -_RATE_LIMIT_RETRIES = 5 +type SearchCall = Callable[[str, float], StreamingResponse] + + +def _seconds(value: str | None) -> float | None: + if value is None: + return None + try: + seconds: Final = float(value) + except ValueError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + +def _rate_limit_delay(headers: Mapping[str, str]) -> float: + delays: Final = tuple( + delay + for name in ("x-ratelimit-reset", "retry-after") + if (delay := _seconds(headers.get(name))) is not None + ) + return max(1.0, max(delays, default=DD_SEARCH_INTERVAL)) class _DdAuthHeaders(Headers): @@ -90,6 +110,10 @@ class DdLogsReader: site: str api_key: str = field(repr=False) app_key: str = field(repr=False) + search: SearchCall | None = field(default=None, repr=False) + now: Callable[[], float] = field(default=time.monotonic, repr=False) + sleep: Callable[[float], None] = field(default=time.sleep, repr=False) + jitter: Callable[[], float] = field(default=random.random, repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog @@ -108,25 +132,28 @@ class DdLogsReader: a single event. A 429 backs off and retries - the search budget is org-wide, so another consumer can empty it under us - while any other failure stays a hard fail.""" - for _ in range(_RATE_LIMIT_RETRIES): - result = post( - URL(f"https://api.{self.site}/api/v2/logs/events/search"), - headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=query)), - response_type=_SearchResponse, - timeout=30.0, - ) - match result: - case Success(data=page): - return [event.attributes for event in page.data] - case RateLimitedError(retry_after_seconds=retry_after): - time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL) - case failure: - pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}") + return self._events_for_query(query, self.now() + POLL_TIMEOUT) + + def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]: + search: Final = self.search or self._search_page + while (remaining := deadline - self.now()) > 0: + if (result := search(query, min(30.0, remaining))).ok: + return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data] + if result.status_code != 429: + pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}") + if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0: + self.sleep(delay) pytest.fail( - f"DataDog Logs Search API at api.{self.site} still rate-limited after " - f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide " - "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer" + f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; " + "the org-wide logs_public_search_api budget is exhausted" + ) + + def _search_page(self, query: str, timeout: float) -> StreamingResponse: + return send( + URL(f"https://api.{self.site}/api/v2/logs/events/search"), + headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), + json=_SearchRequest(filter=_SearchFilter(query=query)), + timeout=timeout, ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: @@ -140,33 +167,42 @@ class DdLogsReader: hide from the exactly-one assertion - real-DataDog jitter can surface one call's two events tens of seconds apart. Searches pace at DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's - request budget. At the deadline the last result is returned as-is.""" - deadline = time.monotonic() + POLL_TIMEOUT - while time.monotonic() < deadline: - events = self.events_for_query(query) + request budget. Discovery, quota retries, and duplicate detection share + one POLL_TIMEOUT deadline; an incomplete settle window fails closed.""" + deadline: Final = self.now() + POLL_TIMEOUT + while (remaining := deadline - self.now()) > 0: + events = self._events_for_query(query, deadline) if events: - return self._settled_events_for_query(query, events) - time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_query(query) + return self._settled_events_for_query(query, events, deadline) + if (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + return [] - def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. Keep the last non-empty result: a transient empty search (index lag) must not erase events already confirmed earlier in the settle window. + A successful final search must reach the full settle window before the + shared read-back deadline; otherwise duplicate detection is incomplete. """ - settle_deadline = time.monotonic() + DD_SETTLE_SECONDS + settle_deadline: Final = self.now() + DD_SETTLE_SECONDS last_nonempty = events - while time.monotonic() < settle_deadline: - time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_query(query) - if not latest: - continue + if len(events) > 1: + return events + while (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + if self.now() >= deadline: + break + latest = self._events_for_query(query, deadline) if len(latest) > 1: return latest - last_nonempty = latest - return last_nonempty + if latest: + last_nonempty = latest + if self.now() >= settle_deadline: + return last_nonempty + pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s") def build_dd_logs_reader() -> DdLogsReader: diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py index 00624bd3c84..910a1cefd42 100644 --- a/tests/e2e/logging/test_datadog_reader.py +++ b/tests/e2e/logging/test_datadog_reader.py @@ -1,7 +1,14 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass from typing import Final +import pytest + from datadog_reader import DdLogsReader from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization +from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT +from e2e_http import StreamingResponse def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: @@ -18,3 +25,199 @@ def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> "DD-API-KEY": api_key, "DD-APPLICATION-KEY": app_key, } + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +@dataclass +class Search: + responses: Iterator[StreamingResponse] + calls: tuple[tuple[str, float], ...] = () + + def __call__(self, query: str, timeout: float) -> StreamingResponse: + self.calls += ((query, timeout),) + return next(self.responses) + + +def _page(*event_ids: str) -> StreamingResponse: + return StreamingResponse( + status_code=200, + body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}), + ) + + +def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]: + search: Final = Search(iter(responses)) + return DdLogsReader( + site="us5.datadoghq.com", + api_key="test-api-secret", + app_key="test-app-secret", + search=search, + now=clock.now, + sleep=clock.sleep, + jitter=lambda: 0.25, + ), search + + +def test_429_honors_server_reset_and_preserves_duplicate_events() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")), + clock, + ) + + events: Final = reader.events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 6.25 + assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0)) + + +@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1")) +def test_invalid_reset_uses_search_interval(reset: str) -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25 + + +def test_zero_reset_cannot_create_a_busy_retry_loop() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 1.25 + + +def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 8.25 + + +def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls == (("test-marker", 30.0),) + + +def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75)) + + +@pytest.mark.parametrize("status", (-1, 401, 403, 500)) +def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None: + clock: Final = Clock() + reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock) + + with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"): + reader.events_for_query("test-marker") + + assert search.calls == (("test-marker", 30.0),) + assert clock.elapsed == 0 + + +def test_polling_quota_retries_share_the_original_deadline() -> None: + clock: Final = Clock() + reader, search = _reader( + (_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == 2 + + +def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None: + clock: Final = Clock() + attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) + reader, search = _reader((_page(),) * attempts, clock) + + assert reader.poll_events_for_query("test-marker") == [] + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == attempts + + +def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader( + (_page(),) * empty_reads + + (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL) + assert len(search.calls) == empty_reads + 2 + + +def test_settlement_detects_a_duplicate_on_the_final_search() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 30 + + +def test_settlement_keeps_confirmed_events_through_empty_searches() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first",) + assert clock.elapsed == 30 + + +def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock) + + with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == empty_reads + 2 From 163847333242b7e75b1a15c27519beef2b1365ed Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 19:02:08 -0700 Subject: [PATCH 13/13] fix(bedrock): keep deletion response IDs in request context --- litellm/llms/bedrock/files/transformation.py | 18 ++++++------- .../test_bedrock_files_transformation.py | 27 +++++++++++++++++++ 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 90b539ff37c..9875ac2b9c3 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -13,7 +13,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from typing_extensions import ReadOnly from litellm._logging import verbose_logger @@ -61,7 +61,11 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" -S3_DELETE_FILE_ID_PARAM: Final = "_s3_delete_file_id" + + +class _S3DeleteContext(BaseModel): + file_id: str = Field(min_length=1) + # litellm_params key carrying the size of the body uploaded to S3, handed from # `transform_create_file_request` to `transform_create_file_response`. @@ -1187,11 +1191,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: MutableMapping[str, object], ) -> tuple[str, dict[str, str]]: - request: Final = self._transform_s3_file_request( + return self._transform_s3_file_request( file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params ) - litellm_params[S3_DELETE_FILE_ID_PARAM] = file_id - return request def transform_delete_file_response( self, @@ -1205,10 +1207,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}", headers=raw_response.headers, ) - file_id: Final = litellm_params.get(S3_DELETE_FILE_ID_PARAM) - if not isinstance(file_id, str) or not file_id: - raise ValueError("Missing file id for Bedrock file deletion response") - return FileDeleted(id=file_id, deleted=True, object="file") + context: Final = _S3DeleteContext.model_validate(logging_obj.model_call_details.get("additional_args")) + return FileDeleted(id=context.file_id, deleted=True, object="file") def transform_list_files_request( self, diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 2c02a58663e..3b01a4f2054 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1861,6 +1861,33 @@ class TestBedrockFileDeletion: S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl" URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl" + def test_interleaved_deletions_keep_their_own_file_ids(self, monkeypatch: pytest.MonkeyPatch) -> None: + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + config: Final = BedrockFilesConfig() + params: Final = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + } + file_ids: Final = (self.S3_URI, "s3://my-bucket/litellm-bedrock-files-model-second.jsonl") + for file_id in file_ids: + config.transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params=params) + + deleted: Final = tuple( + config.transform_delete_file_response( + raw_response=httpx.Response(204), + logging_obj=MagicMock(model_call_details={"additional_args": {"file_id": file_id}}), + litellm_params=params, + ).id + for file_id in file_ids + ) + + assert deleted == file_ids + def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None: import httpx import respx