Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_mcp_connection_errors_31318

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -1801,7 +1801,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)
 
@@ -1812,7 +1821,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

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",
 ]
 

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

litellm/llms/bedrock/files/transformation.py

@@ -7,13 +7,13 @@ 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
 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
@@ -60,11 +60,12 @@ 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"
+
+
+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`.
@@ -291,7 +292,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 +1188,27 @@ 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]]:
+        return self._transform_s3_file_request(
+            file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params
+        )
 
     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,
+            )
+        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,
@@ -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},
         )

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},

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

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)

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

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,

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(

tests/e2e/batches/COVERAGE.md

@@ -120,6 +120,36 @@ 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. Managed batch cancellation polls for up to eleven minutes
+before input deletion: the ten-minute provider window plus a propagation margin.
+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
+
+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.
+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

tests/e2e/batches/batch_cleanup.py

@@ -0,0 +1,140 @@
+from builtins import ExceptionGroup
+from collections.abc import Callable
+from itertools import count
+from time import monotonic, sleep
+from typing import Final, Protocol
+
+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
+
+
+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 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(
+    client: BatchCleanupClient,
+    batch_id: str,
+    *,
+    key: str,
+    provider: str | None = None,
+    delete_output_files: bool = False,
+    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:
+        if delete_output_files:
+            _cleanup_batch_outputs(client, fetched, key=key, provider=provider)
+        return
+    if fetched.status == "cancelling" and 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(
+            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:
+            if delete_output_files:
+                _cleanup_batch_outputs(client, current, key=key, provider=provider)
+            return
+        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

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):
@@ -85,7 +99,7 @@ class BatchList(BaseModel):
 class FileDeleteResponse(BaseModel):
     id: str
     object: str | None = None
-    deleted: bool
+    deleted: bool | None = None
 
 
 class BatchCreateBody(BaseModel):

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

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

tests/e2e/batches/test_batch_cleanup.py

@@ -0,0 +1,313 @@
+from builtins import ExceptionGroup
+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
+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
+
+MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE="
+MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x"
+
+
+class ExpectedCalls[T]:
+    def __init__(self, values: tuple[T, ...]) -> None:
+        self.values: Final = values
+        self.recorder: Final = Mock()
+
+    def __call__(self, value: T) -> None:
+        self.recorder(value)
+
+    def assert_done(self) -> None:
+        assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values)
+
+
+class CleanupClient:
+    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 self.file_response()
+
+    def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
+        self.calls(f"retrieve {provider} {batch_id}")
+        return self.batch_response()
+
+    def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
+        self.calls(f"cancel {provider} {batch_id}")
+        return self.cancel_response()
+
+    def generate_key(self, body: KeyGenerateBody) -> str:
+        return "test-key"
+
+    def delete_key(self, key: str) -> None:
+        self.calls(f"delete key {key}")
+
+    def delete_customers(self, user_ids: list[str]) -> None:
+        self.calls(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:
+    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(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((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")
+        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:
+        expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None
+        client: Final = CleanupClient(
+            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(("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()
+        manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure"))
+        with pytest.raises(ExceptionGroup) as caught:
+            manager.teardown()
+        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(
+            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")
+        client.calls.assert_done()
+
+    def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None:
+        client: Final = CleanupClient(
+            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(("delete None file-1", "delete key test-key")),
+            files=(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()
+        client.calls.assert_done()
+
+
+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:
+        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 = 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 outcomes.call_count == len(CLEANUP_DELAYS) + 1
+
+    def test_permanent_error_is_not_retried(self) -> None:
+        failure: Final = UnknownApiError(status_code=403, body="forbidden")
+        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 outcomes.call_count == 1
+
+
+class TestBatchCancellation:
+    def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None:
+        client: Final = CleanupClient(
+            calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3),
+            batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")),
+        )
+        delays: Final = ExpectedCalls((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(
+            calls=ExpectedCalls(
+                (
+                    f"retrieve None {MANAGED_BATCH_ID}",
+                    f"retrieve None {MANAGED_BATCH_ID}",
+                    "delete None file-1",
+                    "delete key test-key",
+                )
+            ),
+            batches=(batch("cancelling"), batch("cancelling")),
+            files=(deleted_file(),),
+        )
+        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=ticks))
+        with pytest.raises(ExceptionGroup) as caught:
+            manager.teardown()
+        assert "cancellation did not finish" in str(caught.value.exceptions[0])
+        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(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(("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()
+
+    @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(
+                (
+                    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=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")),
+            cancellations=(batch(pending_status),),
+            files=(deleted_file(),),
+        )
+        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"))
+        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(("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=(
+                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(
+            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")
+        else:
+            with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"):
+                cleanup_batch(client, "batch-1", key="test-key")
+        client.calls.assert_done()
+
+
+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"}

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 e2e_config import MASTER_KEY, 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,9 @@ 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, delete_output_files=cap.provider in {"openai", "azure"}
+        )
     )
 
     assert batch.id, f"create returned no batch id (body={created.body[:200]})"
@@ -339,7 +338,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 +355,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 +382,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 +458,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 +517,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 +559,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 +626,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 +690,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 +760,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)
 
@@ -803,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(
@@ -813,7 +813,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=cleanup_key or key))
         return file
 
     def _generate_enqueued_key(
@@ -850,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)
 
@@ -861,7 +861,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=MASTER_KEY, delete_output_files=True))
 
     @pytest.mark.covers(
         "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted",
@@ -904,7 +904,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 +928,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 +984,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 +1044,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 +1099,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 +1192,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 +1243,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 +1258,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 +1307,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 +1315,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}"

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)

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"

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

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,
         )

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,29 @@ 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
+    if not (result := call()).ok:
+        return result
+    while (
+        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
+
+
 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.
 

tests/e2e/guardrails/test_guardrails_client.py

@@ -0,0 +1,66 @@
+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)
+    responses: Final = iter((missing, missing, missing))
+
+    result: Final = poll_until_guardrail_applied(
+        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))
+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

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]}"

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

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(

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:

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

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

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

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

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,

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

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

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,
     )

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=[

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"<html/>"), NoBody)
+        assert isinstance(result, ValidationError)

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

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<void> {
+  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,
+  });
+}

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<string> {
-  const res = await request.post(`${rootPath()}/v1/chat/completions`, {
+const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise<APIResponse> =>
+  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<string> {
+  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<ChatAttempt> {
+  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<KeyInfo> {
+  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<void> {
+  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,

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<string, unknown>,
+): Promise<unknown> {
+  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] },
+      });
+    }
+  });
+});

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<void> {
+  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<void> {
+  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<boolean> {
+  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 });
+  });
+});

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<string, unknown>;
+
+async function readStoredParams(
+  page: PlaywrightPage,
+  modelId: string,
+): Promise<StoredParams> {
+  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<void> {
+  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<void> {
+  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}`,
+    });
+  });
+});

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<boolean> {
+  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<Locator> {
+  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<void> {
+  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<void> {
+  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<void> {
+  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<void>,
+): Promise<void> {
+  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");
+  });
+});

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) });
+  });
+});

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);
+  });
+});

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

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,104 @@ 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_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
+
+        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"<Error><Code>AccessDenied</Code></Error>")
+                )
+                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 +1973,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 +1989,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 +2239,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 +2254,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 +2279,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 +2479,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 +2502,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 +2557,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 +2604,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,

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"
+    }

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(

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}"}
+    ]

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.