From 6c21be619441dbd24879e1f8a4b897c6dbd667fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 12:15:06 -0700 Subject: [PATCH 01/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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 From 0721163cacfbd9fbfeee5dd25205cb76840111dd Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:49:40 -0700 Subject: [PATCH 14/19] test(e2e/ui): cover team-scoped model visibility, re-editing litellm params, and model health checks (#40039) * test(e2e/ui): cover team-scoped model visibility, re-editing litellm params, and model health checks Three Models and Endpoints flows had no end-to-end coverage, and all three keep coming back as bug reports. modelsByTeam walks an internal user through the Current team control and asserts the table lists exactly what each team grants. It creates one deployment that belongs to no team, proves that deployment is visible under Personal, then proves it is absent under both seeded teams, so an empty table cannot pass the same assertions. editLitellmParams adds a temperature and a custom pair to a deployment, saves, then re-edits the temperature and drops the custom pair. It checks both update request bodies, polls the stored deployment until the new temperature is there, reloads the page to confirm the second save is what renders, and sends one chat completion to prove the deployment still serves. modelHealthStatus runs the health check on a reachable deployment and on one pointed at a dead port, asserts the healthy and unhealthy cells and the two detail dialogs, and reloads to confirm both statuses are stored. Every deployment these specs create carries a unique name and is deleted in afterEach, including on the failure path. * test(e2e/ui): find health rows across every page of the health table The health table pages server-side at 50 rows with no search box, so on a proxy carrying more deployments than that the two deployments the spec creates can land on a later page and the lookup finds nothing. Row lookups now walk the pages, using the table's own page indicator to know when to advance and when to wrap back to the first page. * test(e2e/ui): build the created deployment ids without mutating the array * test(ui): scope model deployments to Playwright fixtures --- .../tests/internal-user/modelsByTeam.spec.ts | 196 ++++++++++++++ .../modelsPage/editLitellmParams.spec.ts | 252 ++++++++++++++++++ .../modelsPage/modelHealthStatus.spec.ts | 245 +++++++++++++++++ 3 files changed, 693 insertions(+) create mode 100644 tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts create mode 100644 tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts create mode 100644 tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts diff --git a/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts new file mode 100644 index 00000000000..5e2c80b5845 --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts @@ -0,0 +1,196 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_ORG_ALIAS, + INTERNAL_USER_STORAGE_PATH, +} from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, CHAT_MODEL_B, masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CURRENT_TEAM_VIEW = "Current Team Models"; +const ALL_MODELS_VIEW = "All Available Models"; +const PERSONAL_TEAM = "Personal"; + +const teamSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "Current team", exact: true }); +const viewSelector = (page: PlaywrightPage): Locator => + page.getByRole("combobox", { name: "View", exact: true }); + +async function chooseOption( + page: PlaywrightPage, + selector: Locator, + optionName: string, +): Promise { + await selector.click(); + const option = page.getByRole("option", { name: optionName, exact: true }); + await expect(option, `option ${optionName} is offered`).toBeVisible({ + timeout: 10_000, + }); + await option.click(); + await expect( + selector, + `${optionName} is the selection the control now reports`, + ).toContainText(optionName, { + timeout: 10_000, + }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +function modelRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ ungrantedModelName: string }>({ + ungrantedModelName: async ({ page }, use) => { + const ungrantedModelName = `e2e-ungranted-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: ungrantedModelName, + litellm_params: { + model: `openai/${ungrantedModelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const ungrantedModelId = (await created.json()).model_info?.id; + expect(ungrantedModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll(async () => await isRegistered(page, ungrantedModelName), { + message: `deployment ${ungrantedModelName} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await use(ungrantedModelName); + } finally { + await deleteDeployment(page, ungrantedModelId); + } + }, +}); + +test.describe("Models and Endpoints for an internal user", () => { + test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); + + test("shows an internal user exactly the models of the team they select", async ({ + page, + ungrantedModelName, + }) => { + await navigateToPage(page, Page.Models); + + await expect( + page.getByRole("tab", { name: "Your Models" }), + "an internal user lands on their own models tab, not an admin-only view", + ).toBeVisible({ timeout: 15_000 }); + await expect( + viewSelector(page), + "the models table opens scoped to the selected team", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + `the personal view lists ${ungrantedModelName}, so it is on the proxy and reachable from this page`, + ).toHaveCount(1, { timeout: 30_000 }); + + await chooseOption(page, teamSelector(page), E2E_TEAM_CRUD_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_B}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is on the proxy but not granted to ${E2E_TEAM_CRUD_ALIAS}, so it must not be listed`, + ).toHaveCount(0); + + await chooseOption(page, teamSelector(page), E2E_TEAM_ORG_ALIAS); + await expect( + modelRow(page, CHAT_MODEL_A), + `${E2E_TEAM_ORG_ALIAS} lists ${CHAT_MODEL_A}`, + ).toHaveCount(1, { + timeout: 15_000, + }); + await expect( + page.getByTestId("pagination-range"), + `${E2E_TEAM_ORG_ALIAS} lists the one model it grants and nothing else`, + ).toHaveText("Showing 1-1 of 1", { timeout: 15_000 }); + await expect( + modelRow(page, CHAT_MODEL_B), + `${CHAT_MODEL_B} belongs to another team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + await expect( + modelRow(page, ungrantedModelName), + `${ungrantedModelName} is granted to no team and must not leak into ${E2E_TEAM_ORG_ALIAS}`, + ).toHaveCount(0); + + await chooseOption(page, viewSelector(page), ALL_MODELS_VIEW); + await expect( + modelRow(page, CHAT_MODEL_A), + `switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`, + ).toHaveCount(1, { timeout: 15_000 }); + + await page.reload(); + await expect( + teamSelector(page), + "the team selection is not persisted across a reload, so the table returns to the personal view", + ).toContainText(PERSONAL_TEAM, { timeout: 15_000 }); + await expect( + viewSelector(page), + "the view selection is not persisted across a reload either", + ).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 }); + await expect( + modelRow(page, ungrantedModelName), + "the personal view still renders models after a reload rather than coming back empty", + ).toHaveCount(1, { timeout: 30_000 }); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts new file mode 100644 index 00000000000..4480515ae59 --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts @@ -0,0 +1,252 @@ +import { + test as base, + expect, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { masterKey, sendChatCompletion } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const CUSTOM_PARAM = "extra_headers"; +const CUSTOM_PARAM_VALUE = { "X-E2E-Edit-Probe": "one" }; + +type StoredParams = Record; + +async function readStoredParams( + page: PlaywrightPage, + modelId: string, +): Promise { + const body = await readBack<{ data: { litellm_params: StoredParams }[] }>( + page, + `/model/info?litellm_model_id=${modelId}`, + ); + return body.data[0]?.litellm_params ?? {}; +} + +function paramsEditor(page: PlaywrightPage) { + return page.getByPlaceholder('"rpm": 100'); +} + +async function editParams( + page: PlaywrightPage, + mutate: (params: StoredParams) => StoredParams, +): Promise { + await page.getByRole("button", { name: "Edit Settings" }).click(); + const editor = paramsEditor(page); + await expect( + editor, + "the LiteLLM Params editor is reachable on every visit to the edit form", + ).toBeVisible({ + timeout: 15_000, + }); + const shown = JSON.parse(await editor.inputValue()) as StoredParams; + await editor.fill(JSON.stringify(mutate(shown), null, 2)); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const test = base.extend<{ + deployment: { readonly modelName: string; readonly createdModelId: string }; +}>({ + deployment: async ({ page, request }, use) => { + const modelName = `e2e-edit-params-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: modelName, + litellm_params: { + model: `openai/${modelName}`, + api_base: MOCK_LLM_BASE, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const createdModelId = (await created.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + + try { + await expect + .poll( + async () => { + try { + await sendChatCompletion(request, { + model: modelName, + prompt: `warmup ${modelName}`, + }); + return true; + } catch { + return false; + } + }, + { + message: `deployment ${modelName} never became routable after /model/new`, + timeout: 60_000, + }, + ) + .toBe(true); + await use({ modelName, createdModelId }); + } finally { + await deleteDeployment(page, createdModelId); + } + }, +}); + +test.describe("Edit LiteLLM Params on a deployment", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("params added on a deployment can be re-edited, and the deployment keeps serving", async ({ + page, + request, + deployment: { modelName, createdModelId }, + }) => { + await navigateToPage(page, Page.Models); + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect( + modelIdCell, + `the Models table lists ${modelName}`, + ).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ + timeout: 15_000, + }); + + await editParams(page, (params) => ({ + ...params, + temperature: 0.2, + [CUSTOM_PARAM]: CUSTOM_PARAM_VALUE, + })); + const firstSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + firstSave.litellm_params?.temperature, + "the added temperature goes on the wire", + ).toBe(0.2); + expect( + firstSave.litellm_params?.[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} goes on the wire`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + firstSave.litellm_params?.model, + "a params edit does not rewrite the upstream model", + ).toBe(`openai/${modelName}`); + expect( + firstSave.litellm_params?.api_base, + "a params edit does not rewrite the api base", + ).toBe(MOCK_LLM_BASE); + expect( + firstSave.litellm_params, + "the credential is never re-sent, so a masked placeholder cannot overwrite the stored key", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: "the added temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.2); + const afterFirstSave = await readStoredParams(page, createdModelId); + expect( + afterFirstSave[CUSTOM_PARAM], + `the added ${CUSTOM_PARAM} reached the stored deployment`, + ).toEqual(CUSTOM_PARAM_VALUE); + expect( + afterFirstSave.model, + "the stored upstream model survived the edit", + ).toBe(`openai/${modelName}`); + expect( + afterFirstSave.api_base, + "the stored api base survived the edit", + ).toBe(MOCK_LLM_BASE); + + await editParams(page, (params) => ({ + ...Object.fromEntries( + Object.entries(params).filter(([key]) => key !== CUSTOM_PARAM), + ), + temperature: 0.7, + })); + const secondSave = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect( + secondSave.litellm_params?.temperature, + "a param set by an earlier save can be edited again", + ).toBe(0.7); + expect( + secondSave.litellm_params, + `dropping ${CUSTOM_PARAM} from the editor drops it from the request the UI sends`, + ).not.toHaveProperty(CUSTOM_PARAM); + expect( + secondSave.litellm_params?.model, + "a second params edit still leaves the upstream model alone", + ).toBe(`openai/${modelName}`); + expect( + secondSave.litellm_params?.api_base, + "a second params edit still leaves the api base alone", + ).toBe(MOCK_LLM_BASE); + expect( + secondSave.litellm_params, + "the credential is still never re-sent", + ).not.toHaveProperty("api_key"); + + await expect + .poll( + async () => (await readStoredParams(page, createdModelId)).temperature, + { + message: + "the re-edited temperature never reached the stored deployment", + timeout: 20_000, + }, + ) + .toBe(0.7); + + await page.reload(); + await expect( + page + .getByRole("tabpanel", { name: "Overview" }) + .getByText('"temperature": 0.7'), + "reopening the deployment renders the re-edited value, not the one from the first save", + ).toBeVisible({ timeout: 20_000 }); + + await sendChatCompletion(request, { + model: modelName, + prompt: `still serving ${modelName}`, + }); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts new file mode 100644 index 00000000000..247cce1b85d --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts @@ -0,0 +1,245 @@ +import { + test as base, + expect, + type Locator, + type Page as PlaywrightPage, +} from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; + +const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; +const UNREACHABLE_BASE = "http://127.0.0.1:9/v1"; + +async function isRegistered( + page: PlaywrightPage, + modelName: string, +): Promise { + const body = await readBack<{ data: { model_name?: string }[] }>( + page, + "/v2/model/info", + ); + return body.data.some((row) => row.model_name === modelName); +} + +function healthRow(page: PlaywrightPage, modelName: string): Locator { + return page.getByRole("row").filter({ hasText: modelName }); +} + +function pageOf(label: string): { current: number; total: number } { + const [current, total] = label + .replace("Page ", "") + .split(" of ") + .map((part) => Number(part.trim())); + return { current, total }; +} + +async function locateHealthRow( + page: PlaywrightPage, + modelName: string, +): Promise { + const pageLabel = page.getByTestId("pagination-page"); + await expect( + pageLabel, + "the health table reports which page it is showing", + ).toBeVisible({ timeout: 20_000 }); + + const deadline = Date.now() + 60_000; + while (Date.now() < deadline) { + const row = healthRow(page, modelName); + const onThisPage = await row + .first() + .waitFor({ state: "visible", timeout: 3_000 }) + .then(() => true) + .catch(() => false); + if (onThisPage) return row; + + const { current, total } = pageOf(await pageLabel.innerText()); + const goTo = current < total ? current + 1 : 1; + if (total === 1) continue; + await page + .getByRole("button", { + name: current < total ? "Go to next page" : "Go to first page", + }) + .click(); + await expect(pageLabel).toContainText(`Page ${goTo} of`, { + timeout: 15_000, + }); + } + return healthRow(page, modelName); +} + +async function openHealthTab(page: PlaywrightPage): Promise { + await page.getByRole("tab", { name: "Health Status" }).click(); + await expect( + page.getByRole("heading", { name: "Model Health Status" }), + ).toBeVisible({ timeout: 15_000 }); +} + +async function expectStatus( + page: PlaywrightPage, + modelName: string, + status: string, +): Promise { + const row = await locateHealthRow(page, modelName); + await expect(row, `${modelName} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await expect( + row.getByText(status, { exact: true }), + `the Health Status cell for ${modelName} reads ${status}`, + ).toHaveCount(1, { timeout: 60_000 }); +} + +async function deleteDeployment( + page: PlaywrightPage, + id: string, +): Promise { + const post = () => + page.request.post("/model/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { id }, + }); + const deleted = await post().catch(() => post()); + expect( + deleted.ok(), + `cleanup: /model/delete ${id} returned ${deleted.status()}`, + ).toBe(true); +} + +const uniqueSuffix = (): string => + `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +async function withDeployment( + page: PlaywrightPage, + prefix: string, + apiBase: string, + use: (name: string) => Promise, +): Promise { + const name = `${prefix}-${uniqueSuffix()}`; + const created = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: name, + litellm_params: { + model: `openai/${name}`, + api_base: apiBase, + api_key: "fake-key", + }, + model_info: {}, + }, + }); + expect( + created.ok(), + `/model/new for ${name} failed: ${created.status()} ${await created.text()}`, + ).toBe(true); + const id = (await created.json()).model_info?.id; + expect(id, `model id from /model/new for ${name}`).toBeTruthy(); + try { + await expect + .poll(() => isRegistered(page, name), { + message: `deployment ${name} never appeared in /v2/model/info after create`, + timeout: 60_000, + }) + .toBe(true); + await use(name); + } finally { + await deleteDeployment(page, id); + } +} + +const test = base.extend<{ reachableName: string; unreachableName: string }>({ + reachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-up", MOCK_LLM_BASE, use); + }, + unreachableName: async ({ page }, use) => { + await withDeployment(page, "e2e-health-down", UNREACHABLE_BASE, use); + }, +}); + +test.describe("Model health status", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Run Health Check reports a reachable deployment healthy and an unreachable one unhealthy", async ({ + page, + reachableName, + unreachableName, + }) => { + await navigateToPage(page, Page.Models); + await openHealthTab(page); + + for (const name of [reachableName, unreachableName]) { + const row = await locateHealthRow(page, name); + await expect(row, `${name} has one row in the health table`).toHaveCount( + 1, + { timeout: 20_000 }, + ); + await row + .getByRole("button", { name: "Run Health Check", exact: true }) + .click(); + } + + await expectStatus(page, reachableName, "healthy"); + await expect( + healthRow(page, reachableName).getByText("unhealthy", { exact: true }), + "a reachable deployment is never reported unhealthy", + ).toHaveCount(0); + await expectStatus(page, unreachableName, "unhealthy"); + + const successDetail = ( + await locateHealthRow(page, reachableName) + ).getByRole("button", { + name: "View response details", + }); + await expect( + successDetail, + `${reachableName} offers its health check response for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await successDetail.click(); + const successDialog = page.getByRole("dialog"); + await expect( + successDialog.getByRole("heading", { + name: `Health Check Response - ${reachableName}`, + }), + "the healthy deployment's detail opens its own response dialog", + ).toBeVisible({ timeout: 10_000 }); + await successDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(successDialog).toBeHidden({ timeout: 10_000 }); + + const errorDetail = ( + await locateHealthRow(page, unreachableName) + ).getByRole("button", { + name: "View full error details", + }); + await expect( + errorDetail, + `${unreachableName} offers its health check error for inspection`, + ).toBeVisible({ timeout: 60_000 }); + await errorDetail.click(); + const errorDialog = page.getByRole("dialog"); + await expect( + errorDialog.getByRole("heading", { + name: `Health Check Error - ${unreachableName}`, + }), + "the unreachable deployment's detail opens its own error dialog", + ).toBeVisible({ timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog carries the upstream connection failure, not a generic message", + ).toContainText(/connection error/i, { timeout: 10_000 }); + await expect( + errorDialog, + "the error dialog names the endpoint that could not be reached", + ).toContainText(UNREACHABLE_BASE); + await errorDialog.getByRole("button", { name: "Close" }).last().click(); + await expect(errorDialog).toBeHidden({ timeout: 10_000 }); + + await page.reload(); + await openHealthTab(page); + await expectStatus(page, reachableName, "healthy"); + await expectStatus(page, unreachableName, "unhealthy"); + }); +}); From b3151073d2a8ee274fd067c1d4cfb17da0a82459 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:49:54 -0700 Subject: [PATCH 15/19] test(e2e/ui): cover key budget window, non-admin model scope edit, and key blocking (#40027) * test(e2e/ui): cover key budget window, non-admin model scope edit, and key blocking Three Playwright specs for the Virtual Keys flows customers hit most, each reading its result back through /key/info and /v1/chat/completions rather than trusting the toast: - a monthly spend cap and reset window set through Edit Settings, surviving a reload, with clearing the window leaving the cap in place - a team member narrowing their own team key's models, and the proxy refusing the model they dropped - blocking a key from its detail page, then unblocking it Each test owns the key it edits and deletes it on teardown, so retries and --repeat-each never run out of fixtures. * test(e2e/ui): tighten virtual key specs from review feedback Replace the mutable suite-level key state with a Playwright fixture, so the alias and token are never reassigned and cleanup stays tied to the test. Assert /key/delete succeeded instead of discarding the response, so a failed cleanup surfaces rather than leaving rows behind. Drop the explanatory JSDoc the repo's comment policy disallows, keeping only the one line explaining why Date.now() alone is not unique enough. Type the master-key POST helper against a real guard instead of casting to Record. Assert the unblocked key is served with a 200, not just the response text, and that clearing the reset window also clears budget_reset_at. * test(ui): assert the team response through Playwright --- tests/e2e/ui/helpers/navigation.ts | 10 + tests/e2e/ui/helpers/traffic.ts | 51 ++++- .../internalUserKeyScope.spec.ts | 208 ++++++++++++++++++ .../ui/tests/proxy-admin/keyBlocking.spec.ts | 112 ++++++++++ .../tests/proxy-admin/keyBudgetWindow.spec.ts | 101 +++++++++ 5 files changed, 478 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts create mode 100644 tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 4a7c4e7baa9..e0e7b4da396 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -73,3 +73,13 @@ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise await cell.click(); await expect(page.getByText("Back to Teams")).toBeVisible({ timeout: 10_000 }); } + +export async function openKeyDetail(page: PlaywrightPage, alias: string): Promise { + await page.getByPlaceholder("Search by key alias or ID").fill(alias); + const row = page.getByRole("row").filter({ hasText: alias }); + await expect(row, `key row "${alias}" never appeared on the Virtual Keys page`).toBeVisible({ timeout: 15_000 }); + await row.getByRole("button", { name: alias }).click(); + await expect(page.getByText("Back to Keys"), `key detail for "${alias}" never opened`).toBeVisible({ + timeout: 15_000, + }); +} diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 7f8417cdffb..cb68747b364 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -1,4 +1,4 @@ -import { APIRequestContext, expect } from "@playwright/test"; +import { APIRequestContext, APIResponse, expect } from "@playwright/test"; /** Model names served by fixtures/config.yml, both backed by the mock LLM server. */ export const CHAT_MODEL_A = "fake-openai-gpt-4"; @@ -15,6 +15,9 @@ export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-123 export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; +/** Date.now() alone collides: `--repeat-each` starts its copies inside the same millisecond. */ +export const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + interface ChatOptions { model: string; prompt: string; @@ -25,9 +28,8 @@ interface ChatOptions { traceId?: string; } -/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ -export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { - const res = await request.post(`${rootPath()}/v1/chat/completions`, { +const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise => + request.post(`${rootPath()}/v1/chat/completions`, { headers: { Authorization: `Bearer ${opts.apiKey ?? masterKey()}`, "Content-Type": "application/json", @@ -39,12 +41,26 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}), }, }); + +/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ +export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); const body = await res.json(); expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); return body.id as string; } +export interface ChatAttempt { + status: number; + body: string; +} + +export async function attemptChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); + return { status: res.status(), body: await res.text() }; +} + /** `key` is the sk- value to authenticate with; `token` is its hash, which spend aggregates are keyed by. */ export async function createVirtualKey( request: APIRequestContext, @@ -66,6 +82,33 @@ export async function createVirtualKey( }; } +export interface KeyInfo { + key_alias: string | null; + max_budget: number | null; + budget_duration: string | null; + budget_reset_at: string | null; + blocked: boolean | null; + models: string[]; + team_id: string | null; +} + +export async function readKeyInfo(request: APIRequestContext, token: string): Promise { + const res = await request.get(`${rootPath()}/key/info?key=${encodeURIComponent(token)}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /key/info for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); + const body = await res.json(); + return body.info as KeyInfo; +} + +export async function deleteVirtualKey(request: APIRequestContext, token: string): Promise { + const res = await request.post(`${rootPath()}/key/delete`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { keys: [token] }, + }); + expect(res.ok(), `key delete for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true); +} + /** Spend logs are flushed on a timer, so an assertion straight after a completion races the writer. */ export async function waitForSpendLog( request: APIRequestContext, diff --git a/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts new file mode 100644 index 00000000000..f923841257a --- /dev/null +++ b/tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts @@ -0,0 +1,208 @@ +import { test, expect, type APIRequestContext } from "@playwright/test"; +import { Page } from "../../fixtures/pages"; +import { + dismissFeedbackPopup, + navigateToPage, + openKeyDetail, +} from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + masterKey, + readKeyInfo, + rootPath, + uniqueSuffix, +} from "../../helpers/traffic"; + +const MEMBER_PASSWORD = "E2e-Team-Member-Pass-1!"; + +interface CreatedTeam { + readonly team_id: string; +} + +function assertCreatedTeam(body: unknown): asserts body is CreatedTeam { + expect(body, "/team/new returned no team_id").toMatchObject({ + team_id: expect.any(String), + }); +} + +async function postAsMaster( + request: APIRequestContext, + path: string, + data: Record, +): Promise { + const res = await request.post(`${rootPath()}${path}`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data, + }); + expect( + res.ok(), + `POST ${path} failed (${res.status()}): ${await res.text()}`, + ).toBe(true); + return res.json(); +} + +test.describe("Internal User - own team key model scope", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("a team member narrows their own key's models and the proxy enforces it", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const email = `team-member-${suffix}@test.local`; + const userId = `e2e-key-scope-user-${suffix}`; + const alias = `e2e-key-scope-${suffix}`; + + const team = await postAsMaster(request, "/team/new", { + team_alias: `E2E Key Scope ${suffix}`, + models: [CHAT_MODEL_A, CHAT_MODEL_B], + team_member_permissions: ["/key/generate", "/key/update", "/key/info"], + }); + assertCreatedTeam(team); + const teamId = team.team_id; + + try { + await postAsMaster(request, "/user/new", { + user_id: userId, + user_email: email, + user_role: "internal_user", + auto_create_key: false, + }); + await postAsMaster(request, "/user/update", { + user_id: userId, + password: MEMBER_PASSWORD, + }); + await postAsMaster(request, "/team/member_add", { + team_id: teamId, + member: { role: "user", user_id: userId }, + }); + + const created = await createVirtualKey(request, { + key_alias: alias, + team_id: teamId, + user_id: userId, + models: [], + }); + + try { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page + .getByPlaceholder("Enter your password") + .fill(MEMBER_PASSWORD); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect( + page.locator("a", { hasText: "Virtual Keys" }), + `${email} never reached the dashboard`, + ).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await expect( + page.getByRole("option", { name: CHAT_MODEL_A, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_A} to a team member`, + ).toBeVisible({ timeout: 15_000 }); + await expect( + page.getByRole("option", { name: CHAT_MODEL_B, exact: true }), + `the Models dropdown does not offer ${CHAT_MODEL_B} to a team member`, + ).toBeVisible(); + + await page + .getByRole("option", { name: CHAT_MODEL_A, exact: true }) + .click(); + await page.keyboard.press("Escape"); + + const updated = page.waitForResponse( + (res) => + res.url().includes("/key/update") && + res.request().method() === "POST", + ); + await page.getByRole("button", { name: "Save Changes" }).click(); + const updateStatus = (await updated).status(); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeGreaterThanOrEqual(200); + expect( + updateStatus, + "a team member's own-key edit was refused", + ).toBeLessThan(300); + await expect( + page.getByText("Key updated successfully").first(), + ).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => (await readKeyInfo(request, created.token)).models, + { + message: `the narrowed model scope never reached /key/info for ${alias}`, + timeout: 20_000, + }, + ) + .toEqual([CHAT_MODEL_A]); + + await expect + .poll( + async () => + await attemptChatCompletion(request, { + model: CHAT_MODEL_B, + prompt: `out of scope ${suffix}`, + apiKey: created.key, + }), + { + message: `${CHAT_MODEL_B} was still served after the key was narrowed to ${CHAT_MODEL_A}`, + timeout: 30_000, + }, + ) + .toMatchObject({ + status: 403, + body: expect.stringContaining(CHAT_MODEL_B), + }); + + const inScope = await attemptChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `in scope ${suffix}`, + apiKey: created.key, + }); + expect( + inScope, + `${CHAT_MODEL_A} is no longer served by the narrowed key`, + ).toMatchObject({ + status: 200, + body: expect.stringContaining(MOCK_RESPONSE_TEXT), + }); + } finally { + await deleteVirtualKey(request, created.token); + } + } finally { + await request.post(`${rootPath()}/user/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { user_ids: [userId] }, + }); + await request.post(`${rootPath()}/team/delete`, { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { team_ids: [teamId] }, + }); + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts new file mode 100644 index 00000000000..99a8065a797 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts @@ -0,0 +1,112 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + attemptChatCompletion, + createVirtualKey, + deleteVirtualKey, + readKeyInfo, + sendChatCompletion, + uniqueSuffix, +} from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; + apiKey: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-block-key-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token, apiKey: created.key }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key blocking", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("blocking a key stops it serving and unblocking restores it", async ({ page, scopedKey }) => { + const { alias, token, apiKey } = scopedKey; + + await sendChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: `pre-block ${alias}`, + apiKey, + }); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Block Key" }).click(); + const blockDialog = page.getByRole("dialog", { name: "Block Key" }); + await expect(blockDialog, "the Block Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await blockDialog.getByRole("button", { name: "Block", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back blocked from /key/info", + timeout: 20_000, + }) + .toBe(true); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "blocked", + apiKey, + }), + { + message: "a blocked key was still served by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 401, body: expect.stringContaining("blocked") }); + + await page.reload(); + await expect( + page.getByText("Blocked", { exact: true }), + "the reloaded key detail does not show the key as blocked", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("button", { name: "More key actions" }).click(); + await page.getByRole("menuitem", { name: "Unblock Key" }).click(); + const unblockDialog = page.getByRole("dialog", { name: "Unblock Key" }); + await expect(unblockDialog, "the Unblock Key confirmation never opened").toBeVisible({ timeout: 10_000 }); + await unblockDialog.getByRole("button", { name: "Unblock", exact: true }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).blocked, { + message: "the key never came back unblocked from /key/info", + timeout: 20_000, + }) + .toBe(false); + + await expect + .poll( + async () => + await attemptChatCompletion(page.request, { + model: CHAT_MODEL_A, + prompt: "unblocked", + apiKey, + }), + { + message: "an unblocked key is still refused by /v1/chat/completions", + timeout: 30_000, + }, + ) + .toMatchObject({ status: 200, body: expect.stringContaining(MOCK_RESPONSE_TEXT) }); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts new file mode 100644 index 00000000000..4e4d0a395c3 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -0,0 +1,101 @@ +import { test as base, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, createVirtualKey, deleteVirtualKey, readKeyInfo, uniqueSuffix } from "../../helpers/traffic"; + +interface ScopedKey { + alias: string; + token: string; +} + +const test = base.extend<{ scopedKey: ScopedKey }>({ + scopedKey: async ({ page }, use) => { + const alias = `e2e-budget-window-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + team_id: E2E_TEAM_CRUD_ID, + models: [CHAT_MODEL_A], + }); + await use({ alias, token: created.token }); + await deleteVirtualKey(page.request, created.token); + }, +}); + +test.describe("Proxy Admin - Key budget window", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => { + const { alias, token } = scopedKey; + + const before = await readKeyInfo(page.request, token); + expect(before.max_budget, "a freshly generated key starts with no budget").toBeNull(); + + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "monthly", exact: true }).click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { + message: "the $12.50 cap never reached /key/info", + timeout: 20_000, + }) + .toBe(12.5); + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the monthly reset window never reached /key/info", + timeout: 20_000, + }) + .toBe("30d"); + + const capped = await readKeyInfo(page.request, token); + const resetAt = new Date(capped.budget_reset_at ?? ""); + expect(Number.isNaN(resetAt.getTime()), "a monthly window left the key with no budget_reset_at").toBe(false); + expect(resetAt.getTime(), "budget_reset_at was set in the past").toBeGreaterThan(Date.now()); + expect(resetAt.getUTCDate(), "a monthly window resets on the 1st, a daily one would not").toBe(1); + + await page.reload(); + await expect( + page.getByRole("paragraph").filter({ hasText: "of $12.50" }), + "the reloaded key detail does not render the $12.50 cap", + ).toBeVisible({ timeout: 15_000 }); + + await page.getByRole("tab", { name: "Settings" }).click(); + await expect( + page.getByTestId("budget-reset-value"), + "the reloaded key detail does not name the 30d reset window", + ).toHaveText(/Every 30d/, { timeout: 15_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByLabel("Reset Budget", { exact: true }).click(); + await page.getByRole("option", { name: "Never resets", exact: true }).click(); + + const cleared = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(cleared).toHaveProperty("budget_duration"); + expect(cleared.budget_duration, "clearing the window must send budget_duration: null explicitly").toBeNull(); + + await expect + .poll(async () => (await readKeyInfo(page.request, token)).budget_duration, { + message: "the reset window was never cleared on /key/info", + timeout: 20_000, + }) + .toBeNull(); + + const after = await readKeyInfo(page.request, token); + expect(after.budget_reset_at, "clearing the reset window left a stale next-reset timestamp").toBeNull(); + expect(after.max_budget, "clearing the reset window also wiped the spend cap").toBe(12.5); + expect(after.models, "editing the budget left the key's models untouched").toEqual(before.models); + expect(after.team_id, "editing the budget left the key's team untouched").toEqual(before.team_id); + }); +}); From 36bd7f113837caadd3f3d400dd246cd4649b7b33 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 22:50:13 -0700 Subject: [PATCH 16/19] fix(mcp): honor an explicit null on toolset update, cover MCP lifecycle e2e (#40022) * fix(mcp): honor an explicit null on toolset update, cover MCP lifecycle e2e PUT /v1/mcp/toolset dumped its payload with exclude_none, so a field sent as null looked exactly like one the caller left out and the stored value survived. An admin could not clear a toolset's description: the save reported success and the old text came straight back. It now dumps with exclude_unset, so absent keeps and null clears, which is what PUT /v1/mcp/server already did. A null tools list clears the selection to empty, and a null toolset_name is ignored because a toolset always has a name. Adds create, read, partial-update, clear and delete e2e coverage for MCP servers and toolsets, with every read-back polled on every replica so an edit that lands on one replica and not another fails the test, plus an enforcement test proving a key granted a toolset lists exactly that toolset's tools against the real Datadog upstream. * fix(e2e): refuse a read-back that no replica serves A read-back over an empty replica mapping satisfied every predicate and returned as if it had converged, so it would have asserted nothing and passed. No wiring can produce that today, since the replica list always falls back to at least one URL, but a helper whose whole job is proving a write reached every replica should not have a shape that passes vacuously. * fix(mcp): keep a null tools list a no-op on toolset update Treating a null tools list as a clear meant an existing client that sends tools=null during a partial update, meaning "leave the selection alone", silently lost every tool the toolset grants. That is a permission surface, so the quiet version of it is the worst version. A toolset always has a tool list, the same way it always has a name, so a null on either is now a no-op. Emptying the selection is an explicit [], which cannot be confused with a field the caller left out, and which is what the dashboard already sends. * fix(e2e): keep MCP admin routes on the data plane /v1/mcp/* is a lazily mounted feature, so a gateway registers it on the first matching request, which happens after the startup route trim that drops management endpoints. Routing it to the control plane therefore sent every MCP call to the one backend process: the new lifecycle read-backs proved a single process rather than every replica, and mcp_client's await_registered barrier waited on a registry that does not serve the tools/list call it guards, so the existing MCP suites polled a gateway that had not synced yet until poll_timeout Verified against a two-gateway split stack (backend on 4001, gateways on 4010 and 4011, one postgres): both gateways answer /v1/mcp/server and /v1/mcp/toolset, and each served 6 server reads and 7 toolset reads over the run * fix(e2e): grant the toolset by the tool's own name, not the wire name tools/list serves a tool as , but a toolset grants by the tool's own name: resolve_toolset_permissions reads toolset.tools[].tool_name straight through, and the prefix is added on the way out. The test built the toolset from the names tools/list reported, so the grant matched nothing, the scoped key listed no tools, and await_tools ran out its whole poll_timeout before failing Measure the prefix off search_datadog_logs, whose own name is known, rather than guessing it from the alias, since the proxy can be configured to prefix with a short server id instead. The expectation compared against tools/list stays in wire names; only what the toolset stores crosses back * test(mcp): build immutable lifecycle updates and replica results * test: validate opaque stream IDs and hide log-reader credentials * test: isolate auto-router scenarios and clean partial setup * test: honor Datadog search rate-limit reset headers * test: share the Datadog read-back deadline across retries * test: preserve captured MCP toolset update fields --- .../_experimental/mcp_server/toolset_db.py | 16 +- .../mcp_management_endpoints.py | 4 + tests/e2e/coverage_registry/mcp.yaml | 8 + tests/e2e/coverage_registry/mgmt.yaml | 10 + tests/e2e/e2e_http.py | 64 +++- .../test_responses_bridge_streaming_e2e.py | 24 +- tests/e2e/logging/datadog_reader.py | 124 +++++--- tests/e2e/logging/test_datadog_reader.py | 223 +++++++++++++ tests/e2e/management/management_client.py | 35 +++ .../e2e/management/test_mcp_lifecycle_e2e.py | 294 ++++++++++++++++++ tests/e2e/mcp/datadog_mcp.py | 7 +- tests/e2e/mcp/mcp_client.py | 39 ++- .../mcp/test_mcp_toolset_enforcement_e2e.py | 95 ++++++ tests/e2e/models.py | 73 ++++- tests/e2e/proxy_client.py | 216 ++++++++++++- .../test_auto_router_regressions_e2e.py | 239 +++++++------- tests/e2e/test_e2e_http.py | 91 +++++- tests/e2e/test_proxy_client.py | 95 +++++- .../mcp_server/test_mcp_partial_update.py | 71 ++++- 19 files changed, 1501 insertions(+), 227 deletions(-) create mode 100644 tests/e2e/logging/test_datadog_reader.py create mode 100644 tests/e2e/management/test_mcp_lifecycle_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index ecaaf35e817..48bad178927 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -132,10 +132,18 @@ async def update_mcp_toolset( data: UpdateMCPToolsetRequest, touched_by: str, ) -> MCPToolset | None: - data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"}) - if "tools" in data_dict: - data_dict["tools"] = json.dumps(data_dict["tools"]) - data_dict["updated_by"] = touched_by + """A partial update: absent keeps, null clears. A toolset always has a name and a + tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear; + emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a + caller that left the field out.""" + data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization + ( + (field, json.dumps(value) if field == "tools" else value) + for field, value in data.model_dump(exclude_unset=True).items() + if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) + ), + updated_by=touched_by, + ) try: row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d5c3427f29a..2aa7fdc7393 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2673,6 +2673,8 @@ if MCP_AVAILABLE: """ Updates the MCP Server in the db. + Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared. + Parameters: - payload: UpdateMCPServerRequest - Required. The updated mcp server data. ``` @@ -3098,6 +3100,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: str | None = Header(None), ): + """Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except + ``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit [].""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..f853e9ff8d6 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,11 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "user_api_key_auth_mcp.py:2137" + rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 860d96a50b4..c8d7037d2fd 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,3 +76,13 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} +- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} +- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} +- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"} +- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"} +- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"} +- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"} +- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} +- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} +- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..415c72bbb3c 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders): anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version") +class PartialBody(BaseModel): + """A body for a partial-update route (absent = keep, null = clear): a field left + unset is omitted from the wire, and a field set to None is sent as JSON null.""" + + class NoBody(BaseModel): """Empty body/query for routes that take none.""" @@ -252,6 +257,13 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) + +def wire_body(json: BaseModel) -> dict[str, object]: + if isinstance(json, PartialBody): + return json.model_dump(by_alias=True, exclude_unset=True) + return json.model_dump(by_alias=True, exclude_none=True) + + def _headers(headers: BaseModel) -> dict[str, str]: dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} @@ -307,9 +319,26 @@ def request_with_retry[T: RetryableResponse]( return issue() -def _classify[R: BaseModel]( - resp: requests.Response, response_type: type[R] -) -> Result[R]: +class ClassifiableResponse(Protocol): + """What classifying an outcome reads off a response. requests.Response satisfies + it, and so does a fake, so the classification rules are testable on their own.""" + + @property + def status_code(self) -> int: ... + + @property + def ok(self) -> bool: ... + + @property + def text(self) -> str: ... + + @property + def content(self) -> bytes: ... + + def json(self) -> object: ... + + +def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]: if resp.status_code == 401: return UnauthorizedError(body=resp.text) if resp.status_code == 429: @@ -317,7 +346,8 @@ def _classify[R: BaseModel]( if not resp.ok: return UnknownApiError(status_code=resp.status_code, body=resp.text) try: - return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json())) + payload: Final[object] = resp.json() if resp.content else {} + return Success(status_code=resp.status_code, data=response_type.model_validate(payload)) except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value return ValidationError(message=str(exc)) @@ -335,13 +365,13 @@ def post[R: BaseModel]( lambda: requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get[R: BaseModel]( @@ -363,7 +393,7 @@ def get[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get_external[R: BaseModel]( @@ -383,7 +413,7 @@ def get_external[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def delete[R: BaseModel]( @@ -400,14 +430,14 @@ def delete[R: BaseModel]( lambda: requests.delete( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), params=_params(params), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def patch[R: BaseModel]( @@ -423,13 +453,13 @@ def patch[R: BaseModel]( lambda: requests.patch( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def put[R: BaseModel]( @@ -445,13 +475,13 @@ def put[R: BaseModel]( lambda: requests.put( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def probe( @@ -555,7 +585,7 @@ def send( str(url), headers=_headers(headers), params=_params(params), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=stream, timeout=timeout, ) @@ -605,7 +635,7 @@ def upload[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def stream_binary( @@ -623,7 +653,7 @@ def stream_binary( resp = requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=True, timeout=timeout, ) 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..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 dataclasses import dataclass +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,17 +31,33 @@ 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): - 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 +108,12 @@ 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) + 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 new file mode 100644 index 00000000000..910a1cefd42 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,223 @@ +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: + 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, + } + + +@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 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 2b897f5f07f..1ef0d89a8f9 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -43,6 +43,9 @@ from models import ( KeyResetSpendBody, KeyResetSpendResponse, KeyUpdateBody, + McpServerCreateBody, + McpServerRow, + McpServerUpdateBody, ModelDeleteBody, OrgDeleteBody, OrgInfoParams, @@ -537,6 +540,38 @@ class ManagementClient: ).root ) + def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow: + return unwrap( + self.proxy.transport.post( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow: + """PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial + update where a field left unset keeps its stored value and None clears it.""" + return unwrap( + self.proxy.transport.put( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def delete_mcp_server(self, server_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can + unwrap it while a deferred teardown can ignore an already-deleted server.""" + return self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", diff --git a/tests/e2e/management/test_mcp_lifecycle_e2e.py b/tests/e2e/management/test_mcp_lifecycle_e2e.py new file mode 100644 index 00000000000..9257d697647 --- /dev/null +++ b/tests/e2e/management/test_mcp_lifecycle_e2e.py @@ -0,0 +1,294 @@ +"""Live e2e: the MCP server and toolset management routes' lifecycle contract. + +Two customer defects sit on these routes, and each step here is the read-back that +would have caught one of them: a dashboard edit that took several saves to stick +because the read landed on a replica the write had not reached, and a toolset whose +tools were stored under one name and read back under another, so it granted +nothing. Every read-back therefore polls every replica that serves the route +(ProxyClient.read_back_everywhere) and asserts the exact values written, and both +update routes are held to the same partial-update contract: a field left out of the +payload keeps its stored value, a field sent as null is cleared. The server URL is +unreachable on purpose; only persistence is under test, never a tool call. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + McpInfo, + McpServerCreateBody, + McpServerListResponse, + McpServerRow, + McpServerUpdateBody, + ToolsetCreateBody, + ToolsetListResponse, + ToolsetRow, + ToolsetTool, + ToolsetUpdateBody, +) + +pytestmark = pytest.mark.e2e + +UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp" + + +def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]: + name: Final = f"e2e_mcp_lifecycle_{unique_marker()}" + body: Final = McpServerCreateBody( + server_name=name, + alias=name, + url=UNREACHABLE_URL, + transport="http", + description="e2e lifecycle server", + mcp_info=McpInfo( + server_name=f"{name} (display)", + description="shown on the MCP page", + logo_url="https://e2e.test.local/logo.png", + ), + ) + server_id: Final = client.create_mcp_server(body).server_id + resources.defer(lambda: client.delete_mcp_server(server_id)) + return body, server_id + + +def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None: + stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info) + expected: Final = ( + written.server_name, + written.alias, + written.url, + written.transport, + written.description, + written.mcp_info, + ) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _server_everywhere( + client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool] +) -> Mapping[str, McpServerRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled) + + +def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]: + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: any(row.server_id == server_id for row in rows.root), + ) + return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()} + + +class TestMcpServerLifecycle: + @pytest.mark.covers("mgmt.mcp_server.new.persists") + def test_create_persists_every_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id) + for replica, row in by_id.items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}") + + @pytest.mark.skip( + reason=( + "product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose " + "_build_mcp_server_table sets description from mcp_info['description'], so the list " + "reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the " + "stored description column. A server created with both set to different text reads " + "back with two different descriptions depending on the route" + ) + ) + @pytest.mark.covers("mgmt.mcp_server.list.persists") + def test_created_server_is_listed_with_every_field( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + for replica, row in _listed_server_everywhere(client, server_id).items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}") + + @pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields") + def test_updating_only_the_alias_keeps_every_other_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + renamed: Final = f"{body.alias}_renamed" + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed)) + + after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed) + for replica, row in after_one_put.items(): + _assert_server_matches( + row, + body.model_copy(update={"alias": renamed}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias", + ) + + @pytest.mark.covers("mgmt.mcp_server.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None)) + + cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_server_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_server.delete.persists") + def test_delete_removes_the_server_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + + _ = unwrap(client.delete_mcp_server(server_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}") + assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: all(row.server_id != server_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.server_id != server_id for row in rows.root), ( + f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}" + ) + + +def _create_toolset( + client: ManagementClient, resources: ResourceManager, server_id: str +) -> tuple[ToolsetCreateBody, str]: + body: Final = ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="e2e lifecycle toolset", + tools=[ + ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"), + ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"), + ], + ) + toolset_id: Final = client.proxy.create_toolset(body).toolset_id + resources.defer(lambda: client.proxy.delete_toolset(toolset_id)) + return body, toolset_id + + +def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None: + stored: Final = (row.toolset_name, row.description, row.tools) + expected: Final = (written.toolset_name, written.description, written.tools) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _toolset_everywhere( + client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool] +) -> Mapping[str, ToolsetRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled) + + +class TestMcpToolsetLifecycle: + @pytest.mark.covers("mgmt.mcp_toolset.new.persists") + def test_create_persists_both_tools_under_the_exact_names_written( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id) + for replica, row in by_id.items(): + _assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}") + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + _assert_toolset_matches( + next(row for row in rows.root if row.toolset_id == toolset_id), + body, + where=f"GET /v1/mcp/toolset on {replica}", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields") + def test_updating_only_the_description_keeps_the_tools_and_name( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited")) + + edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited") + for replica, row in edited.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": "edited"}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.persists") + def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + kept: Final = body.tools[:1] + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept)) + + narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept) + for replica, row in narrowed.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"tools": kept}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None)) + + cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.delete.persists") + def test_delete_removes_the_toolset_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + _, toolset_id = _create_toolset(client, resources, server_id) + + _ = unwrap(client.proxy.delete_toolset(toolset_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}") + assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.toolset_id != toolset_id for row in rows.root), ( + f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}" + ) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..352b4446cfd 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the core Datadog toolset with its credentials from the env. By default + the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose + every tool the core toolset serves.""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=None if allowed_tools is None else list(allowed_tools), mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..210fc7a1e98 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -16,11 +16,11 @@ import time from collections.abc import Mapping from dataclasses import dataclass -from pydantic import BaseModel, ConfigDict, Field, RootModel +from pydantic import BaseModel, ConfigDict, Field from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel): server_id: str -class McpServerRow(BaseModel): - server_id: str - alias: str | None = None - url: str | None = None - - -class McpServersListResponse(RootModel[list[McpServerRow]]): - pass - - class McpToolMcpInfo(BaseModel): server_id: str | None = None alias: str | None = None @@ -193,7 +183,7 @@ class McpClient: "/v1/mcp/server", headers=self.proxy.transport.master, params=NoBody(), - response_type=McpServersListResponse, + response_type=McpServerListResponse, ) ).root @@ -224,11 +214,16 @@ class McpClient: user_id: str, mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, + mcp_toolsets: list[str] | None = None, models: list[str] | None = None, ) -> str: object_permission = ( - ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) - if mcp_servers is not None or mcp_access_groups is not None + ObjectPermission( + mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, + mcp_toolsets=mcp_toolsets, + ) + if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None else None ) return self.proxy.generate_key( @@ -272,6 +267,20 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]: + """Poll tools/list until `server_id`'s tools as `key` sees them are exactly + `expected`, and return the last listing either way, so the caller's equality + assertion names the difference. Fails at poll_timeout only when the read + itself never succeeded.""" + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected: + return expected + if time.monotonic() >= deadline: + return unwrap(result).tool_names_for_server(server_id) + time.sleep(self.proxy.poll_interval) + def await_call_tool( self, key: str, diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..6b901145eb1 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,95 @@ +"""Live e2e: a key granted a toolset lists exactly the toolset's tools. + +An admin registers the real Datadog remote MCP server with its whole core toolset +exposed, discovers two of its tool names through a key granted the server outright, +and curates a toolset naming exactly those two. A second key is granted the server +plus that toolset, and its tools/list must come back as exactly those two names: no +more, so the rest of the server's catalog stays hidden behind the toolset, and no +fewer, so a tool stored under one name and read under another (which granted +nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP +upstream). +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ToolsetCreateBody, ToolsetTool + +pytestmark = pytest.mark.e2e + + +def _key( + client: McpClient, + resources: ResourceManager, + label: str, + *, + server_id: str, + toolset_id: str | None = None, +) -> str: + key: Final = client.generate_key( + user_id=f"e2e-mcp-{label}-{unique_marker()}", + mcp_servers=[server_id], + mcp_toolsets=None if toolset_id is None else [toolset_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _wire_prefix(wire_name: str, tool_name: str, catalog: frozenset[str]) -> str: + """The prefix tools/list puts in front of one server's tool names, measured off a + tool whose own name is known rather than guessed from the alias. A toolset grants + by the tool's own name, never the wire name, and the prefix is whatever the proxy + is configured to build (the alias, or a short server id), so measuring it is the + only way to cross between the two.""" + assert wire_name.endswith(tool_name), f"tools/list served {wire_name!r}, expected it to end with {tool_name!r}" + prefix: Final = wire_name[: len(wire_name) - len(tool_name)] + unprefixed: Final = frozenset(name for name in catalog if not name.startswith(prefix)) + assert not unprefixed, ( + f"every tool of one server shares the wire prefix {prefix!r}, so {sorted(unprefixed)} " + f"cannot be reduced to the names a toolset grants by" + ) + return prefix + + +class TestMcpToolsetEnforcement: + @pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped") + def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None: + server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + + catalog_key: Final = _key(client, resources, "catalog", server_id=server_id) + known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL) + catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id) + assert len(catalog) > 2, ( + f"the Datadog core toolset must serve more tools than the toolset names, or the " + f"restriction has nothing to hide; got {sorted(catalog)}" + ) + prefix: Final = _wire_prefix(known_wire, SEARCH_LOGS_TOOL, catalog) + chosen_wire: Final = frozenset(sorted(catalog)[:2]) + chosen: Final = frozenset(name.removeprefix(prefix) for name in chosen_wire) + + toolset: Final = client.proxy.create_toolset( + ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="two Datadog tools", + tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)], + ) + ) + resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id)) + assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, ( + f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim" + ) + + scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id) + listed: Final = client.await_tools(scoped_key, server_id, expected=chosen_wire) + assert listed == chosen_wire, ( + f"a key granted the toolset must list exactly its two tools; " + f"got {sorted(listed)}, expected {sorted(chosen_wire)}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f2654e0eec..62810e6cfd9 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -10,6 +10,7 @@ from collections.abc import Sequence from datetime import datetime from typing import Final, Literal +from e2e_http import PartialBody from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -55,6 +56,7 @@ class KeyMetadata(BaseModel): class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -77,7 +79,7 @@ class KeyGenerateBody(BaseModel): allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None - router_settings: "RouterSettingsOverride | None" = None + router_settings: RouterSettingsOverride | None = None class KeyGenerateResponse(BaseModel): @@ -516,6 +518,15 @@ class CountTokensResponse(BaseModel): # ---------- mcp servers ---------- +class McpInfo(BaseModel): + """The `mcp_info` display block stored on an MCP server; only the fields the + lifecycle test writes and reads back.""" + + server_name: str | None = None + description: str | None = None + logo_url: str | None = None + + class McpServerCreateBody(BaseModel): """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints @@ -530,6 +541,18 @@ class McpServerCreateBody(BaseModel): oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None authorization_url: str | None = None token_url: str | None = None + server_name: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerUpdateBody(PartialBody): + """PUT /v1/mcp/server: a field left unset keeps its stored value, a field set + to None is cleared.""" + + server_id: str + alias: str | None = None + description: str | None = None class McpServerInfo(BaseModel): @@ -543,6 +566,54 @@ class McpServerInfo(BaseModel): allow_all_keys: bool | None = None +class McpServerRow(McpServerInfo): + """A stored MCP server as the create, get, and list routes return it: the + fields the lifecycle test asserts survive the round trip.""" + + server_name: str | None = None + transport: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerListResponse(RootModel[list[McpServerRow]]): + """GET /v1/mcp/server answers with a bare array of servers.""" + + +class ToolsetTool(BaseModel): + server_id: str + tool_name: str + + +class ToolsetCreateBody(BaseModel): + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] + + +class ToolsetUpdateBody(PartialBody): + """PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set + to None is cleared.""" + + toolset_id: str + description: str | None = None + tools: list[ToolsetTool] | None = None + + +class ToolsetRow(BaseModel): + """A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id}, + and each row of GET /v1/mcp/toolset return it.""" + + toolset_id: str + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] = Field(default_factory=list) + + +class ToolsetListResponse(RootModel[list[ToolsetRow]]): + """GET /v1/mcp/toolset answers with a bare array of toolsets.""" + + class EmbedBody(BaseModel): model: str input: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 520cbfde5a9..1bac5116a9d 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable, Mapping from dataclasses import dataclass +from functools import reduce from datetime import datetime from types import MappingProxyType from typing import Final @@ -26,6 +27,7 @@ from e2e_http import ( Result, StreamingResponse, Success, + UnknownApiError, is_ok, unwrap, ) @@ -70,6 +72,9 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + ToolsetCreateBody, + ToolsetRow, + ToolsetUpdateBody, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -82,7 +87,7 @@ from e2e_config import ( SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, ) -from transport import HttpTransport, SplitTransport, Transport +from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -235,6 +240,99 @@ def servable_timeout_message( ) +type ReplicaRead[T] = Callable[[float], T] + + +@dataclass(frozen=True, slots=True) +class EverywhereConverged[T]: + """Every replica answered with something `settled` accepts, keyed by replica.""" + + answers: Mapping[str, T] + + +@dataclass(frozen=True, slots=True) +class NeverConvergedOn[T]: + """`replica` ran out its budget without an answer `settled` accepts; `last` is + its final answer, so the failure can say what that replica still serves.""" + + replica: str + last: T + + +def _last_answer[T]( + read: ReplicaRead[T], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> T: + """Poll `read` until `settled` accepts its answer or `timeout` runs out, and + return the last answer either way. Each read's request timeout is clamped to + the budget left, and the final poll runs even when less than an interval + remains, so a deadline never skips the read that would have settled.""" + deadline: Final = now() + timeout + answer = read(min(request_timeout, timeout)) + while not settled(answer): + remaining = deadline - now() + if remaining <= 0: + return answer + sleep(min(interval, remaining)) + answer = read(min(request_timeout, remaining)) + return answer + + +def await_everywhere[T]( + reads: Mapping[str, ReplicaRead[T]], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> EverywhereConverged[T] | NeverConvergedOn[T]: + """`_last_answer` against every replica in turn, each with the full budget, so a + write counts as visible only once the last replica reflects it, and stop at the + first replica that never converges. Clock and sleep are injected.""" + def read_replica( + outcome: EverywhereConverged[T] | NeverConvergedOn[T], + item: tuple[str, ReplicaRead[T]], + ) -> EverywhereConverged[T] | NeverConvergedOn[T]: + if isinstance(outcome, NeverConvergedOn): + return outcome + replica, read = item + answer: Final = _last_answer( + read, + settled=settled, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ) + if not settled(answer): + return NeverConvergedOn(replica=replica, last=answer) + return EverywhereConverged(answers=MappingProxyType({**outcome.answers, replica: answer})) + + initial: Final[EverywhereConverged[T] | NeverConvergedOn[T]] = EverywhereConverged(answers=MappingProxyType({})) + return reduce(read_replica, reads.items(), initial) + + +def _is_not_found[R: BaseModel](result: Result[R]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _status_of[R: BaseModel](result: Result[R]) -> int: + match result: + case Success(status_code=status_code) | UnknownApiError(status_code=status_code): + return status_code + case _: + return -1 + + type Poller[T] = Callable[[], T] @@ -321,6 +419,7 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re class ProxyClient: transport: Transport replicas: Mapping[str, Transport] + control_replicas: Mapping[str, Transport] poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -569,6 +668,112 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2) + # ---- replica read-back ---------------------------------------------- + + def replicas_for(self, path: str) -> Mapping[str, Transport]: + """The replicas that serve `path`: every data-plane replica for an LLM route, + and for a management route the control-plane replicas, since the data-plane + replicas trim management routes and answer them 404. A monolith serves both + from every replica, so a management read-back polls all of them; a split + deployment exposes one control-plane address (there is one backend process + behind it on the stack these suites run against), so it polls that. A + control plane fronting several backends would need its own replica list to + prove each one converged, the way PROXY_REPLICA_URLS does for the gateways. + Never empty: a read-back against no replica would assert nothing and pass.""" + replicas: Final = self.control_replicas if is_control_plane_path(path) else self.replicas + assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing" + return replicas + + def read_body_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, settled: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica that serves it, polling each to poll_timeout + until `settled` accepts its body, and fail naming the first replica that + never converged. Returns each replica's settled body, keyed by replica, so + the caller can assert the rest of it.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()}, + settled=lambda result: isinstance(result, Success) and settled(result.data), + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: unwrap(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; " + f"last read: {last}" + ) + + def gone_everywhere(self, path: str) -> Mapping[str, int]: + """Poll GET `path` on every replica that serves it until each stops serving + it, and fail naming the first replica that still does at poll_timeout. + Returns each replica's final status, so the caller asserts the 404 itself.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()}, + settled=_is_not_found, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: _status_of(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" + ) + + @staticmethod + def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + return lambda request_timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=request_timeout, + ) + + # ---- mcp toolsets --------------------------------------------------- + + def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow: + return unwrap( + self.transport.post( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow: + """PUT /v1/mcp/toolset: a partial update where a field left unset keeps its + stored value and None clears it.""" + return unwrap( + self.transport.put( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def delete_toolset(self, toolset_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase + can unwrap it while a deferred teardown can ignore an already-deleted row.""" + return self.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def create_credential(self, body: CredentialCreateBody) -> None: unwrap( self.transport.post( @@ -736,7 +941,10 @@ def build_proxy_client( base URLs are the same for a monolithic proxy, so routing is then a no-op. ``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model barrier polls directly; it is the data-plane URL itself unless the stack - exports each gateway's own address. + exports each gateway's own address. Management read-backs poll those same + replicas when the two planes share a base URL (a monolith, where every replica + serves every route) and the control plane alone when they differ (a split + deployment, where the data-plane replicas do not serve management routes). The endpoints are injectable for callers that resolve the proxy some other way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must @@ -764,9 +972,13 @@ def build_proxy_client( for url in replica_urls } ) + control_replicas: Final = ( + replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control}) + ) return ProxyClient( transport=split, replicas=replicas, + control_replicas=control_replicas, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) 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=[ diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 66841725d1d..81cd6c8d3d1 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -13,13 +13,24 @@ monkeypatches anything. from __future__ import annotations from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from types import MappingProxyType from typing import Final import pytest - -from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome +from e2e_http import ( + RETRY_ATTEMPTS, + TRANSIENT_STATUSES, + NoBody, + PartialBody, + Success, + ValidationError, + classify, + request_with_retry, + streaming_outcome, + wire_body, +) +from pydantic import BaseModel, TypeAdapter @dataclass @@ -33,10 +44,10 @@ class FakeResponse: @dataclass class SleepRecorder: - delays: list[float] = field(default_factory=list) + delays: tuple[float, ...] = () def __call__(self, seconds: float) -> None: - self.delays.append(seconds) + self.delays += (seconds,) def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]: @@ -55,7 +66,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_429_is_never_retried(self) -> None: @@ -63,7 +74,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_overloaded_529_retries_with_backoff_then_returns_the_success(self) -> None: @@ -71,7 +82,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[1] - assert sleep.delays == [0.5] + assert sleep.delays == (0.5,) assert responses[0].close_calls == 1 assert responses[1].close_calls == 0 @@ -80,7 +91,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[RETRY_ATTEMPTS - 1] - assert sleep.delays == [0.5, 1.0] + assert sleep.delays == (0.5, 1.0) assert [r.close_calls for r in responses] == [1, 1, 0, 0] @@ -134,3 +145,65 @@ class TestStreamEventArrivals: assert result.stream_events == [] assert result.stream_event_arrivals == [] assert result.body == "bad request" + + +class _ServerUpdate(PartialBody): + server_id: str + alias: str | None = None + description: str | None = None + + +class _ServerCreate(BaseModel): + alias: str + description: str | None = None + + +class TestWireBody: + """A partial-update body must put exactly the caller's choice on the wire: an + omitted field stays off it so the route keeps the stored value, and an explicit + None goes out as JSON null so the route clears it. Plain bodies keep dropping + None, which is what every create route expects.""" + + def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None: + assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None} + assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"} + + def test_plain_body_drops_none_fields(self) -> None: + assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"} + + +_JSON: Final[TypeAdapter[object]] = TypeAdapter(object) + + +@dataclass +class FakeJsonResponse: + """The `classify` view of a response: a status, the raw body bytes, and the + parse that would raise on an empty one.""" + + status_code: int + content: bytes + + @property + def ok(self) -> bool: + return self.status_code < 400 + + @property + def text(self) -> str: + return self.content.decode() + + def json(self) -> object: + return _JSON.validate_json(self.content) + + +class TestClassifyEmptyBody: + """A delete that answers 202 with no body is a success, not a parse failure: + the MCP server and toolset delete routes both answer that way, and reading it + as a failure would hide a delete that did not happen behind one that did.""" + + def test_empty_2xx_body_is_a_success(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody) + assert isinstance(result, Success) and result.status_code == 202 + + def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=200, content=b""), NoBody) + assert isinstance(result, ValidationError) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 2caac58333f..3b84a47e3cc 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -15,28 +15,35 @@ from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat from types import MappingProxyType -from typing import Final +from typing import Final, cast import pytest - from e2e_config import parse_replica_urls from e2e_http import Result, Success from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( - Poller, ConvergeOutcome, Converged, + EverywhereConverged, ModelsPoller, + NeverConvergedOn, NotConverged, NotServableOn, + Poller, + ProxyClient, + ReplicaRead, Servable, await_converged_everywhere, + await_everywhere, await_servable_everywhere, - first_lagging_replica, + build_proxy_client, converge_timeout_message, + first_lagging_replica, ) +from transport import Transport MODEL: Final = "gpt-under-test" +_NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 RPM_BEFORE_UPDATE: Final = 100 @@ -187,3 +194,83 @@ class TestParseReplicaUrls: def test_falls_back_to_the_data_plane_address_when_unset(self) -> None: assert parse_replica_urls("", "http://lb") == ("http://lb",) + + +def _answers(answers: Iterable[str]) -> ReplicaRead[str]: + it: Final = iter(answers) + return lambda _timeout: next(it) + + +def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> EverywhereConverged[str] | NeverConvergedOn[str]: + clock: Final = FakeClock() + return await_everywhere( + reads, + settled=lambda answer: answer == "renamed", + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))), + } + outcome: Final = _await_everywhere(reads) + assert isinstance(outcome, EverywhereConverged) + assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"} + + def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(repeat("stale")), + } + assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale") + + def test_polls_until_the_deadline_before_giving_up(self) -> None: + lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed")) + outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)}) + assert isinstance(outcome, EverywhereConverged), outcome + + +class TestReplicasFor: + def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} + + def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://lb", + replica_urls=("http://pod-1", "http://pod-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://pod-1", "http://pod-2"} + + def test_mcp_admin_routes_read_back_from_every_data_plane_replica(self) -> None: + """/v1/mcp/* is a lazily mounted feature, so a data-plane replica serves it + too and answers from its own in-memory registry. Routing it to the control + plane would leave every replica but that one unproven, and would move the + tools/list barrier in mcp_client off the plane that serves tools/list.""" + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://gateway-1", "http://gateway-2"} + assert set(client.replicas_for("/v1/mcp/toolset/abc")) == {"http://gateway-1", "http://gateway-2"} + + def test_a_route_no_replica_serves_is_refused_rather_than_read_back_vacuously(self) -> None: + """A read-back over zero replicas would satisfy every predicate and assert + nothing, so asking for one fails instead of passing silently.""" + client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) + with pytest.raises(AssertionError, match="no replica is configured"): + _ = client.replicas_for("/v1/models") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index c0d055edb7f..669e094fee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -1,10 +1,11 @@ """ -Tests for partial-update semantics of PUT /v1/mcp/server. +Tests for partial-update semantics of PUT /v1/mcp/server and PUT /v1/mcp/toolset. A partial update must only write the fields the caller explicitly provided. Omitting a field must NOT reset it to its Pydantic schema default (e.g. ``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which -would silently overwrite the existing DB row. +would silently overwrite the existing DB row, and a field the caller sent as null +must be cleared rather than left at its stored value. """ import json @@ -850,3 +851,69 @@ async def test_cf_pair_switch_does_not_clear_dcr_bridge(): data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") assert "dcr_bridge" not in data_dict + + +def _mock_toolset_prisma(): + """A prisma double whose update answers with a row the reader can expand, so the + call under test returns instead of failing inside the row mapper.""" + updated_row = MagicMock() + updated_row.model_dump.return_value = { + "toolset_id": "ts-1", + "toolset_name": "ops", + "description": None, + "tools": "[]", + } + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcptoolsettable = AsyncMock() + mock_prisma.db.litellm_mcptoolsettable.update = AsyncMock(return_value=updated_row) + return mock_prisma + + +async def _run_toolset_update(payload: dict) -> dict: + """The columns PUT /v1/mcp/toolset writes for this payload, minus the audit stamp + every write carries. The prisma double is injected, so nothing is patched.""" + from litellm.proxy._experimental.mcp_server.toolset_db import update_mcp_toolset + from litellm.types.mcp_server.mcp_toolset import UpdateMCPToolsetRequest + + mock_prisma = _mock_toolset_prisma() + await update_mcp_toolset(mock_prisma, UpdateMCPToolsetRequest.model_validate(payload), "test-user") + written = dict(mock_prisma.db.litellm_mcptoolsettable.update.call_args[1]["data"]) + assert written["updated_by"] == "test-user" + return {name: value for name, value in written.items() if name != "updated_by"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_clears_description_on_explicit_null(): + """The dump used to drop None, so a null description could never clear the stored + one: the toolset kept a description its owner had deleted.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "description": None}) == {"description": None} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_omits_the_fields_the_caller_left_out(): + tools = [{"server_id": "s1", "tool_name": "alpha"}] + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": tools}) == {"tools": json.dumps(tools)} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_null_tools_rather_than_revoking_them(): + """A client that sends tools=null means "leave the selection alone", so the grants + survive. Clearing them is an explicit [], which cannot be confused with an omitted + field; treating null as a clear would silently revoke every tool the toolset grants.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": None, "description": "kept"}) == { + "description": "kept" + } + + +@pytest.mark.asyncio +async def test_toolset_partial_update_empties_the_selection_on_an_explicit_empty_list(): + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": []}) == {"tools": "[]"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_a_null_name(): + """A toolset always has a name, so a null toolset_name is a no-op, not a clear + that would write a NOT NULL column to null.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "toolset_name": None, "description": "kept"}) == { + "description": "kept" + } From e8e3172d7d70558929f32f057ebe4c7471c8c352 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 8 Sep 2026 23:10:18 -0700 Subject: [PATCH 17/19] fix(model-management): honor an explicit null as a clear on model update (#40047) * fix(model-management): honor an explicit null as a clear on model update PATCH /model/{model_id}/update merged the patch with exclude_none and then popped explicit nulls only for the mirrored pricing fields, so a null sent for max_input_tokens, mode, supports_vision or any other key was dropped and a value pinned by an earlier save could never be removed. The route now follows JSON Merge Patch over both blobs: a key absent from the body is unchanged, a key sent as null is removed from the stored row, and a key sent with a value is set. Ownership and identity keys keep ignoring a null, as do the fields the stored models require, since clearing one writes a row no reload can rebuild. Mirrored pricing keys still clear from both blobs. Clearing a price also needed the router to stop merging a deployment's cost-map entry onto its previous registration, which left the old rate in place and kept billing at a price the deployment no longer carried. Adds a create, read, partial-update, clear, enforce, delete lifecycle e2e that reads back on every replica, and a harness helper for that read-back. * fix(router): keep a deployment id that names a real model from evicting its catalog entry Deployments are keyed into litellm.model_cost alongside the built-in catalog, so evicting a deployment's stale entry by id could take a real model's entry with it: registering a deployment whose model_info.id is "gpt-4o" stripped that model's pricing, context window and capability flags process-wide, for every other deployment of it, until the next price-map reload. Only evict an entry this registration owns. A colliding id keeps the previous merge, which pollutes the catalog entry rather than emptying it. Also pins the Admin UI round trip: the model edit form echoes the whole /model/info row back on save, and that read reports every key the deployment never stored as an explicit null, so the clear path has to leave those keys alone. * fix(router): decide cost-map eviction by what this registrar created The previous guard read a catalog entry off `litellm_provider`, so a deployment that declares its own provider in model_info was treated as one and kept billing at a price it no longer carried. It also only held for a single registration: a second one under a colliding id saw the id the first merge left behind and evicted the catalog entry anyway. Track the cost-map keys this registrar creates instead. A key it created is evicted before re-registration; one it did not is left to merge, which is what a deployment id colliding with a catalog model name needs. Also folds the required-fields comment into the docstring that already gives the reason. * fix(router): release a deployment's cost-map key when it is deleted The ownership ledger only grew. A deleted deployment kept its claim, so if a later catalog refresh started publishing a model under that same name, the next registration would treat the catalog entry as the deployment's own and evict it. Deleting a deployment now gives the key back, which also stops the ledger growing for the life of the process. * fix(router): hold a cost-map key while another live router still serves it The claim is process-wide but the release was per-deletion, so with two routers serving one deployment id, the first deletion put the survivor back on merging and the price it had just cleared would keep billing. Release the key only once no live router still serves that id. * fix(router): register a router in the live set when it gains a deployment _live_routers was only joined when a router was constructed with a model_list, but a router built empty is populated through add_deployment, and the empty branch exists for exactly that. Such a router was invisible to the live-router scan, so deleting the deployment from another router released the shared cost-map key while it was still serving that id. Joining the set where a deployment enters the list covers every path, and it also lets a price reload rebuild what a dynamically built router serves. * fix(e2e): read the stored model row from the control plane, not each gateway The lifecycle suite polled /model/info on every URL in PROXY_REPLICA_URLS. Those URLs are the stack's gateways, and gateway/routes/allowlist.py trims them to the LLM data-plane surface, so /model/info answers only on the backend and 404s on every replica. All five tests failed at their first read-back in CI while passing against a monolith, where one process serves both planes. The stored row has one answer behind it, so it is read through the shared transport, which routes control-plane paths to the backend. What every gateway must agree on is which models it serves, so the create and delete steps poll /v1/models per replica instead, a route the gateway does serve. read_back_everywhere now rejects a control-plane path outright rather than timing out on it. Two things surfaced behind that. /public/ was missing from the transport's control-plane prefixes, so model_cost_map() was routed to a gateway and 404'd, and the billing steps needed a data-plane wait: a PATCH lands on the backend and each gateway picks it up on its own config reload, measured here at 12-24s, so they now drive calls until the new rate reaches the spend row and let the deadline fail them. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01C1S92J8gSxxKVe1JBzxWBF * test(models): keep polling outcomes immutable and document shared ownership * test: validate opaque stream IDs and hide log-reader credentials * test: isolate auto-router scenarios and clean partial setup --------- Co-authored-by: Claude Opus 5 --- .../model_management_endpoints.py | 70 +++- litellm/router.py | 26 +- tests/e2e/coverage_registry/mgmt.yaml | 2 + .../management/test_model_lifecycle_e2e.py | 365 ++++++++++++++++++ tests/e2e/models.py | 74 +++- tests/e2e/proxy_client.py | 199 +++++++++- tests/e2e/test_proxy_client.py | 57 ++- tests/e2e/transport.py | 1 + .../test_model_management_endpoints.py | 178 ++++++++- .../test_router_model_cost_isolation.py | 192 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 11 files changed, 1130 insertions(+), 39 deletions(-) create mode 100644 tests/e2e/management/test_model_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0f19e9ce149..742b9d9817f 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -119,6 +119,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + LiteLLM_Params, ModelInfo, updateDeployment, ) @@ -728,6 +729,44 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) +_OWNERSHIP_FIELDS: Final = frozenset( + { + "db_model", + "team_id", + "team_public_model_name", + "access_groups", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + } +) + +_STORED_REQUIRED_FIELDS: Final = frozenset( + name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() +) + +_NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS | frozenset(PTU_MODEL_INFO_FIELDS) + + +def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: + """The keys a patch sends as an explicit null, which update_db_model removes from the + stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored + models require, since clearing one writes a row no reload can rebuild through + LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, + whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob + the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. + """ + if patch is None: + return frozenset() + return frozenset( + field + for field in patch.model_fields_set + if field not in _NULL_CLEAR_IGNORED_FIELDS and getattr(patch, field) is None + ) + + def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -748,25 +787,15 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.model_info: merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) - # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI - # passes through (which today re-sends the OLD pricing on every save) cannot - # silently undo a litellm_params clear via .update(). - # - # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character - # and cache read/write costs) so this path cannot be used to null out privileged - # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are - # mirrored between litellm_params and model_info by Deployment.__init__, so the - # clear propagates to both blobs. - if updated_patch.litellm_params: - for field in updated_patch.litellm_params.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_litellm_params.pop(field, None) - merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.litellm_params): + merged_litellm_params.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_model_info.pop(field, None) + for field in _explicitly_cleared_fields(updated_patch.model_info): + merged_model_info.pop(field, None) + if field in SPECIAL_MODEL_INFO_PARAMS: + merged_litellm_params.pop(field, None) if updated_patch.model_info: - for field in updated_patch.model_info.model_fields_set: - if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_model_info.pop(field, None) - merged_litellm_params.pop(field, None) for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) @@ -816,8 +845,9 @@ async def patch_model( """ PATCH Endpoint for partial model updates. - Only updates the fields specified in the request while preserving other existing values. - Follows proper PATCH semantics by only modifying provided fields. + JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + value is set (identity and ownership keys such as `id` and `team_id` ignore a null). Args: model_id: The ID of the model to update diff --git a/litellm/router.py b/litellm/router.py index 934a4ac86a9..1252d7e7487 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -628,6 +628,15 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat +# namespace with the built-in model catalog. Only a key it created may be evicted, or a +# deployment whose id names a real model would strip that model's pricing and +# capabilities for every other deployment of it. delete_deployment gives a key back once no +# live router still serves that id, so a later catalog refresh that starts serving the name +# is not treated as a deployment's own. +_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys + + class FallbackAwareStreamWrapper(CustomStreamWrapper): """Base for the Router's chat-completion stream wrappers, which are built around the attempt the Router picked first and have to repoint themselves when a fallback takes over.""" @@ -9761,6 +9770,7 @@ class Router: """ idx: Final = len(self.model_list) self.model_list.append(model) + _live_routers.add(self) # mutable-ok: track dynamic routers without extending their lifetimes self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() @@ -9929,7 +9939,12 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. + the entries a refresh rebuilds are the ones a fresh boot would produce. An + entry this function created is replaced rather than merged, so a price cleared + from the deployment does not linger from an earlier registration and keep + billing at the old rate. An entry it did not create is left to merge, because + a deployment id that collides with a catalog model name shares that model's + entry with every other deployment of it. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9946,6 +9961,10 @@ class Router: } if model_id is not None: + if model_id in _DEPLOYMENT_COST_MAP_KEYS: + litellm.model_cost.pop(model_id, None) # mutable-ok: remove cleared prices from the shared entry + elif model_id not in litellm.model_cost: + _DEPLOYMENT_COST_MAP_KEYS.add(model_id) # mutable-ok: retain shared ownership across reloads litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, @@ -10042,6 +10061,11 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) + if not any( + router is not self and id in router.model_id_to_deployment_index_map + for router in tuple(_live_routers) + ): + _DEPLOYMENT_COST_MAP_KEYS.discard(id) # mutable-ok: the last owning router released this key try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index c8d7037d2fd..83f1711a245 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,6 +76,8 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.model.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "model_management_endpoints.py:731", rationale: "A partial PATCH changes only the keys it names; every other stored key reads back byte-for-byte, and the new rate reaches billing"} +- {id: mgmt.model.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "model_management_endpoints.py:731", fail_before_fix: proven, rationale: "An explicit null on PATCH removes the key from the stored row and billing falls back to the cost map; before the fix only mirrored pricing keys could be cleared"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/management/test_model_lifecycle_e2e.py b/tests/e2e/management/test_model_lifecycle_e2e.py new file mode 100644 index 00000000000..99044efeb46 --- /dev/null +++ b/tests/e2e/management/test_model_lifecycle_e2e.py @@ -0,0 +1,365 @@ +"""Live e2e: the lifecycle of a DB-stored deployment through the model management +routes, read back on every gateway replica. + +Each test registers its own gpt-4o-mini mock deployment through /model/new (deleted +on teardown) with non-default pricing, context window, mode, and api_base pinned, then +walks the lifecycle up to the step it proves: the create reads back field for field, +a partial PATCH changes only the key it names, an explicit null on PATCH removes the +key from the stored row (JSON Merge Patch), a call after the price clear is billed at +the cost map's rate rather than the cleared override, and a delete removes the +deployment from /model/info and makes the model name unknown to /chat/completions. + +The stored row is read back from /model/info, a control-plane route with one answer +behind it. What every gateway must agree on is which models it serves, so the create +and delete steps poll /v1/models on every URL in PROXY_REPLICA_URLS through +ProxyClient.read_model_back_everywhere, failing by name on the gateway that never converged. +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + ChatBody, + ChatMessage, + Clear, + LiteLLMParamsBody, + LiteLLMParamsPatch, + ModelInfoBody, + ModelInfoEntry, + ModelInfoResponse, + ModelNewBody, + ModelPatchBody, + ModelsListResponse, + SpendLogRow, +) + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL: Final = "gpt-4o-mini" +PINNED_API_BASE: Final = "https://pinned.example.invalid/v1" +PINNED_MAX_INPUT_TOKENS: Final = 4096 +PINNED_INPUT_RATE: Final = 1e-05 +UPDATED_INPUT_RATE: Final = 2e-05 +PINNED_OUTPUT_RATE: Final = 3e-05 + +# A PATCH lands on the control plane, and each gateway picks it up on its own config +# reload, so the first call after the write can still be billed at the old rate. There +# is no price on the gateway's data-plane surface to poll, so the billing steps drive +# calls until the new rate shows up in the spend row and let the deadline be what fails. +BILLING_CONVERGENCE_TIMEOUT: Final = 90.0 +BILLING_CONVERGENCE_INTERVAL: Final = 5.0 + + +@dataclass(frozen=True, slots=True) +class Registered: + model_name: str + model_id: str + + +class _ErrorDetail(BaseModel): + message: str + + +class _ErrorEnvelope(BaseModel): + error: _ErrorDetail + + +def _register(client: ManagementClient, resources: ResourceManager) -> Registered: + """Register a mock gpt-4o-mini deployment with every field under test pinned to a + non-default value, deleted on teardown. max_input_tokens is pinned in + litellm_params only: a value in model_info is copied into the shared cost-map + entry for the backend model, which would leak into every other gpt-4o-mini + deployment on the proxy.""" + model_name: Final = f"e2e-lifecycle-{unique_marker()}" + model_id: Final = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=BACKEND_MODEL, + mock_response="ok", + api_base=PINNED_API_BASE, + input_cost_per_token=PINNED_INPUT_RATE, + output_cost_per_token=PINNED_OUTPUT_RATE, + max_input_tokens=PINNED_MAX_INPUT_TOKENS, + ), + model_info=ModelInfoBody(mode="chat"), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return Registered(model_name=model_name, model_id=model_id) + + +def _entry(body: ModelInfoResponse, model_name: str) -> ModelInfoEntry | None: + return next((entry for entry in body.data if entry.model_name == model_name), None) + + +def _stored_entry( + client: ManagementClient, + model_name: str, + *, + converged: Callable[[ModelInfoEntry], bool], +) -> ModelInfoEntry: + """The stored /model/info row for `model_name`, once it satisfies `converged`. + + /model/info is a control-plane route: the gateways named in PROXY_REPLICA_URLS + serve the LLM surface only, so the stored row has one answer, not one per + gateway. What every gateway must agree on is which models it serves, and + `_assert_served_everywhere` / `_assert_absent_everywhere` poll /v1/models for + that.""" + + def has_converged(body: ModelInfoResponse) -> bool: + entry: Final = _entry(body, model_name) + return entry is not None and converged(entry) + + body: Final = client.proxy.read_model_back("/model/info", ModelInfoResponse, predicate=has_converged) + entry: Final = _entry(body, model_name) + assert entry is not None, f"/model/info stopped listing {model_name!r} between the poll and the read" + return entry + + +def _serves(body: ModelsListResponse, model_name: str) -> bool: + return any(entry.id == model_name for entry in body.data) + + +def _assert_served_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: _serves(body, model_name) + ) + + +def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: + _ = client.proxy.read_model_back_everywhere( + "/v1/models", ModelsListResponse, predicate=lambda body: not _serves(body, model_name) + ) + + +def _assert_untouched_keys_as_created(entry: ModelInfoEntry, replica: str) -> None: + """The keys no later step names read back byte-for-byte as /model/new wrote them.""" + params: Final = entry.litellm_params + assert params.model == BACKEND_MODEL, f"{replica}: litellm_params.model {params.model!r} != {BACKEND_MODEL!r}" + assert params.api_base == PINNED_API_BASE, f"{replica}: api_base {params.api_base!r} != {PINNED_API_BASE!r}" + assert params.output_cost_per_token == PINNED_OUTPUT_RATE, ( + f"{replica}: output_cost_per_token {params.output_cost_per_token} != {PINNED_OUTPUT_RATE}" + ) + assert entry.model_info.mode == "chat", f"{replica}: model_info.mode {entry.model_info.mode!r} != 'chat'" + + +def _approx_equal(actual: float, expected: float) -> bool: + return math.isclose(actual, expected, rel_tol=1e-2, abs_tol=1e-9) + + +def _priced(rows: list[SpendLogRow]) -> bool: + return any(row.metadata and row.metadata.cost_breakdown and row.metadata.cost_breakdown.input_cost for row in rows) + + +def _billed_input_cost(client: ManagementClient, model_name: str, key: str) -> tuple[int, float]: + """Drive one chat completion through `model_name` and return the prompt tokens and + input cost its spend row recorded, so a test can assert the rate the gateway actually + billed rather than only the rate it stored.""" + chat: Final = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], + max_tokens=16, + ), + ) + ) + assert chat.id is not None, f"chat completion carried no id to find its spend row by: {chat}" + + rows: Final = client.proxy.poll_logs_for_request_id(chat.id, predicate=_priced) + row: Final = next((row for row in rows if row.request_id == chat.id), None) + assert row is not None and row.metadata and row.metadata.cost_breakdown, ( + f"no priced spend row for request {chat.id} before the deadline: {rows}" + ) + prompt_tokens: Final = row.prompt_tokens or 0 + input_cost: Final = row.metadata.cost_breakdown.input_cost + assert prompt_tokens > 0 and input_cost is not None, f"spend row logged no prompt tokens or input cost: {row}" + return prompt_tokens, input_cost + + +def _await_billed_input_cost( + client: ManagementClient, model_name: str, key: str, *, expected_rate: float +) -> tuple[int, float]: + """Drive calls through `model_name` until one is billed at `expected_rate`, and + return the prompt tokens and input cost of the last spend row either way. + + Only the deadline ends the wait unsatisfied: a rate that never reaches the gateway + comes back as the stale cost for the caller to assert on, so the rate the caller + expects is still what decides the test.""" + deadline: Final = time.monotonic() + BILLING_CONVERGENCE_TIMEOUT + while True: + prompt_tokens, input_cost = _billed_input_cost(client, model_name, key) + if _approx_equal(input_cost, prompt_tokens * expected_rate) or time.monotonic() >= deadline: + return prompt_tokens, input_cost + time.sleep(BILLING_CONVERGENCE_INTERVAL) + + +class TestModelLifecycle: + @pytest.mark.covers("mgmt.model.add.persists") + def test_create_reads_back_every_field_and_serves_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + entry = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( + f"{stored}: input_cost_per_token {entry.litellm_params.input_cost_per_token} != {PINNED_INPUT_RATE}" + ) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.id == registered.model_id, ( + f"{stored}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" + ) + + _assert_served_everywhere(client, registered.model_name) + + @pytest.mark.covers("mgmt.model.update.preserves_unrelated_fields") + def test_partial_update_changes_only_the_named_key( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(input_cost_per_token=UPDATED_INPUT_RATE)), + ) + assert stored.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"PATCH response stores input_cost_per_token {stored.litellm_params.input_cost_per_token}, " + f"sent {UPDATED_INPUT_RATE}" + ) + + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( + f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" + ) + assert entry.model_info.input_cost_per_token == UPDATED_INPUT_RATE, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} " + f"did not mirror the updated {UPDATED_INPUT_RATE}" + ) + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=UPDATED_INPUT_RATE + ) + assert _approx_equal(input_cost, prompt_tokens * UPDATED_INPUT_RATE), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * updated rate {UPDATED_INPUT_RATE} " + f"= {prompt_tokens * UPDATED_INPUT_RATE}; the partial update did not reach billing" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_explicit_null_removes_the_key_from_the_stored_row( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + registered = _register(client, resources) + + stored = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + stored_params = stored.litellm_params.model_fields_set + assert "max_input_tokens" not in stored_params, ( + f"stored litellm_params still carries max_input_tokens " + f"{stored.litellm_params.max_input_tokens} after an explicit null" + ) + assert "input_cost_per_token" not in stored_params, ( + f"stored litellm_params still carries input_cost_per_token " + f"{stored.litellm_params.input_cost_per_token} after an explicit null" + ) + assert "max_input_tokens" not in stored.model_info.model_fields_set, ( + f"stored model_info carries max_input_tokens {stored.model_info.max_input_tokens} after the clear" + ) + assert "input_cost_per_token" not in stored.model_info.model_fields_set, ( + f"stored model_info still mirrors input_cost_per_token {stored.model_info.input_cost_per_token}" + ) + + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + entry = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, + ) + stored = "/model/info" + + _assert_untouched_keys_as_created(entry, stored) + served = entry.litellm_params.model_fields_set + assert "max_input_tokens" not in served, ( + f"{stored}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" + ) + assert "input_cost_per_token" not in served, ( + f"{stored}: litellm_params still serves input_cost_per_token {entry.litellm_params.input_cost_per_token}" + ) + assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( + f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} is not the " + f"cost map's {cost_map_input_rate}; the cleared override {PINNED_INPUT_RATE} still resolves" + ) + + @pytest.mark.covers("mgmt.model.update.clear_persists") + def test_cleared_price_is_billed_at_the_cost_map_rate( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = client.proxy.patch_model( + registered.model_id, + ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), + ) + _ = _stored_entry( + client, + registered.model_name, + converged=lambda entry: "input_cost_per_token" not in entry.litellm_params.model_fields_set, + ) + cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token + assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" + + prompt_tokens, input_cost = _await_billed_input_cost( + client, registered.model_name, scoped_key, expected_rate=cost_map_input_rate + ) + + assert _approx_equal(input_cost, prompt_tokens * cost_map_input_rate), ( + f"input_cost {input_cost} != {prompt_tokens} tokens * cost map rate {cost_map_input_rate} " + f"= {prompt_tokens * cost_map_input_rate}" + ) + assert not _approx_equal(input_cost, prompt_tokens * PINNED_INPUT_RATE), ( + f"input_cost {input_cost} is still billed at the cleared override {PINNED_INPUT_RATE}" + ) + + @pytest.mark.covers("mgmt.model.delete.persists") + def test_delete_removes_the_deployment_everywhere( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + registered = _register(client, resources) + _ = _stored_entry(client, registered.model_name, converged=lambda _entry: True) + + client.delete_model_strict(registered.model_id) + + _assert_absent_everywhere(client, registered.model_name) + refused = client.chat_status(scoped_key, registered.model_name, "hi this is a test") + assert refused.status_code == 400, ( + f"chat against the deleted model must be rejected 400, got {refused.status_code}: {refused.body[:300]}" + ) + envelope = _ErrorEnvelope.model_validate_json(refused.body) + assert envelope.error.message, f"400 body must carry an error message: {refused.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..8db37bd25a5 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,6 +655,11 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- +class CostBreakdown(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + + class GuardrailEntityMatch(BaseModel): entity_type: str score: float @@ -672,6 +677,7 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): + cost_breakdown: CostBreakdown | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None @@ -815,15 +821,42 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token +class DeploymentParams(CustomPricing): + """The litellm_params half of a /model/info row: the stored deployment as written, + credentials scrubbed. Unlike model_info it is never back-filled from the cost map, + so a key the store dropped is absent here (check `model_fields_set`).""" + + model: str | None = None + api_base: str | None = None + max_input_tokens: int | None = None + + +class DeploymentModelInfo(CustomPricing): + id: str | None = None + max_input_tokens: int | None = None + + class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for - it - the override merged over the cost-map defaults.""" + it - the override merged over the cost-map defaults, so a key cleared from the + stored blob reads as the cost-map default here.""" model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: CustomPricing = CustomPricing() - model_info: CustomPricing = CustomPricing() + litellm_params: DeploymentParams = DeploymentParams() + model_info: DeploymentModelInfo = DeploymentModelInfo() + + +class StoredDeployment(BaseModel): + """PATCH /model/{model_id}/update answers with the row as stored: both blobs raw, + nothing back-filled, so a cleared key is absent from `model_fields_set` of the + blob it was cleared from.""" + + model_config = ConfigDict(protected_namespaces=()) + model_name: str + litellm_params: DeploymentParams + model_info: DeploymentModelInfo class ModelInfoResponse(BaseModel): @@ -924,9 +957,10 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None + max_input_tokens: int | None = None -ModelMode = Literal["batch", "realtime", "image_generation"] +ModelMode = Literal["chat", "batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): @@ -936,6 +970,7 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None + max_input_tokens: int | None = None access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None @@ -964,6 +999,37 @@ class ModelUpdateBody(BaseModel): model_info: ModelInfoBody +class Clear(BaseModel): + """Serializes to JSON null. The transport dumps every body with exclude_none, so a + field set to this is how a patch carries the explicit null that removes a stored key.""" + + @model_serializer + def _as_null(self) -> None: + return None + + +class LiteLLMParamsPatch(BaseModel): + api_base: str | Clear | None = None + max_input_tokens: int | Clear | None = None + input_cost_per_token: float | Clear | None = None + output_cost_per_token: float | Clear | None = None + + +class ModelInfoPatch(BaseModel): + mode: ModelMode | Clear | None = None + max_input_tokens: int | Clear | None = None + + +class ModelPatchBody(BaseModel): + """PATCH /model/{model_id}/update body, JSON Merge Patch over the stored deployment: + a field left None is dropped from the body and unchanged, a field set to `Clear()` + is sent as null and removed, a field with a value is set.""" + + model_config = ConfigDict(protected_namespaces=()) + litellm_params: LiteLLMParamsPatch | None = None + model_info: ModelInfoPatch | None = None + + class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1bac5116a9d..fa1b06fe7ed 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,10 +10,10 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass -from functools import reduce from datetime import datetime +from functools import reduce from types import MappingProxyType from typing import Final @@ -62,6 +62,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -72,6 +73,7 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + StoredDeployment, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -132,6 +134,103 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None +type BodyReader[R: BaseModel] = Callable[[float], Result[R]] + + +@dataclass(frozen=True, slots=True) +class BodyNotConverged[R: BaseModel]: + """The deadline passed without a read the predicate accepted; `last_result` is the + final read, so the caller can tell a body that never matched from a read that + failed.""" + + last_result: Result[R] | None + + +@dataclass(frozen=True, slots=True) +class BodyConverged[R: BaseModel]: + """Every replica answered a body the predicate accepted; `bodies` is the last read + per replica.""" + + bodies: Mapping[str, R] + + +@dataclass(frozen=True, slots=True) +class BodyNeverConvergedOn[R: BaseModel]: + """`BodyNotConverged` labeled with the replica whose reads never satisfied the predicate.""" + + replica: str + last_result: Result[R] | None + + +def await_body_converged[R: BaseModel]( + read: BodyReader[R], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Success[R] | BodyNotConverged[R]: + """Poll `read` until it answers a body `predicate` accepts, or `timeout` passes. + + Each read's request timeout is clamped to the remaining budget, and the sleep + between reads to the time left, so the last read before the deadline is never + skipped. Clock and sleep are injected.""" + deadline: Final = now() + timeout + + def reads() -> Iterator[Result[R]]: + while (remaining := deadline - now()) > 0: + yield read(min(request_timeout, remaining)) + sleep(min(interval, max(deadline - now(), 0.0))) + + def attempts() -> Iterator[Success[R] | BodyNotConverged[R]]: + for result in reads(): + if isinstance(result, Success) and predicate(result.data): + yield result + return + yield BodyNotConverged(last_result=result) + + initial: Final[Success[R] | BodyNotConverged[R]] = BodyNotConverged(last_result=None) + return reduce(lambda _previous, result: result, attempts(), initial) + + +def await_body_converged_everywhere[R: BaseModel]( + readers: Mapping[str, BodyReader[R]], + *, + predicate: Callable[[R], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + """`await_body_converged` against every replica in turn, each with the full budget, so a + write counts as landed only once every replica serves it.""" + def read_replica( + outcome: BodyConverged[R] | BodyNeverConvergedOn[R], + item: tuple[str, BodyReader[R]], + ) -> BodyConverged[R] | BodyNeverConvergedOn[R]: + if isinstance(outcome, BodyNeverConvergedOn): + return outcome + replica, read = item + match await_body_converged( + read, + predicate=predicate, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ): + case Success(data=data): + return BodyConverged(bodies=MappingProxyType({**outcome.bodies, replica: data})) + case BodyNotConverged(last_result=last_result): + return BodyNeverConvergedOn(replica=replica, last_result=last_result) + initial: Final[BodyConverged[R] | BodyNeverConvergedOn[R]] = BodyConverged(bodies=MappingProxyType({})) + return reduce(read_replica, readers.items(), initial) + + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -658,6 +757,102 @@ class ProxyClient: ) ) + def patch_model(self, model_id: str, body: ModelPatchBody) -> StoredDeployment: + """JSON Merge Patch the deployment `model_id` via PATCH /model/{model_id}/update: + a field the body omits is unchanged, one sent as null is removed from the stored + row, one sent with a value is set. See ModelPatchBody for how a null is sent. + Returns the row as stored after the write.""" + return unwrap( + self.transport.patch( + f"/model/{model_id}/update", + headers=self.transport.master, + json=body, + response_type=StoredDeployment, + ) + ) + + def read_model_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, predicate: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica until each answers a body `predicate` accepts, + polling to poll_timeout, and return the last body per replica. + + Fails naming the replica that never converged, so a write that reached one + gateway but not the others is caught instead of passing on whichever gateway + the balancer answered from. Falls back to the single proxy address when no + replica list is configured. + + `path` must be a data-plane route. The replicas are gateways, which serve only + the LLM surface, so a control-plane path answers on exactly one service and + 404s on every replica in a split deployment: asking each replica for one is + never the question the caller means. Read those through `self.transport` + instead, which routes them to the control plane.""" + if is_control_plane_path(path): + raise AssertionError( + f"read_model_back_everywhere({path!r}) asks every data-plane replica for a control-plane route. " + "The replicas are gateways and do not serve it; poll a data-plane path such as /v1/models " + "here, and read the control plane through the shared transport." + ) + readers: Final = { + url: self._body_reader(transport, path, response_type) + for url, transport in self._read_back_replicas().items() + } + outcome: Final = await_body_converged_everywhere( + readers, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies + case BodyNeverConvergedOn(replica=replica, last_result=last_result): + raise AssertionError( + f"GET {path} on {replica} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def read_model_back[R: BaseModel](self, path: str, response_type: type[R], *, predicate: Callable[[R], bool]) -> R: + """GET `path` through the shared transport until the body satisfies `predicate`, + polling to poll_timeout, and return that body. + + The counterpart to `read_model_back_everywhere` for a control-plane route such as + /model/info: the stored row lives in one database behind one control plane, so + there is a single answer to converge on rather than one per gateway.""" + outcome: Final = await_body_converged_everywhere( + {CONTROL_PLANE_BASE_URL: self._body_reader(self.transport, path, response_type)}, + predicate=predicate, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case BodyConverged(bodies=bodies): + return bodies[CONTROL_PLANE_BASE_URL] + case BodyNeverConvergedOn(last_result=last_result): + raise AssertionError( + f"GET {path} never answered the expected body within " + f"{self.poll_timeout}s; last read: {last_result}" + ) + + def _read_back_replicas(self) -> Mapping[str, Transport]: + return self.replicas or MappingProxyType({CONTROL_PLANE_BASE_URL: self.transport}) + + @staticmethod + def _body_reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> BodyReader[R]: + return lambda timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=timeout, + ) + def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 3b84a47e3cc..cbf7f5648d4 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -20,8 +20,12 @@ from typing import Final, cast import pytest from e2e_config import parse_replica_urls from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from models import KeyInfo, KeyInfoResponse, ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( + BodyReader, + BodyConverged, + BodyNeverConvergedOn, + await_body_converged_everywhere, ConvergeOutcome, Converged, EverywhereConverged, @@ -274,3 +278,54 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +def _info(*model_names: str) -> Success[ModelInfoResponse]: + entries: Final = [ModelInfoEntry(model_name=model_name) for model_name in model_names] + return Success(status_code=200, data=ModelInfoResponse(data=entries)) + + +def _reader(results: Iterable[Success[ModelInfoResponse]]) -> BodyReader[ModelInfoResponse]: + it: Final = iter(results) + return lambda _timeout: next(it) + + +def _lists_model(body: ModelInfoResponse) -> bool: + return any(entry.model_name == MODEL for entry in body.data) + + +def _read_back( + readers: Mapping[str, BodyReader[ModelInfoResponse]], +) -> tuple[BodyConverged[ModelInfoResponse] | BodyNeverConvergedOn[ModelInfoResponse], FakeClock]: + clock: Final = FakeClock() + outcome: Final = await_body_converged_everywhere( + readers, + predicate=_lists_model, + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + return outcome, clock + + +class TestAwaitBodyConvergedEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_body(self) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(chain(repeat(_info(), 2), repeat(_info(MODEL)))), + } + outcome, clock = _read_back(readers) + assert outcome == BodyConverged(bodies={"gateway-1": _info(MODEL).data, "gateway-2": _info(MODEL).data}) + assert clock.elapsed == 2 * INTERVAL + + @pytest.mark.parametrize("lagging", ["gateway-1", "gateway-2"]) + def test_fails_naming_the_replica_that_never_converges(self, lagging: str) -> None: + readers: Final = { + "gateway-1": _reader(repeat(_info(MODEL))), + "gateway-2": _reader(repeat(_info(MODEL))), + } | {lagging: _reader(repeat(_info()))} + outcome, clock = _read_back(readers) + assert outcome == BodyNeverConvergedOn(replica=lagging, last_result=_info()) + assert clock.elapsed >= TIMEOUT diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..804e073a4a0 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,6 +306,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/config", "/guardrails", "/openapi.json", + "/public/", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index c02f886fc31..1376727e296 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3115,9 +3115,6 @@ class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored between the two by Deployment.__init__). - - Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) - cannot be cleared via this path. """ def test_clear_input_cost_removes_from_both_blobs(self): @@ -3193,10 +3190,10 @@ class TestUpdateDBModelClearPricing: assert params["input_cost_per_token"] == 0.000001 assert params["output_cost_per_token"] == 0.000007 - def test_null_on_non_pricing_field_does_not_clear(self): - """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. - Privileged or unrelated model_info fields (e.g. team_id) must be unaffected - by the null-clearing path so a team admin can't ungate a team-scoped model. + def test_null_on_one_field_leaves_other_fields_alone(self): + """A null clears only the key it names: pricing the patch never mentions and + the ownership key team_id stay put, so a team admin can't ungate a + team-scoped model through the clear path. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, @@ -3217,8 +3214,6 @@ class TestUpdateDBModelClearPricing: model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), ) - # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id - # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, updated_patch=updateDeployment( @@ -3394,6 +3389,171 @@ class TestUpdateDBModelClearPricing: assert info["cache_creation_input_token_cost"] == 0.000003 +_PROTECTED_MODEL_INFO_VALUES = { + "team_id": "team-keep-me", + "team_public_model_name": "team-facing-name", + "access_groups": ["group-a"], + "created_at": "2026-01-01T00:00:00+00:00", + "created_by": "creator", + "updated_at": "2026-01-02T00:00:00+00:00", + "updated_by": "updater", + "blocked": True, +} + + +def _build_db_model_with_pinned_model_info(): + """Deployment whose stored blobs pin non-pricing keys an earlier save wrote, next to a + pricing override, so a clear can be checked key by key.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name="pinned-gpt-4o-mini", + litellm_params=LiteLLM_Params( + model="gpt-4o-mini", input_cost_per_token=0.000001, max_input_tokens=4096 + ), + model_info=ModelInfo( + id="dep-pinned-0", + max_input_tokens=4096, + mode="chat", + supports_vision=True, + **_PROTECTED_MODEL_INFO_VALUES, + ), + ) + + +class TestUpdateDBModelNullClearsAnyKey: + """JSON Merge Patch on PATCH /model/{id}/update: a key sent as null is removed from the + stored blob it was sent in, whatever the key, except the identity and ownership keys, + whose nulls are ignored.""" + + def test_model_info_nulls_remove_pinned_non_pricing_keys(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"max_input_tokens": None, "mode": None}} + ), + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info + assert "mode" not in info + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + + def test_litellm_params_null_removes_pinned_non_pricing_key(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"litellm_params": {"max_input_tokens": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in params + assert params["model"] == "gpt-4o-mini" + assert params["input_cost_per_token"] == 0.000001 + assert info["max_input_tokens"] == 4096 + + def test_omitted_key_is_untouched_by_a_null_elsewhere(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"mode": None, "supports_vision": False}} + ), + ) + + info = json.loads(result["model_info"]) + assert "mode" not in info + assert info["supports_vision"] is False + assert info["max_input_tokens"] == 4096 + + @pytest.mark.parametrize("field", sorted(_PROTECTED_MODEL_INFO_VALUES)) + def test_null_on_protected_key_is_ignored(self, field): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate({"model_info": {field: None}}), + ) + + info = json.loads(result["model_info"]) + assert info[field] == _PROTECTED_MODEL_INFO_VALUES[field] + assert info["max_input_tokens"] == 4096 + + def test_echoing_the_read_back_blob_preserves_every_stored_key(self): + """The Admin UI edit form submits the whole /model/info row back, and that read reports + every key the deployment never stored as an explicit null. Those nulls have to stay + no-ops: a write drops None before storing, so a null in the echoed blob always names a + key the stored row does not carry. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + db_model = _build_db_model_with_pinned_model_info() + echoed = { + "id": "dep-pinned-0", + "max_input_tokens": 4096, + "mode": "chat", + "supports_vision": True, + "input_cost_per_token": 0.000001, + "team_id": "team-keep-me", + "base_model": None, + "tier": None, + "max_output_tokens": None, + "supports_function_calling": None, + "cache_read_input_token_cost": None, + } + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment.model_validate({"model_info": echoed}), + ) + + info = json.loads(result["model_info"]) + assert info["max_input_tokens"] == 4096 + assert info["mode"] == "chat" + assert info["supports_vision"] is True + assert info["input_cost_per_token"] == 0.000001 + assert info["team_id"] == "team-keep-me" + for never_stored in ("base_model", "tier", "max_output_tokens", "supports_function_calling"): + assert never_stored not in info + + def test_null_on_pricing_key_still_clears_both_blobs(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + + result = update_db_model( + db_model=_build_db_model_with_pinned_model_info(), + updated_patch=updateDeployment.model_validate( + {"model_info": {"input_cost_per_token": None}} + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + assert params["max_input_tokens"] == 4096 + assert info["max_input_tokens"] == 4096 + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 30b265905f3..d22ec60e61a 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,6 +220,198 @@ def test_should_store_full_pricing_under_deployment_model_id(): assert entry["output_cost_per_token"] == 0.0 +def test_should_drop_a_price_the_deployment_no_longer_carries(): + """Re-registering a deployment must replace its model_id entry, not merge onto it. + + A merge left the old rate in the cost map after an operator cleared the override, so + the deployment kept billing at a price its config no longer had. + """ + backend_model = "vertex_ai/gemini-2.5-flash" + model_id = "deployment-cleared-price" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"input_cost_per_token": 0.005, "output_cost_per_token": 0.01}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"mode": "chat"}, + model=backend_model, + custom_llm_provider="vertex_ai", + ) + + entry = litellm.model_cost[model_id] + assert entry.get("input_cost_per_token") != 0.005, ( + "the cleared override survived re-registration, so the deployment still bills at it" + ) + assert entry.get("output_cost_per_token") != 0.01 + finally: + _restore_model_cost_entries(original) + + +def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it(): + """Deployments are keyed into the same cost map as the built-in catalog, so a deployment + whose id happens to name a real model must not evict that model's entry. + + Stripping it would take the pricing and capability flags every other deployment of that + model reads, process-wide, until the next price-map reload. Registering twice, because + the first registration is what would mark the entry as this deployment's own. + """ + colliding_id = "gpt-4o" + original = {colliding_id: litellm.model_cost.get(colliding_id)} + builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] + + try: + for _ in range(2): + Router._register_deployment_in_model_cost( + model_id=colliding_id, + model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + entry = litellm.model_cost[colliding_id] + assert entry["max_tokens"] == builtin_max_tokens, ( + "registering a deployment under a catalog model's name wiped that model's context window" + ) + assert entry["litellm_provider"] == "openai" + assert entry["supports_vision"] is True + finally: + _restore_model_cost_entries(original) + + +def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): + """A deployment may carry `litellm_provider` in its own model_info, which must not be + read as "this is a catalog entry" and stop the stale price from being dropped.""" + model_id = "deployment-provider-tagged" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( + "a deployment that declares its provider kept billing at the price it no longer carries" + ) + finally: + _restore_model_cost_entries(original) + + +def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): + """Deleting a deployment releases its claim on the shared cost-map key. + + Held forever, a later catalog refresh that starts publishing a model under that same + name would be treated as the deleted deployment's own entry and evicted. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-to-delete" + original = {model_id: litellm.model_cost.get(model_id)} + router = Router( + model_list=[ + { + "model_name": "to-delete", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + ] + ) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert router.delete_deployment(id=model_id) is not None + + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS, ( + "a deleted deployment kept its claim on the shared cost-map key" + ) + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): + """Two live routers can serve the same deployment id, and the claim is process-wide. + + Releasing it when only one of them drops the deployment would put the survivor back on + merging, so the price it just cleared would keep billing. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-served-twice" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "served-twice", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + first = Router(model_list=[entry]) + second = Router(model_list=[entry]) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert first.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while another router still served the deployment" + ) + + assert second.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + +def test_should_keep_the_cost_map_key_while_a_dynamically_built_router_serves_it(): + """A router built with no model_list still serves whatever add_deployment gives it, so it + counts when deciding whether the shared cost-map claim can be released.""" + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-added-dynamically" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "added-dynamically", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + configured = Router(model_list=[entry]) + dynamic = Router() + dynamic.add_deployment(deployment=Deployment(**entry)) + + try: + assert configured.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while a dynamically built router still served the deployment" + ) + + assert dynamic.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _restore_model_cost_entries(original) + + def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..93aaf3ca58c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9064,8 +9064,9 @@ export interface paths { * Patch Model * @description PATCH Endpoint for partial model updates. * - * Only updates the fields specified in the request while preserving other existing values. - * Follows proper PATCH semantics by only modifying provided fields. + * JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the + * body is unchanged, a key sent as null is removed from the stored row, and a key sent with a + * value is set (identity and ownership keys such as `id` and `team_id` ignore a null). * * Args: * model_id: The ID of the model to update From ef3a3c16ae02bc6d83e14b09c437ef36081c0498 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:32:31 -0700 Subject: [PATCH 18/19] feat(guardrails): map each guardrail scan id to its guardrail, stage and provider (#40327) * feat(guardrails): map each guardrail scan id to its guardrail, stage and provider Adds the x-litellm-guardrail-scan-metadata response header, a JSON list of {guardrail, stage, provider, scan_id} entries, next to the existing comma-separated x-litellm-guardrail-scan-id header. Prisma AIRS records the execution stage for every scan and OpenAI Moderation now records its moderation id too. The new metadata key is internal: client-supplied values are stripped and it is exposed through the UI CORS allow list. Resolves LIT-6018 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(guardrails): cap the scan metadata response header at a configurable length Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(guardrails): hardcode the scan metadata header cap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 + litellm/proxy/common_utils/callback_utils.py | 62 ++++++++- .../guardrail_hooks/openai/moderations.py | 10 +- .../panw_prisma_airs/panw_prisma_airs.py | 34 +++-- litellm/proxy/litellm_pre_call_utils.py | 2 + .../proxy/common_utils/test_callback_utils.py | 122 +++++++++++++++--- .../openai/test_moderations.py | 32 +++++ .../guardrail_hooks/test_panw_prisma_airs.py | 32 ++++- 8 files changed, 265 insertions(+), 31 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f2384c8c6c..108a914e9c1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -143,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) +MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048 DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000 @@ -197,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-adaptive-router-model", "x-litellm-applied-guardrails", "x-litellm-guardrail-scan-id", + "x-litellm-guardrail-scan-metadata", "x-litellm-cache-key", ] diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 770963a1f24..561a53409f4 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,10 +1,12 @@ import copy +import json import os from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass +from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias -from typing_extensions import assert_never +from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm import get_secret @@ -12,6 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -28,6 +31,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -52,6 +56,15 @@ reset_color_code: Final = "\033[0m" TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY: Final = "_pillar_response_headers_trusted" GUARDRAIL_SCAN_IDS_METADATA_KEY: Final = "guardrail_scan_ids" +GUARDRAIL_SCAN_METADATA_METADATA_KEY: Final = "guardrail_scan_metadata" + + +class GuardrailScanMetadata(TypedDict): + guardrail: ReadOnly[str | None] + stage: ReadOnly[str] + provider: ReadOnly[str] + scan_id: ReadOnly[str] + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -450,6 +463,16 @@ def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, return headers +def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: int) -> str | None: + """Compact JSON list of scan metadata entries, dropping trailing entries so the header fits in max_length.""" + encoded: Final = tuple(json.dumps(entry, separators=(",", ":")) for entry in entries) + lengths: Final = tuple(accumulate(len(item) + 1 for item in encoded)) + kept: Final = sum(1 for length in lengths if length + 1 <= max_length) + if kept == 0: + return None + return f"[{','.join(encoded[:kept])}]" + + def get_logging_caching_headers(request_data: dict) -> dict | None: _metadata: Final[dict] = {} metadata_bucket: Final = request_data.get("metadata") @@ -468,6 +491,15 @@ def get_logging_caching_headers(request_data: dict) -> dict | None: if scan_ids: headers["x-litellm-guardrail-scan-id"] = ",".join(scan_ids) + scan_metadata: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + scan_metadata_header: Final = ( + _serialize_scan_metadata_header(scan_metadata, max_length=MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH) + if isinstance(scan_metadata, (list, tuple)) + else None + ) + if scan_metadata_header: + headers["x-litellm-guardrail-scan-metadata"] = scan_metadata_header + if "applied_policies" in _metadata: headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"]) @@ -501,6 +533,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( "applied_policies", "applied_guardrails", GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, "policy_sources", "guardrails", "guardrail_config", @@ -565,21 +598,40 @@ def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_nam _metadata["applied_guardrails"] = [guardrail_name] -def add_guardrail_scan_id(request_data: dict, scan_id: str | None) -> None: +def add_guardrail_scan_id( + request_data: dict[str, object], + scan_id: str | None, + *, + guardrail_name: str | None, + provider: str, + stage: GuardrailEventHooks, +) -> None: """ - Record a provider scan id so it can be surfaced to the caller. + Record a provider scan id, keyed to the guardrail execution that produced it, so it can be surfaced to the caller. Guardrails only return scan details to the client when they block, so allowed requests carry no - audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header. + audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header, and the + (guardrail, stage, provider, scan_id) entries become the x-litellm-guardrail-scan-metadata header. """ if not scan_id: return _, _metadata = get_or_create_metadata_bucket(request_data) existing: Final = _metadata.get(GUARDRAIL_SCAN_IDS_METADATA_KEY) - scan_ids: Final = tuple(existing) if isinstance(existing, (list, tuple)) else () + scan_ids: Final[tuple[object, ...]] = tuple(existing) if isinstance(existing, (list, tuple)) else () if scan_id not in scan_ids: _metadata[GUARDRAIL_SCAN_IDS_METADATA_KEY] = (*scan_ids, scan_id) + entry: Final[GuardrailScanMetadata] = { + "guardrail": guardrail_name, + "stage": stage.value, + "provider": provider, + "scan_id": scan_id, + } + existing_entries: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY) + entries: Final[tuple[object, ...]] = tuple(existing_entries) if isinstance(existing_entries, (list, tuple)) else () + if entry not in entries: + _metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry) + def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 10683550f85..c22d35509c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -17,7 +17,8 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id +from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from litellm.types.utils import ( GenericGuardrailAPIInputs, GuardrailStatus, @@ -218,6 +219,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): metadata: Final = request_data.get("metadata") or {} request_data["metadata"] = metadata metadata["_openai_moderation_response"] = moderation_response.model_dump() + add_guardrail_scan_id( + request_data=request_data, + scan_id=moderation_response.id, + guardrail_name=self.guardrail_name, + provider=SupportedGuardrailIntegrations.OPENAI_MODERATION.value, + stage=GuardrailEventHooks.post_call if input_type == "response" else GuardrailEventHooks.pre_call, + ) # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index b73d3adb99e..3bc0dfabefc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -721,10 +721,18 @@ class PanwPrismaAirsHandler(CustomGuardrail): } } - def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None: + def _record_scan_id( + self, request_data: dict[str, object], scan_result: Mapping[str, object], stage: GuardrailEventHooks + ) -> None: """Surface the AIRS scan id on the response, so allowed calls are auditable too.""" scan_id: Final = scan_result.get("scan_id") - add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None) + add_guardrail_scan_id( + request_data=request_data, + scan_id=str(scan_id) if scan_id else None, + guardrail_name=self.guardrail_name, + provider=self._PROVIDER_NAME, + stage=stage, + ) def _handle_api_error_with_logging( self, @@ -948,7 +956,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): event_type=GuardrailEventHooks.post_call, ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool: """ @@ -1078,7 +1086,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.pre_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.pre_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1199,7 +1207,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(data, scan_result) + self._record_scan_id(data, scan_result, GuardrailEventHooks.post_call) action: Final = scan_result.get("action", "block") category: Final = scan_result.get("category", "unknown") @@ -1401,7 +1409,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - self._record_scan_id(request_data, scan_result) + self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call) # Add guardrail to applied guardrails header for observability add_guardrail_to_applied_guardrails_header( @@ -1475,7 +1483,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_args = self._masked_tool_call_arguments( @@ -1829,7 +1841,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): new_texts.append(text) continue - self._record_scan_id(request_data, scan_result) + self._record_scan_id( + request_data, + scan_result, + GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call, + ) action = scan_result.get("action", "block") masked_text = self._get_masked_text(scan_result, is_response=is_response) @@ -1901,7 +1917,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) # If we reach here, fallback_on_error="allow" else: - self._record_scan_id(request_data, mcp_scan_result) + self._record_scan_id(request_data, mcp_scan_result, GuardrailEventHooks.pre_call) action = mcp_scan_result.get("action", "block") masked_text = self._get_masked_text(mcp_scan_result, is_response=False) if action == "allow": diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 924f84be5f4..3250ae5cca9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -235,6 +235,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, "pillar_response_headers", @@ -291,6 +292,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "guardrail_scan_ids", + "guardrail_scan_metadata", "routing_decision", GATEWAY_INJECTED_CACHE_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 66f77db6da9..ecb2375d495 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,30 +1,33 @@ import copy +import json import sys from types import ModuleType, SimpleNamespace +from typing import Final +from unittest.mock import patch import pytest - +import litellm +from litellm.caching.caching import DualCache +from litellm.constants import MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( + _serialize_scan_metadata_header, add_guardrail_scan_id, add_policy_to_applied_policies_header, decrypt_callback_vars, encrypt_callback_vars, get_logging_caching_headers, - initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, + initialize_callbacks_on_proxy, normalize_callback_names, + process_callback, sanitize_openai_provider_metadata, strip_callback_config, ) -import litellm -from litellm.caching.caching import DualCache -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging - -from unittest.mock import patch -from litellm.proxy.common_utils.callback_utils import process_callback +from litellm.types.guardrails import GuardrailEventHooks def test_get_remaining_tokens_and_requests_from_request_data(): @@ -189,20 +192,109 @@ def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata(): assert headers["x-litellm-policy-sources"] == "global-baseline=team_default" +def _record( + request_data: dict[str, object], + scan_id: str | None, + guardrail_name: str = "airs", + provider: str = "panw_prisma_airs", + stage: GuardrailEventHooks = GuardrailEventHooks.pre_call, +) -> None: + add_guardrail_scan_id( + request_data=request_data, scan_id=scan_id, guardrail_name=guardrail_name, provider=provider, stage=stage + ) + + def test_add_guardrail_scan_id_dedupes_and_becomes_response_header(): request_data = {"litellm_metadata": {}} - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-1") - add_guardrail_scan_id(request_data=request_data, scan_id="scan-2") - add_guardrail_scan_id(request_data=request_data, scan_id=None) + _record(request_data, "scan-1") + _record(request_data, "scan-1") + _record(request_data, "scan-2") + _record(request_data, None) assert request_data["litellm_metadata"]["guardrail_scan_ids"] == ("scan-1", "scan-2") assert get_logging_caching_headers(request_data)["x-litellm-guardrail-scan-id"] == "scan-1,scan-2" -def test_get_logging_caching_headers_omits_scan_id_header_without_scans(): - assert "x-litellm-guardrail-scan-id" not in get_logging_caching_headers({"litellm_metadata": {}}) +def test_scan_metadata_header_maps_each_id_to_its_guardrail_stage_and_provider(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + + _record( + request_data, "scan-1", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "mod-1", guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.pre_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record( + request_data, "scan-2", guardrail_name="airs", provider="panw_prisma_airs", stage=GuardrailEventHooks.post_call + ) + _record(request_data, None, guardrail_name="mod", provider="openai_moderation", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1,mod-1,scan-2" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "airs", "stage": "pre_call", "provider": "panw_prisma_airs", "scan_id": "scan-1"}, + {"guardrail": "mod", "stage": "pre_call", "provider": "openai_moderation", "scan_id": "mod-1"}, + {"guardrail": "airs", "stage": "post_call", "provider": "panw_prisma_airs", "scan_id": "scan-2"}, + ] + + +def test_scan_metadata_keeps_same_id_reused_across_stages(): + request_data: Final[dict[str, object]] = {"metadata": {}} + + _record(request_data, "scan-1", stage=GuardrailEventHooks.pre_call) + _record(request_data, "scan-1", stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + assert [entry["stage"] for entry in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + "pre_call", + "post_call", + ] + + +def test_scan_metadata_header_drops_trailing_entries_to_stay_within_length_limit(): + request_data: Final[dict[str, object]] = {"litellm_metadata": {}} + scan_ids: Final = tuple(f"0f9c4b7e-3d2a-4c1b-9e8f-{index:012d}" for index in range(40)) + for scan_id in scan_ids: + _record(request_data, scan_id, stage=GuardrailEventHooks.post_call) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == ",".join(scan_ids) + header: Final = headers["x-litellm-guardrail-scan-metadata"] + assert len(header) <= MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH + kept: Final = json.loads(header) + assert 1 < len(kept) < len(scan_ids) + assert [entry["scan_id"] for entry in kept] == list(scan_ids[: len(kept)]) + + +def test_serialize_scan_metadata_header_keeps_exactly_the_entries_that_fit(): + entries: Final = ({"scan_id": "a"}, {"scan_id": "b"}, {"scan_id": "c"}) + two_entries: Final = '[{"scan_id":"a"},{"scan_id":"b"}]' + + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries)) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) - 1) == '[{"scan_id":"a"}]' + assert _serialize_scan_metadata_header(entries, max_length=len(two_entries) + 1) == two_entries + assert _serialize_scan_metadata_header(entries, max_length=1000) == json.dumps(entries, separators=(",", ":")) + assert _serialize_scan_metadata_header(entries, max_length=5) is None + assert _serialize_scan_metadata_header((), max_length=1000) is None + + +def test_scan_metadata_is_an_internal_metadata_key(): + assert sanitize_openai_provider_metadata({"guardrail_scan_metadata": "x", "keep": "y"}) == {"keep": "y"} + + +def test_get_logging_caching_headers_omits_scan_headers_without_scans(): + headers: Final = get_logging_caching_headers({"litellm_metadata": {}}) + assert headers is not None + assert "x-litellm-guardrail-scan-id" not in headers + assert "x-litellm-guardrail-scan-metadata" not in headers def test_initialize_callbacks_on_proxy_instantiates_compression_interception( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 2b43720a126..615d06b0f42 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -3,14 +3,19 @@ Test OpenAI Moderation Guardrail """ +import json import os +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) @@ -989,3 +994,30 @@ async def test_openai_moderation_initialize_guardrail_forwards_streaming_flags() assert guardrail.streaming_sampling_rate == 2 finally: litellm.logging_callback_manager._reset_all_callbacks() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("input_type", "stage"), [("request", "pre_call"), ("response", "post_call")]) +async def test_openai_moderation_records_moderation_id_as_scan_metadata(input_type: str, stage: str): + """Each moderation call's id is exposed with the guardrail name, stage and provider that produced it.""" + payload: Final = { + "id": f"modr-{stage}", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}, "category_applied_input_types": {}}], + } + http_client: Final = AsyncHTTPHandler() + http_client.client = httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail: Final = OpenAIModerationGuardrail(guardrail_name="openai-mod") + guardrail.async_handler = http_client + request_data: Final[dict[str, object]] = {"metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data=request_data, input_type=input_type) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + assert headers["x-litellm-guardrail-scan-id"] == f"modr-{stage}" + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + {"guardrail": "openai-mod", "stage": stage, "provider": "openai_moderation", "scan_id": f"modr-{stage}"} + ] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 8f29ba66814..3d7c6e06d94 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -12,6 +12,7 @@ This test file follows LiteLLM's testing patterns and covers: import copy import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5785,7 +5786,14 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123" - assert "x-litellm-guardrail-scan-metadata" not in headers + assert json.loads(headers["x-litellm-guardrail-scan-metadata"]) == [ + { + "guardrail": handler.guardrail_name, + "stage": "pre_call", + "provider": "panw_prisma_airs", + "scan_id": "scan-abc-123", + } + ] @pytest.mark.asyncio async def test_request_and_response_scan_ids_are_both_exposed(self, user_api_key_dict): @@ -5809,6 +5817,26 @@ class TestPanwAirsScanIdExposure: headers = get_logging_caching_headers(data) assert headers["x-litellm-guardrail-scan-id"] == "scan-abc-123,scan-response-456" + assert [(e["stage"], e["scan_id"]) for e in json.loads(headers["x-litellm-guardrail-scan-metadata"])] == [ + ("pre_call", "scan-abc-123"), + ("post_call", "scan-response-456"), + ] + + @pytest.mark.asyncio + async def test_apply_guardrail_response_scan_is_tagged_post_call(self): + from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers + + handler: Final = self._handler(self.ALLOW_SCAN_RESULT) + request_data: Final[dict[str, object]] = {"litellm_call_id": "test-call-id", "model": "gpt-4", "metadata": {}} + + await handler.apply_guardrail( + inputs={"texts": ["Hello world"]}, request_data=request_data, input_type="response" + ) + + headers: Final = get_logging_caching_headers(request_data) + assert headers is not None + entries: Final = json.loads(headers["x-litellm-guardrail-scan-metadata"]) + assert [(e["stage"], e["provider"]) for e in entries] == [("post_call", "panw_prisma_airs")] @pytest.mark.asyncio async def test_repeated_scan_id_is_not_duplicated(self, user_api_key_dict): @@ -5850,6 +5878,8 @@ class TestPanwAirsScanIdExposure: assert "guardrail_scan_ids" in _UNTRUSTED_METADATA_CONTROL_FIELDS assert "guardrail_scan_ids" in _UNTRUSTED_ROOT_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_METADATA_CONTROL_FIELDS + assert "guardrail_scan_metadata" in _UNTRUSTED_ROOT_CONTROL_FIELDS class TestPanwAirsBlockedErrorDetailPassthrough: """Regression tests for the full AIRS scan response on blocks. From 9d90a544916c6f38397862dd18b6a4c12de98827 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 9 Sep 2026 00:06:26 -0700 Subject: [PATCH 19/19] revert(model-management): roll back #40047 This reverts commit e8e3172d7d70558929f32f057ebe4c7471c8c352 Restore the previous model update and router cost registration behavior while pricing compatibility is investigated --- .../model_management_endpoints.py | 70 +--- litellm/router.py | 26 +- tests/e2e/coverage_registry/mgmt.yaml | 2 - .../management/test_model_lifecycle_e2e.py | 365 ------------------ tests/e2e/models.py | 74 +--- tests/e2e/proxy_client.py | 199 +--------- tests/e2e/test_proxy_client.py | 57 +-- tests/e2e/transport.py | 1 - .../test_model_management_endpoints.py | 178 +-------- .../test_router_model_cost_isolation.py | 192 --------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 11 files changed, 39 insertions(+), 1130 deletions(-) delete mode 100644 tests/e2e/management/test_model_lifecycle_e2e.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 742b9d9817f..0f19e9ce149 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -119,7 +119,6 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, - LiteLLM_Params, ModelInfo, updateDeployment, ) @@ -729,44 +728,6 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) -_OWNERSHIP_FIELDS: Final = frozenset( - { - "db_model", - "team_id", - "team_public_model_name", - "access_groups", - "created_at", - "created_by", - "updated_at", - "updated_by", - "blocked", - } -) - -_STORED_REQUIRED_FIELDS: Final = frozenset( - name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() -) - -_NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS | frozenset(PTU_MODEL_INFO_FIELDS) - - -def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: - """The keys a patch sends as an explicit null, which update_db_model removes from the - stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored - models require, since clearing one writes a row no reload can rebuild through - LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, - whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob - the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. - """ - if patch is None: - return frozenset() - return frozenset( - field - for field in patch.model_fields_set - if field not in _NULL_CLEAR_IGNORED_FIELDS and getattr(patch, field) is None - ) - - def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -787,15 +748,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.model_info: merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) - for field in _explicitly_cleared_fields(updated_patch.litellm_params): - merged_litellm_params.pop(field, None) - if field in SPECIAL_MODEL_INFO_PARAMS: - merged_model_info.pop(field, None) - for field in _explicitly_cleared_fields(updated_patch.model_info): - merged_model_info.pop(field, None) - if field in SPECIAL_MODEL_INFO_PARAMS: - merged_litellm_params.pop(field, None) + # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI + # passes through (which today re-sends the OLD pricing on every save) cannot + # silently undo a litellm_params clear via .update(). + # + # Restricted to SPECIAL_MODEL_INFO_PARAMS (input/output cost per token/character + # and cache read/write costs) so this path cannot be used to null out privileged + # model_info fields like team_id or access groups. SPECIAL_MODEL_INFO_PARAMS are + # mirrored between litellm_params and model_info by Deployment.__init__, so the + # clear propagates to both blobs. + if updated_patch.litellm_params: + for field in updated_patch.litellm_params.model_fields_set: + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: + for field in updated_patch.model_info.model_fields_set: + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) @@ -845,9 +816,8 @@ async def patch_model( """ PATCH Endpoint for partial model updates. - JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the - body is unchanged, a key sent as null is removed from the stored row, and a key sent with a - value is set (identity and ownership keys such as `id` and `team_id` ignore a null). + Only updates the fields specified in the request while preserving other existing values. + Follows proper PATCH semantics by only modifying provided fields. Args: model_id: The ID of the model to update diff --git a/litellm/router.py b/litellm/router.py index 1252d7e7487..934a4ac86a9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -628,15 +628,6 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 -# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat -# namespace with the built-in model catalog. Only a key it created may be evicted, or a -# deployment whose id names a real model would strip that model's pricing and -# capabilities for every other deployment of it. delete_deployment gives a key back once no -# live router still serves that id, so a later catalog refresh that starts serving the name -# is not treated as a deployment's own. -_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys - - class FallbackAwareStreamWrapper(CustomStreamWrapper): """Base for the Router's chat-completion stream wrappers, which are built around the attempt the Router picked first and have to repoint themselves when a fallback takes over.""" @@ -9770,7 +9761,6 @@ class Router: """ idx: Final = len(self.model_list) self.model_list.append(model) - _live_routers.add(self) # mutable-ok: track dynamic routers without extending their lifetimes self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() @@ -9939,12 +9929,7 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. An - entry this function created is replaced rather than merged, so a price cleared - from the deployment does not linger from an earlier registration and keep - billing at the old rate. An entry it did not create is left to merge, because - a deployment id that collides with a catalog model name shares that model's - entry with every other deployment of it. + the entries a refresh rebuilds are the ones a fresh boot would produce. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9961,10 +9946,6 @@ class Router: } if model_id is not None: - if model_id in _DEPLOYMENT_COST_MAP_KEYS: - litellm.model_cost.pop(model_id, None) # mutable-ok: remove cleared prices from the shared entry - elif model_id not in litellm.model_cost: - _DEPLOYMENT_COST_MAP_KEYS.add(model_id) # mutable-ok: retain shared ownership across reloads litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, @@ -10061,11 +10042,6 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) - if not any( - router is not self and id in router.model_id_to_deployment_index_map - for router in tuple(_live_routers) - ): - _DEPLOYMENT_COST_MAP_KEYS.discard(id) # mutable-ok: the last owning router released this key try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 83f1711a245..c8d7037d2fd 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,8 +76,6 @@ - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} - {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} -- {id: mgmt.model.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "model_management_endpoints.py:731", rationale: "A partial PATCH changes only the keys it names; every other stored key reads back byte-for-byte, and the new rate reaches billing"} -- {id: mgmt.model.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "model_management_endpoints.py:731", fail_before_fix: proven, rationale: "An explicit null on PATCH removes the key from the stored row and billing falls back to the cost map; before the fix only mirrored pricing keys could be cleared"} - {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} - {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} - {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} diff --git a/tests/e2e/management/test_model_lifecycle_e2e.py b/tests/e2e/management/test_model_lifecycle_e2e.py deleted file mode 100644 index 99044efeb46..00000000000 --- a/tests/e2e/management/test_model_lifecycle_e2e.py +++ /dev/null @@ -1,365 +0,0 @@ -"""Live e2e: the lifecycle of a DB-stored deployment through the model management -routes, read back on every gateway replica. - -Each test registers its own gpt-4o-mini mock deployment through /model/new (deleted -on teardown) with non-default pricing, context window, mode, and api_base pinned, then -walks the lifecycle up to the step it proves: the create reads back field for field, -a partial PATCH changes only the key it names, an explicit null on PATCH removes the -key from the stored row (JSON Merge Patch), a call after the price clear is billed at -the cost map's rate rather than the cleared override, and a delete removes the -deployment from /model/info and makes the model name unknown to /chat/completions. - -The stored row is read back from /model/info, a control-plane route with one answer -behind it. What every gateway must agree on is which models it serves, so the create -and delete steps poll /v1/models on every URL in PROXY_REPLICA_URLS through -ProxyClient.read_model_back_everywhere, failing by name on the gateway that never converged. -""" - -from __future__ import annotations - -import math -import time -from collections.abc import Callable -from dataclasses import dataclass -from typing import Final - -import pytest -from pydantic import BaseModel - -from e2e_config import unique_marker -from e2e_http import unwrap -from lifecycle import ResourceManager -from management_client import ManagementClient -from models import ( - ChatBody, - ChatMessage, - Clear, - LiteLLMParamsBody, - LiteLLMParamsPatch, - ModelInfoBody, - ModelInfoEntry, - ModelInfoResponse, - ModelNewBody, - ModelPatchBody, - ModelsListResponse, - SpendLogRow, -) - -pytestmark = pytest.mark.e2e - -BACKEND_MODEL: Final = "gpt-4o-mini" -PINNED_API_BASE: Final = "https://pinned.example.invalid/v1" -PINNED_MAX_INPUT_TOKENS: Final = 4096 -PINNED_INPUT_RATE: Final = 1e-05 -UPDATED_INPUT_RATE: Final = 2e-05 -PINNED_OUTPUT_RATE: Final = 3e-05 - -# A PATCH lands on the control plane, and each gateway picks it up on its own config -# reload, so the first call after the write can still be billed at the old rate. There -# is no price on the gateway's data-plane surface to poll, so the billing steps drive -# calls until the new rate shows up in the spend row and let the deadline be what fails. -BILLING_CONVERGENCE_TIMEOUT: Final = 90.0 -BILLING_CONVERGENCE_INTERVAL: Final = 5.0 - - -@dataclass(frozen=True, slots=True) -class Registered: - model_name: str - model_id: str - - -class _ErrorDetail(BaseModel): - message: str - - -class _ErrorEnvelope(BaseModel): - error: _ErrorDetail - - -def _register(client: ManagementClient, resources: ResourceManager) -> Registered: - """Register a mock gpt-4o-mini deployment with every field under test pinned to a - non-default value, deleted on teardown. max_input_tokens is pinned in - litellm_params only: a value in model_info is copied into the shared cost-map - entry for the backend model, which would leak into every other gpt-4o-mini - deployment on the proxy.""" - model_name: Final = f"e2e-lifecycle-{unique_marker()}" - model_id: Final = client.proxy.register_model( - ModelNewBody( - model_name=model_name, - litellm_params=LiteLLMParamsBody( - model=BACKEND_MODEL, - mock_response="ok", - api_base=PINNED_API_BASE, - input_cost_per_token=PINNED_INPUT_RATE, - output_cost_per_token=PINNED_OUTPUT_RATE, - max_input_tokens=PINNED_MAX_INPUT_TOKENS, - ), - model_info=ModelInfoBody(mode="chat"), - ) - ) - resources.defer(lambda: client.proxy.delete_model(model_id)) - return Registered(model_name=model_name, model_id=model_id) - - -def _entry(body: ModelInfoResponse, model_name: str) -> ModelInfoEntry | None: - return next((entry for entry in body.data if entry.model_name == model_name), None) - - -def _stored_entry( - client: ManagementClient, - model_name: str, - *, - converged: Callable[[ModelInfoEntry], bool], -) -> ModelInfoEntry: - """The stored /model/info row for `model_name`, once it satisfies `converged`. - - /model/info is a control-plane route: the gateways named in PROXY_REPLICA_URLS - serve the LLM surface only, so the stored row has one answer, not one per - gateway. What every gateway must agree on is which models it serves, and - `_assert_served_everywhere` / `_assert_absent_everywhere` poll /v1/models for - that.""" - - def has_converged(body: ModelInfoResponse) -> bool: - entry: Final = _entry(body, model_name) - return entry is not None and converged(entry) - - body: Final = client.proxy.read_model_back("/model/info", ModelInfoResponse, predicate=has_converged) - entry: Final = _entry(body, model_name) - assert entry is not None, f"/model/info stopped listing {model_name!r} between the poll and the read" - return entry - - -def _serves(body: ModelsListResponse, model_name: str) -> bool: - return any(entry.id == model_name for entry in body.data) - - -def _assert_served_everywhere(client: ManagementClient, model_name: str) -> None: - _ = client.proxy.read_model_back_everywhere( - "/v1/models", ModelsListResponse, predicate=lambda body: _serves(body, model_name) - ) - - -def _assert_absent_everywhere(client: ManagementClient, model_name: str) -> None: - _ = client.proxy.read_model_back_everywhere( - "/v1/models", ModelsListResponse, predicate=lambda body: not _serves(body, model_name) - ) - - -def _assert_untouched_keys_as_created(entry: ModelInfoEntry, replica: str) -> None: - """The keys no later step names read back byte-for-byte as /model/new wrote them.""" - params: Final = entry.litellm_params - assert params.model == BACKEND_MODEL, f"{replica}: litellm_params.model {params.model!r} != {BACKEND_MODEL!r}" - assert params.api_base == PINNED_API_BASE, f"{replica}: api_base {params.api_base!r} != {PINNED_API_BASE!r}" - assert params.output_cost_per_token == PINNED_OUTPUT_RATE, ( - f"{replica}: output_cost_per_token {params.output_cost_per_token} != {PINNED_OUTPUT_RATE}" - ) - assert entry.model_info.mode == "chat", f"{replica}: model_info.mode {entry.model_info.mode!r} != 'chat'" - - -def _approx_equal(actual: float, expected: float) -> bool: - return math.isclose(actual, expected, rel_tol=1e-2, abs_tol=1e-9) - - -def _priced(rows: list[SpendLogRow]) -> bool: - return any(row.metadata and row.metadata.cost_breakdown and row.metadata.cost_breakdown.input_cost for row in rows) - - -def _billed_input_cost(client: ManagementClient, model_name: str, key: str) -> tuple[int, float]: - """Drive one chat completion through `model_name` and return the prompt tokens and - input cost its spend row recorded, so a test can assert the rate the gateway actually - billed rather than only the rate it stored.""" - chat: Final = unwrap( - client.proxy.chat( - key, - ChatBody( - model=model_name, - messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], - max_tokens=16, - ), - ) - ) - assert chat.id is not None, f"chat completion carried no id to find its spend row by: {chat}" - - rows: Final = client.proxy.poll_logs_for_request_id(chat.id, predicate=_priced) - row: Final = next((row for row in rows if row.request_id == chat.id), None) - assert row is not None and row.metadata and row.metadata.cost_breakdown, ( - f"no priced spend row for request {chat.id} before the deadline: {rows}" - ) - prompt_tokens: Final = row.prompt_tokens or 0 - input_cost: Final = row.metadata.cost_breakdown.input_cost - assert prompt_tokens > 0 and input_cost is not None, f"spend row logged no prompt tokens or input cost: {row}" - return prompt_tokens, input_cost - - -def _await_billed_input_cost( - client: ManagementClient, model_name: str, key: str, *, expected_rate: float -) -> tuple[int, float]: - """Drive calls through `model_name` until one is billed at `expected_rate`, and - return the prompt tokens and input cost of the last spend row either way. - - Only the deadline ends the wait unsatisfied: a rate that never reaches the gateway - comes back as the stale cost for the caller to assert on, so the rate the caller - expects is still what decides the test.""" - deadline: Final = time.monotonic() + BILLING_CONVERGENCE_TIMEOUT - while True: - prompt_tokens, input_cost = _billed_input_cost(client, model_name, key) - if _approx_equal(input_cost, prompt_tokens * expected_rate) or time.monotonic() >= deadline: - return prompt_tokens, input_cost - time.sleep(BILLING_CONVERGENCE_INTERVAL) - - -class TestModelLifecycle: - @pytest.mark.covers("mgmt.model.add.persists") - def test_create_reads_back_every_field_and_serves_on_every_replica( - self, client: ManagementClient, resources: ResourceManager - ) -> None: - registered = _register(client, resources) - - entry = _stored_entry(client, registered.model_name, converged=lambda _entry: True) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - assert entry.litellm_params.input_cost_per_token == PINNED_INPUT_RATE, ( - f"{stored}: input_cost_per_token {entry.litellm_params.input_cost_per_token} != {PINNED_INPUT_RATE}" - ) - assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( - f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" - ) - assert entry.model_info.id == registered.model_id, ( - f"{stored}: model_info.id {entry.model_info.id!r} != {registered.model_id!r}" - ) - - _assert_served_everywhere(client, registered.model_name) - - @pytest.mark.covers("mgmt.model.update.preserves_unrelated_fields") - def test_partial_update_changes_only_the_named_key( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - - stored = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(input_cost_per_token=UPDATED_INPUT_RATE)), - ) - assert stored.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, ( - f"PATCH response stores input_cost_per_token {stored.litellm_params.input_cost_per_token}, " - f"sent {UPDATED_INPUT_RATE}" - ) - - entry = _stored_entry( - client, - registered.model_name, - converged=lambda entry: entry.litellm_params.input_cost_per_token == UPDATED_INPUT_RATE, - ) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - assert entry.litellm_params.max_input_tokens == PINNED_MAX_INPUT_TOKENS, ( - f"{stored}: max_input_tokens {entry.litellm_params.max_input_tokens} != {PINNED_MAX_INPUT_TOKENS}" - ) - assert entry.model_info.input_cost_per_token == UPDATED_INPUT_RATE, ( - f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} " - f"did not mirror the updated {UPDATED_INPUT_RATE}" - ) - - prompt_tokens, input_cost = _await_billed_input_cost( - client, registered.model_name, scoped_key, expected_rate=UPDATED_INPUT_RATE - ) - assert _approx_equal(input_cost, prompt_tokens * UPDATED_INPUT_RATE), ( - f"input_cost {input_cost} != {prompt_tokens} tokens * updated rate {UPDATED_INPUT_RATE} " - f"= {prompt_tokens * UPDATED_INPUT_RATE}; the partial update did not reach billing" - ) - - @pytest.mark.covers("mgmt.model.update.clear_persists") - def test_explicit_null_removes_the_key_from_the_stored_row( - self, client: ManagementClient, resources: ResourceManager - ) -> None: - registered = _register(client, resources) - - stored = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), - ) - stored_params = stored.litellm_params.model_fields_set - assert "max_input_tokens" not in stored_params, ( - f"stored litellm_params still carries max_input_tokens " - f"{stored.litellm_params.max_input_tokens} after an explicit null" - ) - assert "input_cost_per_token" not in stored_params, ( - f"stored litellm_params still carries input_cost_per_token " - f"{stored.litellm_params.input_cost_per_token} after an explicit null" - ) - assert "max_input_tokens" not in stored.model_info.model_fields_set, ( - f"stored model_info carries max_input_tokens {stored.model_info.max_input_tokens} after the clear" - ) - assert "input_cost_per_token" not in stored.model_info.model_fields_set, ( - f"stored model_info still mirrors input_cost_per_token {stored.model_info.input_cost_per_token}" - ) - - cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token - assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" - entry = _stored_entry( - client, - registered.model_name, - converged=lambda entry: "max_input_tokens" not in entry.litellm_params.model_fields_set, - ) - stored = "/model/info" - - _assert_untouched_keys_as_created(entry, stored) - served = entry.litellm_params.model_fields_set - assert "max_input_tokens" not in served, ( - f"{stored}: litellm_params still serves max_input_tokens {entry.litellm_params.max_input_tokens}" - ) - assert "input_cost_per_token" not in served, ( - f"{stored}: litellm_params still serves input_cost_per_token {entry.litellm_params.input_cost_per_token}" - ) - assert entry.model_info.input_cost_per_token == cost_map_input_rate, ( - f"{stored}: model_info.input_cost_per_token {entry.model_info.input_cost_per_token} is not the " - f"cost map's {cost_map_input_rate}; the cleared override {PINNED_INPUT_RATE} still resolves" - ) - - @pytest.mark.covers("mgmt.model.update.clear_persists") - def test_cleared_price_is_billed_at_the_cost_map_rate( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - _ = client.proxy.patch_model( - registered.model_id, - ModelPatchBody(litellm_params=LiteLLMParamsPatch(max_input_tokens=Clear(), input_cost_per_token=Clear())), - ) - _ = _stored_entry( - client, - registered.model_name, - converged=lambda entry: "input_cost_per_token" not in entry.litellm_params.model_fields_set, - ) - cost_map_input_rate = client.proxy.model_cost_map()[BACKEND_MODEL].input_cost_per_token - assert cost_map_input_rate is not None, f"cost map has no input rate for {BACKEND_MODEL}" - - prompt_tokens, input_cost = _await_billed_input_cost( - client, registered.model_name, scoped_key, expected_rate=cost_map_input_rate - ) - - assert _approx_equal(input_cost, prompt_tokens * cost_map_input_rate), ( - f"input_cost {input_cost} != {prompt_tokens} tokens * cost map rate {cost_map_input_rate} " - f"= {prompt_tokens * cost_map_input_rate}" - ) - assert not _approx_equal(input_cost, prompt_tokens * PINNED_INPUT_RATE), ( - f"input_cost {input_cost} is still billed at the cleared override {PINNED_INPUT_RATE}" - ) - - @pytest.mark.covers("mgmt.model.delete.persists") - def test_delete_removes_the_deployment_everywhere( - self, client: ManagementClient, resources: ResourceManager, scoped_key: str - ) -> None: - registered = _register(client, resources) - _ = _stored_entry(client, registered.model_name, converged=lambda _entry: True) - - client.delete_model_strict(registered.model_id) - - _assert_absent_everywhere(client, registered.model_name) - refused = client.chat_status(scoped_key, registered.model_name, "hi this is a test") - assert refused.status_code == 400, ( - f"chat against the deleted model must be rejected 400, got {refused.status_code}: {refused.body[:300]}" - ) - envelope = _ErrorEnvelope.model_validate_json(refused.body) - assert envelope.error.message, f"400 body must carry an error message: {refused.body[:300]}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 8db37bd25a5..62810e6cfd9 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,11 +655,6 @@ class OcrResponse(BaseModel): # ---------- spend logs ---------- -class CostBreakdown(BaseModel): - input_cost: float | None = None - output_cost: float | None = None - - class GuardrailEntityMatch(BaseModel): entity_type: str score: float @@ -677,7 +672,6 @@ class GuardrailRunRecord(BaseModel): class SpendLogMetadata(BaseModel): - cost_breakdown: CostBreakdown | None = None applied_guardrails: list[str] | None = None guardrail_information: list[GuardrailRunRecord] | None = None @@ -821,42 +815,15 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token -class DeploymentParams(CustomPricing): - """The litellm_params half of a /model/info row: the stored deployment as written, - credentials scrubbed. Unlike model_info it is never back-filled from the cost map, - so a key the store dropped is absent here (check `model_fields_set`).""" - - model: str | None = None - api_base: str | None = None - max_input_tokens: int | None = None - - -class DeploymentModelInfo(CustomPricing): - id: str | None = None - max_input_tokens: int | None = None - - class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for - it - the override merged over the cost-map defaults, so a key cleared from the - stored blob reads as the cost-map default here.""" + it - the override merged over the cost-map defaults.""" model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: DeploymentParams = DeploymentParams() - model_info: DeploymentModelInfo = DeploymentModelInfo() - - -class StoredDeployment(BaseModel): - """PATCH /model/{model_id}/update answers with the row as stored: both blobs raw, - nothing back-filled, so a cleared key is absent from `model_fields_set` of the - blob it was cleared from.""" - - model_config = ConfigDict(protected_namespaces=()) - model_name: str - litellm_params: DeploymentParams - model_info: DeploymentModelInfo + litellm_params: CustomPricing = CustomPricing() + model_info: CustomPricing = CustomPricing() class ModelInfoResponse(BaseModel): @@ -957,10 +924,9 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None - max_input_tokens: int | None = None -ModelMode = Literal["chat", "batch", "realtime", "image_generation"] +ModelMode = Literal["batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): @@ -970,7 +936,6 @@ class ModelInfoBody(BaseModel): # constraint when a prior run's teardown had not removed the row. id: str | None = None mode: ModelMode | None = None - max_input_tokens: int | None = None access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None @@ -999,37 +964,6 @@ class ModelUpdateBody(BaseModel): model_info: ModelInfoBody -class Clear(BaseModel): - """Serializes to JSON null. The transport dumps every body with exclude_none, so a - field set to this is how a patch carries the explicit null that removes a stored key.""" - - @model_serializer - def _as_null(self) -> None: - return None - - -class LiteLLMParamsPatch(BaseModel): - api_base: str | Clear | None = None - max_input_tokens: int | Clear | None = None - input_cost_per_token: float | Clear | None = None - output_cost_per_token: float | Clear | None = None - - -class ModelInfoPatch(BaseModel): - mode: ModelMode | Clear | None = None - max_input_tokens: int | Clear | None = None - - -class ModelPatchBody(BaseModel): - """PATCH /model/{model_id}/update body, JSON Merge Patch over the stored deployment: - a field left None is dropped from the body and unchanged, a field set to `Clear()` - is sent as null and removed, a field with a value is set.""" - - model_config = ConfigDict(protected_namespaces=()) - litellm_params: LiteLLMParamsPatch | None = None - model_info: ModelInfoPatch | None = None - - class ModelListEntry(BaseModel): id: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index fa1b06fe7ed..1bac5116a9d 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,10 +10,10 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass -from datetime import datetime from functools import reduce +from datetime import datetime from types import MappingProxyType from typing import Final @@ -62,7 +62,6 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, - ModelPatchBody, ModelsListParams, ModelsListResponse, ModelUpdateBody, @@ -73,7 +72,6 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, - StoredDeployment, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -134,103 +132,6 @@ class NotServableOn: last_result: Result[ModelsListResponse] | None -type BodyReader[R: BaseModel] = Callable[[float], Result[R]] - - -@dataclass(frozen=True, slots=True) -class BodyNotConverged[R: BaseModel]: - """The deadline passed without a read the predicate accepted; `last_result` is the - final read, so the caller can tell a body that never matched from a read that - failed.""" - - last_result: Result[R] | None - - -@dataclass(frozen=True, slots=True) -class BodyConverged[R: BaseModel]: - """Every replica answered a body the predicate accepted; `bodies` is the last read - per replica.""" - - bodies: Mapping[str, R] - - -@dataclass(frozen=True, slots=True) -class BodyNeverConvergedOn[R: BaseModel]: - """`BodyNotConverged` labeled with the replica whose reads never satisfied the predicate.""" - - replica: str - last_result: Result[R] | None - - -def await_body_converged[R: BaseModel]( - read: BodyReader[R], - *, - predicate: Callable[[R], bool], - timeout: float, - interval: float, - request_timeout: float, - now: Callable[[], float], - sleep: Callable[[float], None], -) -> Success[R] | BodyNotConverged[R]: - """Poll `read` until it answers a body `predicate` accepts, or `timeout` passes. - - Each read's request timeout is clamped to the remaining budget, and the sleep - between reads to the time left, so the last read before the deadline is never - skipped. Clock and sleep are injected.""" - deadline: Final = now() + timeout - - def reads() -> Iterator[Result[R]]: - while (remaining := deadline - now()) > 0: - yield read(min(request_timeout, remaining)) - sleep(min(interval, max(deadline - now(), 0.0))) - - def attempts() -> Iterator[Success[R] | BodyNotConverged[R]]: - for result in reads(): - if isinstance(result, Success) and predicate(result.data): - yield result - return - yield BodyNotConverged(last_result=result) - - initial: Final[Success[R] | BodyNotConverged[R]] = BodyNotConverged(last_result=None) - return reduce(lambda _previous, result: result, attempts(), initial) - - -def await_body_converged_everywhere[R: BaseModel]( - readers: Mapping[str, BodyReader[R]], - *, - predicate: Callable[[R], bool], - timeout: float, - interval: float, - request_timeout: float, - now: Callable[[], float], - sleep: Callable[[float], None], -) -> BodyConverged[R] | BodyNeverConvergedOn[R]: - """`await_body_converged` against every replica in turn, each with the full budget, so a - write counts as landed only once every replica serves it.""" - def read_replica( - outcome: BodyConverged[R] | BodyNeverConvergedOn[R], - item: tuple[str, BodyReader[R]], - ) -> BodyConverged[R] | BodyNeverConvergedOn[R]: - if isinstance(outcome, BodyNeverConvergedOn): - return outcome - replica, read = item - match await_body_converged( - read, - predicate=predicate, - timeout=timeout, - interval=interval, - request_timeout=request_timeout, - now=now, - sleep=sleep, - ): - case Success(data=data): - return BodyConverged(bodies=MappingProxyType({**outcome.bodies, replica: data})) - case BodyNotConverged(last_result=last_result): - return BodyNeverConvergedOn(replica=replica, last_result=last_result) - initial: Final[BodyConverged[R] | BodyNeverConvergedOn[R]] = BodyConverged(bodies=MappingProxyType({})) - return reduce(read_replica, readers.items(), initial) - - def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], *, @@ -757,102 +658,6 @@ class ProxyClient: ) ) - def patch_model(self, model_id: str, body: ModelPatchBody) -> StoredDeployment: - """JSON Merge Patch the deployment `model_id` via PATCH /model/{model_id}/update: - a field the body omits is unchanged, one sent as null is removed from the stored - row, one sent with a value is set. See ModelPatchBody for how a null is sent. - Returns the row as stored after the write.""" - return unwrap( - self.transport.patch( - f"/model/{model_id}/update", - headers=self.transport.master, - json=body, - response_type=StoredDeployment, - ) - ) - - def read_model_back_everywhere[R: BaseModel]( - self, path: str, response_type: type[R], *, predicate: Callable[[R], bool] - ) -> Mapping[str, R]: - """GET `path` on every replica until each answers a body `predicate` accepts, - polling to poll_timeout, and return the last body per replica. - - Fails naming the replica that never converged, so a write that reached one - gateway but not the others is caught instead of passing on whichever gateway - the balancer answered from. Falls back to the single proxy address when no - replica list is configured. - - `path` must be a data-plane route. The replicas are gateways, which serve only - the LLM surface, so a control-plane path answers on exactly one service and - 404s on every replica in a split deployment: asking each replica for one is - never the question the caller means. Read those through `self.transport` - instead, which routes them to the control plane.""" - if is_control_plane_path(path): - raise AssertionError( - f"read_model_back_everywhere({path!r}) asks every data-plane replica for a control-plane route. " - "The replicas are gateways and do not serve it; poll a data-plane path such as /v1/models " - "here, and read the control plane through the shared transport." - ) - readers: Final = { - url: self._body_reader(transport, path, response_type) - for url, transport in self._read_back_replicas().items() - } - outcome: Final = await_body_converged_everywhere( - readers, - predicate=predicate, - timeout=self.poll_timeout, - interval=self.poll_interval, - request_timeout=REQUEST_TIMEOUT, - now=time.monotonic, - sleep=time.sleep, - ) - match outcome: - case BodyConverged(bodies=bodies): - return bodies - case BodyNeverConvergedOn(replica=replica, last_result=last_result): - raise AssertionError( - f"GET {path} on {replica} never answered the expected body within " - f"{self.poll_timeout}s; last read: {last_result}" - ) - - def read_model_back[R: BaseModel](self, path: str, response_type: type[R], *, predicate: Callable[[R], bool]) -> R: - """GET `path` through the shared transport until the body satisfies `predicate`, - polling to poll_timeout, and return that body. - - The counterpart to `read_model_back_everywhere` for a control-plane route such as - /model/info: the stored row lives in one database behind one control plane, so - there is a single answer to converge on rather than one per gateway.""" - outcome: Final = await_body_converged_everywhere( - {CONTROL_PLANE_BASE_URL: self._body_reader(self.transport, path, response_type)}, - predicate=predicate, - timeout=self.poll_timeout, - interval=self.poll_interval, - request_timeout=REQUEST_TIMEOUT, - now=time.monotonic, - sleep=time.sleep, - ) - match outcome: - case BodyConverged(bodies=bodies): - return bodies[CONTROL_PLANE_BASE_URL] - case BodyNeverConvergedOn(last_result=last_result): - raise AssertionError( - f"GET {path} never answered the expected body within " - f"{self.poll_timeout}s; last read: {last_result}" - ) - - def _read_back_replicas(self) -> Mapping[str, Transport]: - return self.replicas or MappingProxyType({CONTROL_PLANE_BASE_URL: self.transport}) - - @staticmethod - def _body_reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> BodyReader[R]: - return lambda timeout: transport.get( - path, - headers=transport.master, - params=NoBody(), - response_type=response_type, - timeout=timeout, - ) - def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index cbf7f5648d4..3b84a47e3cc 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -20,12 +20,8 @@ from typing import Final, cast import pytest from e2e_config import parse_replica_urls from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelInfoEntry, ModelInfoResponse, ModelListEntry, ModelsListResponse +from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( - BodyReader, - BodyConverged, - BodyNeverConvergedOn, - await_body_converged_everywhere, ConvergeOutcome, Converged, EverywhereConverged, @@ -278,54 +274,3 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") - - -def _info(*model_names: str) -> Success[ModelInfoResponse]: - entries: Final = [ModelInfoEntry(model_name=model_name) for model_name in model_names] - return Success(status_code=200, data=ModelInfoResponse(data=entries)) - - -def _reader(results: Iterable[Success[ModelInfoResponse]]) -> BodyReader[ModelInfoResponse]: - it: Final = iter(results) - return lambda _timeout: next(it) - - -def _lists_model(body: ModelInfoResponse) -> bool: - return any(entry.model_name == MODEL for entry in body.data) - - -def _read_back( - readers: Mapping[str, BodyReader[ModelInfoResponse]], -) -> tuple[BodyConverged[ModelInfoResponse] | BodyNeverConvergedOn[ModelInfoResponse], FakeClock]: - clock: Final = FakeClock() - outcome: Final = await_body_converged_everywhere( - readers, - predicate=_lists_model, - timeout=TIMEOUT, - interval=INTERVAL, - request_timeout=5.0, - now=clock.now, - sleep=clock.sleep, - ) - return outcome, clock - - -class TestAwaitBodyConvergedEverywhere: - def test_waits_for_the_lagging_replica_and_returns_every_body(self) -> None: - readers: Final = { - "gateway-1": _reader(repeat(_info(MODEL))), - "gateway-2": _reader(chain(repeat(_info(), 2), repeat(_info(MODEL)))), - } - outcome, clock = _read_back(readers) - assert outcome == BodyConverged(bodies={"gateway-1": _info(MODEL).data, "gateway-2": _info(MODEL).data}) - assert clock.elapsed == 2 * INTERVAL - - @pytest.mark.parametrize("lagging", ["gateway-1", "gateway-2"]) - def test_fails_naming_the_replica_that_never_converges(self, lagging: str) -> None: - readers: Final = { - "gateway-1": _reader(repeat(_info(MODEL))), - "gateway-2": _reader(repeat(_info(MODEL))), - } | {lagging: _reader(repeat(_info()))} - outcome, clock = _read_back(readers) - assert outcome == BodyNeverConvergedOn(replica=lagging, last_result=_info()) - assert clock.elapsed >= TIMEOUT diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 804e073a4a0..44fdbaa3e41 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -306,7 +306,6 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/config", "/guardrails", "/openapi.json", - "/public/", ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 1376727e296..c02f886fc31 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3115,6 +3115,9 @@ class TestUpdateDBModelClearPricing: """Sending an explicit `null` for a pricing field must remove it from both `litellm_params` and `model_info` (SPECIAL_MODEL_INFO_PARAMS are mirrored between the two by Deployment.__init__). + + Restricted to SPECIAL_MODEL_INFO_PARAMS so non-pricing fields (e.g. team_id) + cannot be cleared via this path. """ def test_clear_input_cost_removes_from_both_blobs(self): @@ -3190,10 +3193,10 @@ class TestUpdateDBModelClearPricing: assert params["input_cost_per_token"] == 0.000001 assert params["output_cost_per_token"] == 0.000007 - def test_null_on_one_field_leaves_other_fields_alone(self): - """A null clears only the key it names: pricing the patch never mentions and - the ownership key team_id stay put, so a team admin can't ungate a - team-scoped model through the clear path. + def test_null_on_non_pricing_field_does_not_clear(self): + """Security guard: only SPECIAL_MODEL_INFO_PARAMS can be cleared via null. + Privileged or unrelated model_info fields (e.g. team_id) must be unaffected + by the null-clearing path so a team admin can't ungate a team-scoped model. """ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, @@ -3214,6 +3217,8 @@ class TestUpdateDBModelClearPricing: model_info=ModelInfo(id="dep-pricing-1", team_id="team-keep-me"), ) + # Patch sends a null for api_base (non-SPECIAL field). Must NOT clear team_id + # or any other non-pricing field from the merged dict. result = update_db_model( db_model=db_model, updated_patch=updateDeployment( @@ -3389,171 +3394,6 @@ class TestUpdateDBModelClearPricing: assert info["cache_creation_input_token_cost"] == 0.000003 -_PROTECTED_MODEL_INFO_VALUES = { - "team_id": "team-keep-me", - "team_public_model_name": "team-facing-name", - "access_groups": ["group-a"], - "created_at": "2026-01-01T00:00:00+00:00", - "created_by": "creator", - "updated_at": "2026-01-02T00:00:00+00:00", - "updated_by": "updater", - "blocked": True, -} - - -def _build_db_model_with_pinned_model_info(): - """Deployment whose stored blobs pin non-pricing keys an earlier save wrote, next to a - pricing override, so a clear can be checked key by key.""" - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - - return Deployment( - model_name="pinned-gpt-4o-mini", - litellm_params=LiteLLM_Params( - model="gpt-4o-mini", input_cost_per_token=0.000001, max_input_tokens=4096 - ), - model_info=ModelInfo( - id="dep-pinned-0", - max_input_tokens=4096, - mode="chat", - supports_vision=True, - **_PROTECTED_MODEL_INFO_VALUES, - ), - ) - - -class TestUpdateDBModelNullClearsAnyKey: - """JSON Merge Patch on PATCH /model/{id}/update: a key sent as null is removed from the - stored blob it was sent in, whatever the key, except the identity and ownership keys, - whose nulls are ignored.""" - - def test_model_info_nulls_remove_pinned_non_pricing_keys(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"max_input_tokens": None, "mode": None}} - ), - ) - - info = json.loads(result["model_info"]) - assert "max_input_tokens" not in info - assert "mode" not in info - assert info["supports_vision"] is True - assert info["input_cost_per_token"] == 0.000001 - - def test_litellm_params_null_removes_pinned_non_pricing_key(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"litellm_params": {"max_input_tokens": None}} - ), - ) - - params = json.loads(result["litellm_params"]) - info = json.loads(result["model_info"]) - assert "max_input_tokens" not in params - assert params["model"] == "gpt-4o-mini" - assert params["input_cost_per_token"] == 0.000001 - assert info["max_input_tokens"] == 4096 - - def test_omitted_key_is_untouched_by_a_null_elsewhere(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"mode": None, "supports_vision": False}} - ), - ) - - info = json.loads(result["model_info"]) - assert "mode" not in info - assert info["supports_vision"] is False - assert info["max_input_tokens"] == 4096 - - @pytest.mark.parametrize("field", sorted(_PROTECTED_MODEL_INFO_VALUES)) - def test_null_on_protected_key_is_ignored(self, field): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate({"model_info": {field: None}}), - ) - - info = json.loads(result["model_info"]) - assert info[field] == _PROTECTED_MODEL_INFO_VALUES[field] - assert info["max_input_tokens"] == 4096 - - def test_echoing_the_read_back_blob_preserves_every_stored_key(self): - """The Admin UI edit form submits the whole /model/info row back, and that read reports - every key the deployment never stored as an explicit null. Those nulls have to stay - no-ops: a write drops None before storing, so a null in the echoed blob always names a - key the stored row does not carry. - """ - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - db_model = _build_db_model_with_pinned_model_info() - echoed = { - "id": "dep-pinned-0", - "max_input_tokens": 4096, - "mode": "chat", - "supports_vision": True, - "input_cost_per_token": 0.000001, - "team_id": "team-keep-me", - "base_model": None, - "tier": None, - "max_output_tokens": None, - "supports_function_calling": None, - "cache_read_input_token_cost": None, - } - - result = update_db_model( - db_model=db_model, - updated_patch=updateDeployment.model_validate({"model_info": echoed}), - ) - - info = json.loads(result["model_info"]) - assert info["max_input_tokens"] == 4096 - assert info["mode"] == "chat" - assert info["supports_vision"] is True - assert info["input_cost_per_token"] == 0.000001 - assert info["team_id"] == "team-keep-me" - for never_stored in ("base_model", "tier", "max_output_tokens", "supports_function_calling"): - assert never_stored not in info - - def test_null_on_pricing_key_still_clears_both_blobs(self): - from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_db_model, - ) - - result = update_db_model( - db_model=_build_db_model_with_pinned_model_info(), - updated_patch=updateDeployment.model_validate( - {"model_info": {"input_cost_per_token": None}} - ), - ) - - params = json.loads(result["litellm_params"]) - info = json.loads(result["model_info"]) - assert "input_cost_per_token" not in params - assert "input_cost_per_token" not in info - assert params["max_input_tokens"] == 4096 - assert info["max_input_tokens"] == 4096 - - class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index d22ec60e61a..30b265905f3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -220,198 +220,6 @@ def test_should_store_full_pricing_under_deployment_model_id(): assert entry["output_cost_per_token"] == 0.0 -def test_should_drop_a_price_the_deployment_no_longer_carries(): - """Re-registering a deployment must replace its model_id entry, not merge onto it. - - A merge left the old rate in the cost map after an operator cleared the override, so - the deployment kept billing at a price its config no longer had. - """ - backend_model = "vertex_ai/gemini-2.5-flash" - model_id = "deployment-cleared-price" - original = {model_id: litellm.model_cost.get(model_id)} - - try: - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"input_cost_per_token": 0.005, "output_cost_per_token": 0.01}, - model=backend_model, - custom_llm_provider="vertex_ai", - ) - assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 - - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"mode": "chat"}, - model=backend_model, - custom_llm_provider="vertex_ai", - ) - - entry = litellm.model_cost[model_id] - assert entry.get("input_cost_per_token") != 0.005, ( - "the cleared override survived re-registration, so the deployment still bills at it" - ) - assert entry.get("output_cost_per_token") != 0.01 - finally: - _restore_model_cost_entries(original) - - -def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it(): - """Deployments are keyed into the same cost map as the built-in catalog, so a deployment - whose id happens to name a real model must not evict that model's entry. - - Stripping it would take the pricing and capability flags every other deployment of that - model reads, process-wide, until the next price-map reload. Registering twice, because - the first registration is what would mark the entry as this deployment's own. - """ - colliding_id = "gpt-4o" - original = {colliding_id: litellm.model_cost.get(colliding_id)} - builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] - - try: - for _ in range(2): - Router._register_deployment_in_model_cost( - model_id=colliding_id, - model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - entry = litellm.model_cost[colliding_id] - assert entry["max_tokens"] == builtin_max_tokens, ( - "registering a deployment under a catalog model's name wiped that model's context window" - ) - assert entry["litellm_provider"] == "openai" - assert entry["supports_vision"] is True - finally: - _restore_model_cost_entries(original) - - -def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): - """A deployment may carry `litellm_provider` in its own model_info, which must not be - read as "this is a catalog entry" and stop the stale price from being dropped.""" - model_id = "deployment-provider-tagged" - original = {model_id: litellm.model_cost.get(model_id)} - - try: - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 - - Router._register_deployment_in_model_cost( - model_id=model_id, - model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( - "a deployment that declares its provider kept billing at the price it no longer carries" - ) - finally: - _restore_model_cost_entries(original) - - -def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): - """Deleting a deployment releases its claim on the shared cost-map key. - - Held forever, a later catalog refresh that starts publishing a model under that same - name would be treated as the deleted deployment's own entry and evicted. - """ - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-to-delete" - original = {model_id: litellm.model_cost.get(model_id)} - router = Router( - model_list=[ - { - "model_name": "to-delete", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - ] - ) - - try: - assert model_id in _DEPLOYMENT_COST_MAP_KEYS - - assert router.delete_deployment(id=model_id) is not None - - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS, ( - "a deleted deployment kept its claim on the shared cost-map key" - ) - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - -def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): - """Two live routers can serve the same deployment id, and the claim is process-wide. - - Releasing it when only one of them drops the deployment would put the survivor back on - merging, so the price it just cleared would keep billing. - """ - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-served-twice" - original = {model_id: litellm.model_cost.get(model_id)} - entry = { - "model_name": "served-twice", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - first = Router(model_list=[entry]) - second = Router(model_list=[entry]) - - try: - assert model_id in _DEPLOYMENT_COST_MAP_KEYS - - assert first.delete_deployment(id=model_id) is not None - - assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( - "the claim was released while another router still served the deployment" - ) - - assert second.delete_deployment(id=model_id) is not None - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - -def test_should_keep_the_cost_map_key_while_a_dynamically_built_router_serves_it(): - """A router built with no model_list still serves whatever add_deployment gives it, so it - counts when deciding whether the shared cost-map claim can be released.""" - from litellm.router import _DEPLOYMENT_COST_MAP_KEYS - - model_id = "deployment-added-dynamically" - original = {model_id: litellm.model_cost.get(model_id)} - entry = { - "model_name": "added-dynamically", - "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, - "model_info": {"id": model_id, "input_cost_per_token": 0.005}, - } - configured = Router(model_list=[entry]) - dynamic = Router() - dynamic.add_deployment(deployment=Deployment(**entry)) - - try: - assert configured.delete_deployment(id=model_id) is not None - - assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( - "the claim was released while a dynamically built router still served the deployment" - ) - - assert dynamic.delete_deployment(id=model_id) is not None - assert model_id not in _DEPLOYMENT_COST_MAP_KEYS - finally: - _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) - _restore_model_cost_entries(original) - - def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 93aaf3ca58c..83b0d58f2b2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9064,9 +9064,8 @@ export interface paths { * Patch Model * @description PATCH Endpoint for partial model updates. * - * JSON Merge Patch semantics over `litellm_params` and `model_info`: a key absent from the - * body is unchanged, a key sent as null is removed from the stored row, and a key sent with a - * value is set (identity and ownership keys such as `id` and `team_id` ignore a null). + * Only updates the fields specified in the request while preserving other existing values. + * Follows proper PATCH semantics by only modifying provided fields. * * Args: * model_id: The ID of the model to update