Merge branch 'main' into litellm_strict_provider_identity

.github/e2e-stack/assert_tests_ran.py

@@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET
 from pathlib import Path
 from typing import Final
 
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
+from coverage_registry.management_cases import MANAGEMENT_CASES
+
 
 def main() -> int:
     selected: Final = tuple(sys.argv[2:])
@@ -16,6 +19,17 @@ def main() -> int:
         case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
     )
     missing: Final = tuple(path for path in selected if path not in passed)
+    required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
+    passed_nodes: Final = frozenset(
+        prop.get("value")
+        for case in cases
+        if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
+        for prop in case.findall("./properties/property")
+        if prop.get("name") == "management_node"
+    )
+    missing_nodes: Final = required_nodes - passed_nodes
+    for node in sorted(missing_nodes):
+        _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
     for path in selected:
         collected: Final = sum(case.get("file") == path for case in cases)
         skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
@@ -27,6 +41,7 @@ def main() -> int:
     if (
         selected
         and not missing
+        and not missing_nodes
         and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
     ):
         return 0

.github/e2e-stack/oidc-profile.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+set -euo pipefail
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "${REPO_ROOT}"
+exec uv run --no-sync python tests/e2e/idp.py "$@"

.github/e2e-stack/select_tests.py

@@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile(
 HARNESS: Final = re.compile(
     r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
     r"|^tests/e2e/idp_realm\.json$"
+    r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$"
+    r"|^tests/e2e/coverage_registry/management_cases\.py$"
     r"|^tests/e2e/gateway/"
     r"|^\.github/e2e-stack/"
     r"|^\.github/workflows/test-e2e-changed\.yml$"

.github/workflows/test-code-quality.yml

@@ -186,7 +186,7 @@ jobs:
           version: "0.10.9"
 
       - name: Install dependencies
-        run: uv sync --frozen --extra proxy --python 3.10
+        run: uv sync --frozen --extra proxy --extra cli --python 3.10
 
       - run: uv run --no-sync python --version
 
@@ -195,3 +195,6 @@ jobs:
 
       - name: Check litellm CLI
         run: uv run --no-sync litellm --version
+
+      - name: Check lite CLI
+        run: uv run --no-sync lite version

.github/workflows/test-e2e-changed.yml

@@ -183,7 +183,7 @@ jobs:
             log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
             echo "::group::pass ${pass} of 3"
             set +e
-            uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
+            uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \
               -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
             status=$?
             uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"

litellm/batches/batch_utils.py

@@ -9,6 +9,7 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
+from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
 from litellm.types.llms.openai import Batch
 from litellm.types.utils import ModelInfo, Usage
 from litellm.utils import token_counter
@@ -356,6 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
             prompt_tokens=_prompt,
             completion_tokens=_completion,
             total_tokens=_total,
+            prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata),
         )
 
         try:

litellm/cost_calculator.py

@@ -2278,6 +2278,19 @@ def default_video_cost_calculator(
     return 0.0
 
 
+def _batch_rate(
+    model_info: ModelInfo,
+    key: Literal[
+        "input_cost_per_audio_token_batches",
+        "input_cost_per_image_token_batches",
+        "input_cost_per_video_token_batches",
+    ],
+    fallback: float,
+) -> float:
+    rate: Final = model_info.get(key)
+    return fallback if rate is None else rate
+
+
 def batch_cost_calculator(
     usage: Usage,
     model: str,
@@ -2337,7 +2350,29 @@ def batch_cost_calculator(
     total_prompt_cost = 0.0
     total_completion_cost = 0.0
     if input_cost_per_token_batches is not None:
-        total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches
+        batch_details: Final = parse_prompt_tokens_details(usage)
+        audio_tokens, image_tokens, video_tokens = (
+            batch_details["audio_tokens"],
+            batch_details["image_tokens"],
+            batch_details["video_tokens"],
+        )
+        modality_rates: Final = (
+            _batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches),
+            _batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches),
+            _batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches),
+        )
+        total_prompt_cost = sum(
+            tokens * rate
+            for tokens, rate in zip(
+                (
+                    max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0),
+                    audio_tokens,
+                    image_tokens,
+                    video_tokens,
+                ),
+                (input_cost_per_token_batches, *modality_rates),
+            )
+        )
     elif input_cost_per_token:
         details: Final = parse_prompt_tokens_details(usage)
         cache_read_tokens: Final = details["cache_hit_tokens"]

litellm/integrations/custom_guardrail.py

@@ -884,7 +884,9 @@ class CustomGuardrail(CustomLogger):
         """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
         from litellm.llms import get_guardrail_translation_mapping
 
-        if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
+        if not self.uses_apply_guardrail_interface():
+            return kwargs, result
+        if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
             return kwargs, result
         try:
             translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
@@ -901,8 +903,18 @@ class CustomGuardrail(CustomLogger):
             for key, value in (litellm_params.get("metadata") or {}).items()
             if key != "standard_logging_guardrail_information"
         }
+        response: Final = (
+            kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
+        )
+        from litellm.types.utils import ModelResponse
+
+        output_translation: Final = (
+            get_guardrail_translation_mapping(CallTypes.acompletion)()
+            if isinstance(response, ModelResponse)
+            else translation
+        )
         try:
-            await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
+            await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
         except Exception as e:
             verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
         recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
@@ -919,8 +931,9 @@ class CustomGuardrail(CustomLogger):
     async def _scan_logged_call(
         self,
         kwargs: dict,  # mutable-ok: CustomLogger.async_logging_hook contract
-        result: object,
+        response: object | None,
         translation: "BaseTranslation",
+        output_translation: "BaseTranslation",
         scratch_metadata: dict,  # mutable-ok: apply_guardrail records its verdict into request metadata
     ) -> None:
         optional_params: Final = kwargs.get("optional_params") or {}
@@ -934,8 +947,10 @@ class CustomGuardrail(CustomLogger):
             "metadata": scratch_metadata,
         }
         await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
-        await translation.process_output_response(
-            response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
+        if response is None:
+            return
+        await output_translation.process_output_response(
+            response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
         )
 
     def supports_scan_only_tool_results(self) -> bool:

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -956,12 +956,16 @@ def _calculate_input_cost(
     )
 
     ### AUDIO COST
-    if prompt_tokens_details["audio_tokens"]:
+    if prompt_tokens_details["audio_tokens"] and not (
+        prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
+    ):
         audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
         prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
 
     ### IMAGE TOKEN COST
-    if prompt_tokens_details["image_tokens"]:
+    if prompt_tokens_details["image_tokens"] and not (
+        prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
+    ):
         # For image token costs:
         # First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
         image_token_cost_key = "input_cost_per_image_token"
@@ -970,7 +974,9 @@ def _calculate_input_cost(
         prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
 
     ### VIDEO TOKEN COST
-    if prompt_tokens_details["video_tokens"]:
+    if prompt_tokens_details["video_tokens"] and not (
+        prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
+    ):
         video_token_cost_key = "input_cost_per_video_token"
         if model_info.get(video_token_cost_key) is None:
             video_token_cost_key = "input_cost_per_token"

litellm/llms/vertex_ai/batches/transformation.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import Any, Final
 from urllib.parse import unquote
 
@@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import (
 )
 from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
 from litellm.types.llms.vertex_ai import *
-from litellm.types.utils import LiteLLMBatch
+from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper
+
+
+def vertex_prompt_tokens_details(
+    usage_metadata: Mapping[str, object],
+) -> PromptTokensDetailsWrapper | None:
+    raw_details: Final = usage_metadata.get("promptTokensDetails")
+    if not isinstance(raw_details, list):
+        return None
+
+    def _normalize(detail: object) -> tuple[str, int] | None:
+        if not isinstance(detail, Mapping):
+            return None
+        modality: Final = detail.get("modality")
+        token_count: Final = detail.get("tokenCount")
+        if not isinstance(modality, str) or not isinstance(token_count, int):
+            return None
+        return modality.upper(), token_count
+
+    parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
+    normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
+    if len(normalized) != len(parsed_details):
+        return None
+
+    return PromptTokensDetailsWrapper(
+        text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
+        audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
+        image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
+        video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
+    )
 
 
 class VertexAIBatchTransformation:

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py

@@ -298,8 +298,6 @@ def transform_openai_input_gemini_embed_content(
 
 
 _IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"})
-_VIDEO_TOKENS_PER_SECOND: Final = 258.0
-_AUDIO_TOKENS_PER_SECOND: Final = 32.0
 _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata)
 
 
@@ -339,11 +337,12 @@ def _is_image_element(
     return False
 
 
-def _count_input_images(
+def _is_image_only_input(
     input: GeminiEmbeddingInput,
     resolved_files: Mapping[str, Mapping[str, str]],
-) -> int:
-    return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files))
+) -> bool:
+    elements: Final = _flatten_input(input)
+    return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements)
 
 
 def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
@@ -372,30 +371,29 @@ def _usage_from_embed_content_response(
     total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens
 
     details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or ()
+    if not details:
+        return Usage(
+            prompt_tokens=prompt_tokens,
+            total_tokens=total_tokens,
+            prompt_tokens_details=PromptTokensDetailsWrapper(
+                text_tokens=0,
+                image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0,
+            ),
+        )
+
     text_tokens: Final = _tokens_for_modality(details, "TEXT")
     audio_tokens: Final = _tokens_for_modality(details, "AUDIO")
+    image_tokens: Final = _tokens_for_modality(details, "IMAGE")
     video_tokens: Final = _tokens_for_modality(details, "VIDEO")
-    image_count: Final = _count_input_images(input, resolved_files)
-
-    video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
-    audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
-
-    # generic_cost_per_token rewrites text_tokens to the full prompt minus
-    # other modalities when both text_tokens and image_count are zero. For
-    # video, that misallocates video tokens to text; a 1-token floor sidesteps
-    # the rewrite and keeps billing on input_cost_per_video_per_second.
-    needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0
-    resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens
 
     return Usage(
         prompt_tokens=prompt_tokens,
         total_tokens=total_tokens,
         prompt_tokens_details=PromptTokensDetailsWrapper(
-            text_tokens=resolved_text_tokens,
+            text_tokens=text_tokens,
             audio_tokens=audio_tokens,
-            image_count=image_count,
-            video_length_seconds=video_length_seconds,
-            audio_length_seconds=audio_length_seconds,
+            image_tokens=image_tokens,
+            video_tokens=video_tokens,
         ),
     )
 
@@ -415,8 +413,7 @@ def process_embed_content_response(
         model_response: EmbeddingResponse to populate
         model: Model name
         response_json: Raw JSON response from embedContent endpoint
-        resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
-            used to bill resolved image references at the per-image rate
+        resolved_files: Mapping of file references to resolved metadata
 
     Returns:
         EmbeddingResponse with single embedding

litellm/model_prices_and_context_window_backup.json

@@ -25601,10 +25601,14 @@
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
     },
     "gemini-embedding-2-preview": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai-embedding-models",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25615,13 +25619,14 @@
         "uses_embed_content": true
     },
     "gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
         "input_cost_per_audio_token": 6.5e-06,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token_batches": 3.25e-06,
         "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
         "input_cost_per_token_batches": 1e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai-embedding-models",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25633,10 +25638,14 @@
         "uses_embed_content": true
     },
     "vertex_ai/gemini-embedding-2-preview": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25648,13 +25657,14 @@
         "uses_embed_content": true
     },
     "vertex_ai/gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
         "input_cost_per_audio_token": 6.5e-06,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token_batches": 3.25e-06,
         "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
         "input_cost_per_token_batches": 1e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25693,10 +25703,14 @@
     },
     "gemini/gemini-embedding-2-preview": {
         "deprecation_date": "2026-08-10",
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "gemini",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25709,10 +25723,14 @@
         "tpm": 10000000
     },
     "gemini/gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "gemini",
         "max_input_tokens": 8192,
         "max_tokens": 8192,

litellm/proxy/_lazy_openapi_snapshot.json

@@ -11503,6 +11503,12 @@
               "description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
               "title": "Content Moderation Check"
             },
+            "contextual_grounding_from_messages": {
+              "default": false,
+              "description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
+              "title": "Contextual Grounding From Messages",
+              "type": "boolean"
+            },
             "credentials": {
               "anyOf": [
                 {

litellm/proxy/client/cli/commands/pi.py

@@ -10,7 +10,7 @@ import os
 import tempfile
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from enum import StrEnum
+from enum import Enum
 from pathlib import Path
 from types import MappingProxyType
 from typing import Annotated, Final
@@ -25,7 +25,7 @@ LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
 _REJECTED_STATUSES: Final = frozenset((401, 403))
 
 
-class ListingFailure(StrEnum):
+class ListingFailure(str, Enum):
     """Why a proxy could not be listed, decided once where the HTTP outcome is classified.
 
     `unreachable` means no response at all; the other kinds prove the proxy answered, so callers

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         prompt_attack_threshold: float | None = 0.5,
         pii_confidence_threshold: float | None = 0.5,
         chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
+        contextual_grounding_from_messages: bool = False,
         streaming_buffer_until_moderated: bool | None = None,
         streaming_sampling_rate: int | None = None,
         streaming_end_of_stream_only: bool | None = None,
@@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self.guardrailVersion = guardrailVersion
         self.guardrail_provider = "bedrock"
         self.chunk_budget_chars = chunk_budget_chars
+        self.contextual_grounding_from_messages = contextual_grounding_from_messages
         self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
 
         # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         """
         Flatten a message into text blocks, preserving any contextual-grounding
         qualifier carried by the content-block ``type`` (grounding_source / query).
-        Untagged text keeps ``qualifier=None`` so the payload is unchanged for
-        callers that do not use grounding.
+        Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
+        derive grounding qualifiers from it.
         """
         content: Final = message.get("content")
         if content is None:
@@ -493,6 +495,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         result carrying externally-influenced content can supply fake evidence for the
         contextual-grounding check to grade the response against. ``query`` is accepted
         from any role (it is the user's question).
+
+        With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
+        falls back to the plain messages: system / developer text is the grounding
+        source and the latest user message is the query.
         """
         grounding: Final[list[QualifiedTextBlock]] = []
         for message in messages or []:
@@ -504,7 +510,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                     and role in _GROUNDING_SOURCE_TRUSTED_ROLES
                 ):
                     grounding.append(block)
-        return grounding
+        if grounding or not self.contextual_grounding_from_messages:
+            return grounding
+        return self._derive_grounding_blocks_from_plain_messages(messages)
+
+    def _derive_grounding_blocks_from_plain_messages(
+        self, messages: list[AllMessageValues] | None
+    ) -> list[QualifiedTextBlock]:
+        if not messages:
+            return []
+        latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")
+        if latest_user_index is None:
+            return []
+        sources: Final = tuple(
+            QualifiedTextBlock(text=block.text, qualifier="grounding_source")
+            for message in messages
+            if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES
+            for block in self.get_content_items_for_message(message=message) or []
+            if block.text
+        )
+        queries: Final = tuple(
+            QualifiedTextBlock(text=block.text, qualifier="query")
+            for block in self.get_content_items_for_message(message=messages[latest_user_index]) or []
+            if block.text
+        )
+        if not sources or not queries:
+            return []
+        return [*sources, *queries]
 
     def supports_scan_only_tool_results(self) -> bool:
         return self.experimental_use_latest_role_message_only is not True
@@ -3210,6 +3242,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                     bedrock_response = await self.make_bedrock_api_request(
                         source="OUTPUT",
                         response=synthetic_response,
+                        messages=request_data.get("messages"),
                         request_data=request_data,
                         logging_event_type=_log_hook,
                     )

litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py

@@ -1,11 +1,13 @@
-from collections.abc import AsyncGenerator, Mapping, Sequence
+import time
+from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, Final, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
 
 import httpx
 from fastapi import HTTPException
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 import json
@@ -23,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import (
     get_or_create_metadata_bucket,
 )
 from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
     get_async_httpx_client,
     httpxSpecialProvider,
 )
@@ -52,6 +55,7 @@ from litellm.types.utils import (
     CallTypes,
     CallTypesLiteral,
     Choices,
+    GenericGuardrailAPIInputs,
     GuardrailStatus,
     ModelResponse,
     ModelResponseStream,
@@ -118,8 +122,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
     Supports:
     - Pre-call sanitization (sanitizeUserPrompt)
     - Post-call sanitization (sanitizeModelResponse)
+    - logging_only: scans the completed response after it reaches the client and
+      records the verdict in spend logs without blocking
     """
 
+    use_native_lifecycle_hooks: ClassVar[bool] = True
+
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
         return [
@@ -128,6 +136,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             GuardrailEventHooks.post_call,
             GuardrailEventHooks.pre_mcp_call,
             GuardrailEventHooks.during_mcp_call,
+            GuardrailEventHooks.logging_only,
         ]
 
     def __init__(
@@ -138,6 +147,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         credentials: VERTEX_CREDENTIALS_TYPES | None = None,
         api_endpoint: str | None = None,
         sanitize_error_detail: "bool | None" = True,
+        async_handler: AsyncHTTPHandler | None = None,
+        access_token_provider: Callable[[], Awaitable[tuple[str, str]]] | None = None,
         **kwargs,
     ):
         # Set supported event hooks if not already provided
@@ -154,7 +165,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         VertexBase.__init__(self)
 
         # Then set our attributes (this ensures project_id is not overwritten)
-        self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
+        self.async_handler = async_handler or get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.GuardrailCallback
+        )
+        self.access_token_provider = access_token_provider
         self.template_id = template_id
         self.project_id = project_id
         self.location = location or "us-central1"
@@ -278,11 +292,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         If file_bytes and file_type are provided, file prompt sanitization is performed.
         """
         # Get access token using VertexBase auth
-        access_token, resolved_project_id = await self._ensure_access_token_async(
-            credentials=self.credentials,
-            project_id=self.project_id,
-            custom_llm_provider="vertex_ai",
-        )
+        if self.access_token_provider is not None:
+            access_token, resolved_project_id = await self.access_token_provider()
+        else:
+            access_token, resolved_project_id = await self._ensure_access_token_async(
+                credentials=self.credentials,
+                project_id=self.project_id,
+                custom_llm_provider="vertex_ai",
+            )
 
         # Use resolved project ID if not explicitly set
         if not self.project_id and resolved_project_id:
@@ -1096,6 +1113,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             add_guardrail_to_applied_guardrails_header,
         )
 
+        if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
+            async for chunk in response:
+                yield chunk
+            return
+
         all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response])
 
         if not all_chunks or self._is_terminal_error_stream(all_chunks):
@@ -1213,6 +1235,60 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         for chunk in all_chunks:
             yield chunk
 
+    @log_guardrail_information
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        content: Final = "\n".join(text for text in inputs.get("texts") or () if text)
+        if not content:
+            return inputs
+
+        source: Final[Literal["user_prompt", "model_response"]] = (
+            "user_prompt" if input_type == "request" else "model_response"
+        )
+        start_time: Final = time.time()
+        try:
+            armor_response: Final = await self.make_model_armor_request(
+                content=content, source=source, request_data=request_data
+            )
+        except (ModelArmorAPIError, httpx.HTTPError) as e:
+            error_end_time: Final = time.time()
+            self.add_standard_logging_guardrail_information_to_request_data(
+                guardrail_json_response=str(e),
+                request_data=request_data,
+                guardrail_status="guardrail_failed_to_respond",
+                guardrail_provider="model_armor",
+                start_time=start_time,
+                end_time=error_end_time,
+                duration=error_end_time - start_time,
+            )
+            return inputs
+
+        flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
+        end_time: Final = time.time()
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=self._build_logging_response(armor_response),
+            request_data=request_data,
+            guardrail_status="guardrail_flagged" if flagged else "success",
+            guardrail_provider="model_armor",
+            start_time=start_time,
+            end_time=end_time,
+            duration=end_time - start_time,
+        )
+        if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
+            raise HTTPException(
+                status_code=400,
+                detail=self._build_block_error_detail(
+                    "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor",
+                    armor_response,
+                ),
+            )
+        return inputs
+
     @staticmethod
     def get_config_model() -> type["GuardrailConfigModel"] | None:
         """

litellm/proxy/guardrails/guardrail_initializers.py

@@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
         prompt_attack_threshold=litellm_params.prompt_attack_threshold,
         pii_confidence_threshold=litellm_params.pii_confidence_threshold,
         chunk_budget_chars=litellm_params.chunk_budget_chars,
+        contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages,
         default_on=litellm_params.default_on,
         disable_exception_on_block=litellm_params.disable_exception_on_block,
         mask_request_content=litellm_params.mask_request_content,

litellm/responses/streaming_iterator.py

@@ -221,6 +221,13 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None
     )
 
 
+def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
+    if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
+        return True
+    status_code: Final = getattr(mapped_exception, "status_code", None)
+    return not isinstance(status_code, int) or status_code >= 500 or status_code == 429
+
+
 class BaseResponsesAPIStreamingIterator:
     """
     Base class for streaming iterators that process responses from the Responses API.
@@ -521,15 +528,8 @@ class BaseResponsesAPIStreamingIterator:
             getattr(self.completed_response, "response", None) if self.completed_response else None
         )
         error_info: Final = getattr(response_obj, "error", None) if response_obj else None
-        error_message, error_type, error_code = _error_event_fields(error_info)
         self._record_failed_response_usage(response_obj)
-        exception: Final = litellm.APIError(
-            status_code=_status_code_for_error_fields(error_type, error_code),
-            message=error_message,
-            llm_provider=self.custom_llm_provider or "",
-            model=self.model or "",
-        )
-        self._handle_failure(exception)
+        self._handle_failure(self._map_error_event_exception(error_info))
 
     def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
         if response_obj is None or self.logging_obj is None:
@@ -551,6 +551,28 @@ class BaseResponsesAPIStreamingIterator:
             self.logging_obj._response_cost_calculator(result=response_obj) or 0.0
         )
 
+    def _map_error_event_exception(self, error_obj: object) -> Exception:
+        from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+        error_message, error_type, error_code = _error_event_fields(error_obj)
+        status_code: Final = _status_code_for_error_fields(error_type, error_code)
+        error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
+        provider_exception: Final = BaseLLMException(
+            status_code=status_code,
+            message=f"Error code: {status_code} - {{'error': {error_body}}}",
+            body=error_body,
+        )
+        try:
+            return litellm.exception_type(
+                model=self.model or "",
+                custom_llm_provider=self.custom_llm_provider or "",
+                original_exception=provider_exception,
+                completion_kwargs={},
+                extra_kwargs={},
+            )
+        except Exception as mapped_exception:
+            return mapped_exception
+
     def _maybe_raise_for_error_event(self, result: object) -> None:
         chunk_type: Final = getattr(result, "type", None)
         if chunk_type not in ("error", "response.failed"):
@@ -562,15 +584,8 @@ class BaseResponsesAPIStreamingIterator:
             else getattr(result, "error", None)
         )
 
-        error_message, error_type, error_code = _error_event_fields(error_obj)
-        status_code: Final = _status_code_for_error_fields(error_type, error_code)
-        mapped_exception: Final = litellm.APIError(
-            status_code=status_code,
-            message=error_message,
-            llm_provider=self.custom_llm_provider or "",
-            model=self.model or "",
-        )
-        if 400 <= status_code < 500 and status_code != 429:
+        mapped_exception: Final = self._map_error_event_exception(error_obj)
+        if not _mid_stream_fallback_eligible(mapped_exception):
             raise mapped_exception
         raise MidStreamFallbackError(
             message=str(mapped_exception),

litellm/router.py

@@ -3268,8 +3268,15 @@ class Router:
                         kwargs=initial_kwargs,
                         metadata_variable_name="litellm_metadata",
                     )
+                    # The content-policy dispatch branch matches on the trigger's own type, so a refusal's
+                    # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted.
+                    fallback_trigger: Final[Exception] = (
+                        e.original_exception
+                        if isinstance(e.original_exception, litellm.ContentPolicyViolationError)
+                        else e
+                    )
                     fallback_response = await self.async_function_with_fallbacks_common_utils(
-                        e=e,
+                        e=fallback_trigger,
                         disable_fallbacks=False,
                         fallbacks=fallbacks,
                         context_window_fallbacks=context_window_fallbacks,

litellm/types/guardrails.py

@@ -552,6 +552,16 @@ class BedrockGuardrailConfigModel(BaseModel):
         "still rejects is bisected automatically, so this value only trades round trips against "
         "batch size and cannot fail a request on its own.",
     )
+    contextual_grounding_from_messages: bool = Field(
+        default=False,
+        description="ApplyGuardrail: when True, post-call scans of a request with no grounding_source / "
+        "query content parts send the system and developer messages as the grounding source and "
+        "the latest user message as the query, so the guardrail's contextual grounding policy can "
+        "score the response. Bedrock bills contextual grounding units for these scans and rejects "
+        "queries, sources and responses over its contextual grounding length limits, so leave this "
+        "off for guardrails without a contextual grounding policy. Default False: plain messages "
+        "are never sent as grounding context.",
+    )
 
 
 class BedrockGuardrailStreamingParams(BaseModel):

litellm/types/utils.py

@@ -283,8 +283,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     input_cost_per_video_token: float | None  # for gemini omni models with video input
     input_cost_per_audio_per_second: float | None  # only for vertex ai models
     input_cost_per_video_per_second: float | None  # only for vertex ai models
+    input_cost_per_audio_token_batches: ReadOnly[float | None]
+    input_cost_per_image_token_batches: ReadOnly[float | None]
     input_cost_per_second: float | None  # for OpenAI Speech models
     input_cost_per_token_batches: float | None
+    input_cost_per_video_token_batches: ReadOnly[float | None]
     output_cost_per_token_batches: float | None
     output_cost_per_token: Required[float | None]
     output_cost_per_token_flex: float | None  # OpenAI flex service tier pricing
@@ -3583,7 +3586,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     input_cost_per_video_per_second_above_128k_tokens: float | None = None
     input_cost_per_video_per_second_above_15s_interval: float | None = None
     input_cost_per_video_per_second_above_8s_interval: float | None = None
+    input_cost_per_audio_token_batches: float | None = None
+    input_cost_per_image_token_batches: float | None = None
     input_cost_per_token_batches: float | None = None
+    input_cost_per_video_token_batches: float | None = None
     output_cost_per_token_batches: float | None = None
     output_cost_per_token_flex: float | None = None
     output_cost_per_token_priority: float | None = None

litellm/utils.py

@@ -5923,10 +5923,13 @@ def _get_model_info_helper(
                 input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
                 input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
                 input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
+                input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None),
+                input_cost_per_image_token_batches=_model_info.get("input_cost_per_image_token_batches", None),
                 input_cost_per_image=_model_info.get("input_cost_per_image", None),
                 input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None),
                 input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None),
                 input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"),
+                input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None),
                 output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"),
                 output_cost_per_token=_output_cost_per_token,
                 output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None),

model_prices_and_context_window.json

@@ -25601,10 +25601,14 @@
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
     },
     "gemini-embedding-2-preview": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai-embedding-models",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25615,13 +25619,14 @@
         "uses_embed_content": true
     },
     "gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
         "input_cost_per_audio_token": 6.5e-06,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token_batches": 3.25e-06,
         "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
         "input_cost_per_token_batches": 1e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai-embedding-models",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25633,10 +25638,14 @@
         "uses_embed_content": true
     },
     "vertex_ai/gemini-embedding-2-preview": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25648,13 +25657,14 @@
         "uses_embed_content": true
     },
     "vertex_ai/gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
         "input_cost_per_audio_token": 6.5e-06,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token_batches": 3.25e-06,
         "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
         "input_cost_per_token_batches": 1e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "vertex_ai",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25693,10 +25703,14 @@
     },
     "gemini/gemini-embedding-2-preview": {
         "deprecation_date": "2026-08-10",
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "gemini",
         "max_input_tokens": 8192,
         "max_tokens": 8192,
@@ -25709,10 +25723,14 @@
         "tpm": 10000000
     },
     "gemini/gemini-embedding-2": {
-        "input_cost_per_audio_per_second": 0.00016,
-        "input_cost_per_image": 0.00012,
+        "input_cost_per_audio_token": 6.5e-06,
+        "input_cost_per_audio_token_batches": 3.25e-06,
+        "input_cost_per_image_token": 4.5e-07,
+        "input_cost_per_image_token_batches": 2.25e-07,
         "input_cost_per_token": 2e-07,
-        "input_cost_per_video_per_second": 0.00079,
+        "input_cost_per_token_batches": 1e-07,
+        "input_cost_per_video_token": 1.2e-05,
+        "input_cost_per_video_token_batches": 6e-06,
         "litellm_provider": "gemini",
         "max_input_tokens": 8192,
         "max_tokens": 8192,

model_prices_and_context_window.schema.json

@@ -249,6 +249,10 @@
           "type": "number",
           "minimum": 0
         },
+        "input_cost_per_audio_token_batches": {
+          "type": "number",
+          "minimum": 0
+        },
         "input_cost_per_audio_token_priority": {
           "type": "number",
           "minimum": 0,
@@ -276,6 +280,10 @@
           "type": "number",
           "minimum": 0
         },
+        "input_cost_per_image_token_batches": {
+          "type": "number",
+          "minimum": 0
+        },
         "input_cost_per_pixel": {
           "type": "number",
           "minimum": 0
@@ -375,6 +383,14 @@
           "minimum": 0,
           "description": "Rate applied once the prompt exceeds the token threshold in the field name."
         },
+        "input_cost_per_video_token": {
+          "type": "number",
+          "minimum": 0
+        },
+        "input_cost_per_video_token_batches": {
+          "type": "number",
+          "minimum": 0
+        },
         "input_dbu_cost_per_token": {
           "type": "number",
           "minimum": 0

tests/code_coverage_tests/test_e2e_changed_gate.py

@@ -81,6 +81,25 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None
     assert result.returncode == 1
 
 
+@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer"))
+def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None:
+    suite: Final = ET.Element("testsuite")
+    path: Final = "tests/e2e/management/test_jwt_management_e2e.py"
+    case: Final = ET.SubElement(suite, "testcase", file=path)
+    properties: Final = ET.SubElement(case, "properties")
+    _ = ET.SubElement(
+        properties,
+        "property",
+        name="management_node",
+        value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]",
+    )
+    report: Final = tmp_path / "report.xml"
+    ET.ElementTree(suite).write(report)
+    result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True)
+    assert result.returncode == 1
+    assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout
+
+
 def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None:
     env_path: Final = tmp_path / ".env"
 
@@ -141,6 +160,10 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them(
     (
         "tests/e2e/proxy_client.py",
         "tests/e2e/conftest.py",
+        "tests/e2e/management/management_client.py",
+        "tests/e2e/management/jwt_actors.py",
+        "tests/e2e/management/conftest.py",
+        "tests/e2e/coverage_registry/management_cases.py",
         "tests/e2e/pytest.ini",
         "tests/e2e/gateway/stage_mirror_ci_config.yml",
         ".github/e2e-stack/up.sh",

tests/e2e/CONTRIBUTING.md

@@ -60,7 +60,11 @@ The suites run against a live proxy, so bring one up first by running the litell
 
    Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack.
 
-   Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write.
+   Management tests can bind a credential once with `client.with_caller(Caller(...))`; direct calls, delegated helpers and replica read-backs then retain that caller. Explicit `caller_key` arguments override the binding. Keep the original master-backed client for bootstrap and cleanup. `actor_factory` lazily provisions database roles and tenant memberships, with `database_role` tokens carrying no groups and `group_scoped` actors retaining the existing team route gate. Token minting is explicit through `actor.mint_caller(idp)`. The factory runs requests without backend retries and reports cleanup failures. `coverage_registry/management_cases.py` records exact canary nodes and non-secret actor labels; the CI execution assertion rejects a missing or skipped actor row
+
+   For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" <server-command>`. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts
+
+   `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step
 
    Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`:
 

tests/e2e/coverage_registry/management_cases.py

@@ -0,0 +1,151 @@
+from dataclasses import dataclass
+from typing import Final, Literal
+
+CredentialKind = Literal["master", "idp_admin", "direct_jwt", "virtual_key", "dashboard_session"]
+DependencyProfile = Literal["management_only", "real_oidc_browser", "external_provider_required"]
+
+
+@dataclass(frozen=True, slots=True)
+class ManagementCase:
+    node: str
+    credential_kind: CredentialKind
+    actor: str
+    profile: str
+    method: Literal["GET", "POST"]
+    path: str
+    operation_family: str
+    dependency_profile: DependencyProfile = "management_only"
+
+
+JWT_FILE: Final = "tests/e2e/management/test_jwt_management_e2e.py"
+JWT_CLASS: Final = f"{JWT_FILE}::TestJwtManagement"
+ACTORS: Final = (
+    "proxy_admin",
+    "proxy_admin_viewer",
+    "organization_admin",
+    "team_admin",
+    "team_member",
+    "internal_user",
+    "internal_user_viewer",
+    "unrelated_user",
+)
+MANAGEMENT_CASES: Final = tuple(
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_actor_subject_and_database_role[{role}]",
+        credential_kind="direct_jwt",
+        actor=role,
+        profile="database_role",
+        method="GET",
+        path="/user/info",
+        operation_family="identity",
+    )
+    for role in ACTORS
+) + (
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_admin_viewer_reads_but_cannot_update",
+        credential_kind="direct_jwt",
+        actor="proxy_admin_viewer",
+        profile="database_role",
+        method="POST",
+        path="/key/update",
+        operation_family="denial",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[direct_jwt]",
+        credential_kind="direct_jwt",
+        actor="proxy_admin",
+        profile="group_scoped",
+        method="POST",
+        path="/key/generate",
+        operation_family="key_lifecycle",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]",
+        credential_kind="virtual_key",
+        actor="proxy_admin",
+        profile="group_scoped",
+        method="POST",
+        path="/key/generate",
+        operation_family="key_lifecycle",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_two_actor_sets_keep_tenants_and_keys_isolated",
+        credential_kind="direct_jwt",
+        actor="team_member",
+        profile="group_scoped",
+        method="GET",
+        path="/key/info",
+        operation_family="tenant_isolation",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_member_cannot_write_and_another_team_cannot_read_the_key",
+        credential_kind="direct_jwt",
+        actor="team_member",
+        profile="group_scoped",
+        method="POST",
+        path="/key/update",
+        operation_family="tenant_isolation",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_multi_group_actor_keeps_exact_memberships",
+        credential_kind="master",
+        actor="bootstrap",
+        profile="group_scoped",
+        method="GET",
+        path="/team/info",
+        operation_family="memberships",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_successful_actor_cleanup_removes_owned_state",
+        credential_kind="master",
+        actor="bootstrap",
+        profile="failure_cleanup",
+        method="GET",
+        path="/team/info",
+        operation_family="cleanup",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[group]",
+        credential_kind="idp_admin",
+        actor="idp_admin",
+        profile="failure_cleanup",
+        method="POST",
+        path="/groups",
+        operation_family="cleanup",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[user]",
+        credential_kind="idp_admin",
+        actor="idp_admin",
+        profile="failure_cleanup",
+        method="POST",
+        path="/users",
+        operation_family="cleanup",
+    ),
+    ManagementCase(
+        node=f"{JWT_CLASS}::test_oidc_browser_profile_identity_mapping",
+        credential_kind="direct_jwt",
+        actor="internal_user",
+        profile="oidc_configuration",
+        method="GET",
+        path="/protocol/openid-connect/userinfo",
+        operation_family="oidc_identity",
+    ),
+)
+
+
+def canonical_node(node: str) -> str:
+    return node if node.startswith("tests/e2e/") else f"tests/e2e/{node}"
+
+
+def case_properties(node: str) -> tuple[tuple[str, str], ...]:
+    case: Final = next((case for case in MANAGEMENT_CASES if case.node == canonical_node(node)), None)
+    if case is None:
+        return ()
+    return (
+        ("management_node", case.node),
+        ("credential_kind", case.credential_kind),
+        ("actor", case.actor),
+        ("auth_profile", case.profile),
+        ("dependency_profile", case.dependency_profile),
+    )

tests/e2e/coverage_registry/mgmt.yaml

@@ -90,3 +90,11 @@
 - {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"}
+
+- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"}
+- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"}
+- {id: mgmt.user.oidc.identity_mapping, module: mgmt, tier: P0, surface: api, assertions: [identity_mapping], source: "tests/e2e/idp.py", rationale: "IdP configuration canary only: confidential-client token and userinfo subjects match the seeded user; application SSO is separate"}
+- {id: mgmt.team.jwt.tenant_isolation, module: mgmt, tier: P0, surface: api, assertions: [tenant_isolation], source: "auth/handle_jwt.py", rationale: "Isolated team actors read their own key and receive 403 for the other tenant key"}
+- {id: mgmt.team.jwt.multiple_memberships, module: mgmt, tier: P0, surface: api, assertions: [multiple_memberships], source: "auth/handle_jwt.py", rationale: "A multi-group actor has exactly the configured memberships without admin scope"}
+- {id: mgmt.user.jwt.cleanup, module: mgmt, tier: P0, surface: api, assertions: [cleanup], source: "management_endpoints/internal_user_endpoints.py", rationale: "Owned users teams organizations keys and IdP objects disappear after successful cleanup"}
+- {id: mgmt.user.jwt.partial_cleanup, module: mgmt, tier: P0, surface: api, assertions: [partial_cleanup], source: "auth/handle_jwt.py", rationale: "Partial identity setup removes the group and user created before failure"}

tests/e2e/e2e_http.py

@@ -16,9 +16,11 @@ requests itself imports.
 from __future__ import annotations
 
 import time
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Generator, Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast
+from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast
 
 import pytest
 import requests
@@ -36,8 +38,8 @@ class Headers(BaseModel):
 
 class AuthHeaders(Headers):
     # litellm accepts either; set whichever the call needs, leave the other None.
-    authorization: str | None = None
-    x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key")
+    authorization: str | None = Field(default=None, repr=False)
+    x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False)
 
 
 class AnthropicHeaders(AuthHeaders):
@@ -292,6 +294,22 @@ def _params(params: BaseModel | None) -> dict[str, str]:
 
 TRANSIENT_STATUSES: frozenset[int] = frozenset({529})
 RETRY_ATTEMPTS: int = 3
+_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False)
+
+
+def retry_attempts(default: int) -> int:
+    return 1 if _QUALIFICATION.get() else default
+
+
+@contextmanager
+def without_retries() -> Generator[None]:
+    token: Final = _QUALIFICATION.set(True)
+    try:
+        yield
+    finally:
+        _QUALIFICATION.reset(token)
+
+
 RETRY_BACKOFF_SECONDS: float = 0.5
 
 
@@ -319,7 +337,7 @@ def request_with_retry[T: RetryableResponse](
     hang should surface as a hang instead of doubling the wall clock. Every
     retry prints, so flakiness stays visible in the run log instead of
     vanishing into green."""
-    for attempt in range(1, RETRY_ATTEMPTS):
+    for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)):
         resp = issue()
         if resp.status_code not in TRANSIENT_STATUSES:
             return resp
@@ -414,6 +432,7 @@ def get_external[R: BaseModel](
     url: str,
     *,
     response_type: type[R],
+    headers: BaseModel | None = None,
     timeout: float = 30.0,
 ) -> Result[R]:
     """GET an absolute URL outside the proxy (e.g. a public /.well-known document).
@@ -422,7 +441,7 @@ def get_external[R: BaseModel](
     try:
         resp = requests.get(
             url,
-            headers={"Accept": "application/json"},
+            headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})},
             timeout=timeout,
         )
     except requests.RequestException as exc:

tests/e2e/idp.py

@@ -2,11 +2,18 @@
 
 from __future__ import annotations
 
+import base64
 import os
 import secrets
+import signal
+import subprocess
+import sys
+import time
 import warnings
 from collections.abc import Callable
-from dataclasses import dataclass, field
+from contextlib import ExitStack
+from dataclasses import dataclass, field, replace
+from types import FrameType
 from typing import Final, Literal
 
 import pytest
@@ -14,11 +21,15 @@ from e2e_http import (
     AuthHeaders,
     ExternalWrite,
     NetworkError,
+    NoBody,
     Result,
     Success,
+    UnknownApiError,
     delete_external,
+    get_external,
     post_form_external,
     post_json_external,
+    unwrap,
 )
 from pydantic import BaseModel, Field
 
@@ -46,7 +57,9 @@ class TokenGrantForm(BaseModel):
     grant_type: Literal["password"] = "password"
     client_id: str
     username: str
-    password: str
+    password: str = Field(repr=False)
+    client_secret: str | None = Field(default=None, repr=False)
+    scope: str | None = None
 
 
 class TokenResponse(BaseModel):
@@ -63,7 +76,7 @@ class GroupCreateBody(BaseModel):
 
 class PasswordCredential(BaseModel):
     type: Literal["password"] = "password"
-    value: str
+    value: str = Field(repr=False)
     temporary: bool = False
 
 
@@ -101,8 +114,20 @@ class Identity:
     user_id: str
     username: str
     password: str = field(repr=False)
-    group: str
-    group_id: str
+    groups: tuple[str, ...]
+    group_ids: tuple[str, ...]
+
+    @property
+    def group(self) -> str:
+        if len(self.groups) != 1:
+            raise ValueError("A single-group identity is required")
+        return self.groups[0]
+
+    @property
+    def group_id(self) -> str:
+        if len(self.group_ids) != 1:
+            raise ValueError("A single-group identity is required")
+        return self.group_ids[0]
 
 
 @dataclass(frozen=True, slots=True)
@@ -111,6 +136,10 @@ class Keycloak:
     realm: str
     admin_username: str
     admin_password: str = field(repr=False)
+    strict_cleanup: bool = False
+
+    def with_strict_cleanup(self) -> Keycloak:
+        return replace(self, strict_cleanup=True)
 
     @property
     def issuer(self) -> str:
@@ -150,7 +179,9 @@ class Keycloak:
             f"group {name}",
         )
 
-    def create_user(self, *, username: str, email: str, password: str, group: str) -> str:
+    def create_user(
+        self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = ()
+    ) -> str:
         return created_id(
             post_json_external(
                 self._admin_url("/users"),
@@ -158,7 +189,7 @@ class Keycloak:
                 json=UserCreateBody(
                     username=username,
                     email=email,
-                    groups=(group,),
+                    groups=(group,) if group is not None else groups,
                     credentials=(PasswordCredential(value=password),),
                 ),
             ),
@@ -171,14 +202,28 @@ class Keycloak:
     def delete_group(self, group_id: str) -> None:
         self._delete(f"/groups/{group_id}")
 
+    def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None:
+        result: Final = get_external(
+            self._admin_url(f"/{kind}/{resource_id}"),
+            headers=self._admin_headers(),
+            response_type=NoBody,
+        )
+        assert isinstance(result, UnknownApiError) and result.status_code == 404, (
+            f"Owned IdP {kind} still exists: {result}"
+        )
+
     def _delete(self, path: str) -> None:
         try:
             headers: Final = self._admin_headers()
         except pytest.fail.Exception as exc:
+            if self.strict_cleanup:
+                raise RuntimeError(f"Keycloak cleanup could not authenticate for {path}") from exc
             warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2)
             return
         result: Final = delete_external(self._admin_url(path), headers=headers)
         if result.status_code not in (204, 404):
+            if self.strict_cleanup:
+                raise RuntimeError(f"Keycloak cleanup failed for {path}: HTTP {result.status_code}")
             warnings.warn(
                 f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}",
                 RuntimeWarning,
@@ -188,15 +233,34 @@ class Keycloak:
     def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity:
         """Create `group` and a user in it, credentialed with a password generated
         for this test alone, and hand back the identity a token can be minted for."""
-        group_id: Final = self.create_group(group)
-        defer(lambda: self.delete_group(group_id))
+        return self.provision_groups(marker=marker, groups=(group,), defer=defer)
+
+    def provision_groups(
+        self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None]
+    ) -> Identity:
+        def provision_group(name: str) -> str:
+            created: Final = self.create_group(name)
+            defer(lambda: self.delete_group(created))
+            return created
+
+        group_ids: Final = tuple(provision_group(group) for group in groups)
+        return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer)
+
+    def provision_user(
+        self,
+        *,
+        marker: str,
+        groups: tuple[str, ...],
+        group_ids: tuple[str, ...],
+        defer: Callable[[Callable[[], object]], None],
+    ) -> Identity:
         username: Final = f"e2e-jwt-user-{marker}"
         password: Final = secrets.token_urlsafe(24)
         user_id: Final = self.create_user(
-            username=username, email=f"{username}@example.com", password=password, group=group
+            username=username, email=f"{username}@example.com", password=password, groups=groups
         )
         defer(lambda: self.delete_user(user_id))
-        return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id)
+        return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids)
 
     def access_token(
         self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None
@@ -211,6 +275,65 @@ class Keycloak:
         )
         return self._token(result, f"a token for {identity.username}")
 
+    def discovery(self) -> Discovery:
+        return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery))
+
+    def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient:
+        client: Final = BrowserClient(
+            client_id=f"e2e-browser-{secrets.token_hex(8)}",
+            secret=secrets.token_urlsafe(32),
+            callback_url=callback_url,
+        )
+        resource_id: Final = created_id(
+            post_json_external(
+                self._admin_url("/clients"),
+                headers=self._admin_headers(),
+                json=BrowserClientBody(
+                    clientId=client.client_id,
+                    secret=client.secret,
+                    redirectUris=(callback_url,),
+                ),
+            ),
+            "browser client",
+        )
+        defer(lambda: self._delete(f"/clients/{resource_id}"))
+        configured: Final = unwrap(
+            get_external(
+                self._admin_url(f"/clients/{resource_id}"),
+                headers=self._admin_headers(),
+                response_type=BrowserClientBody,
+            )
+        )
+        assert configured.redirect_uris == (callback_url,)
+        assert configured.standard_flow_enabled and not configured.public_client
+        assert configured.attributes.pkce == "S256"
+        return client
+
+    def browser_token(self, identity: Identity, client: BrowserClient) -> str:
+        return self._token(
+            post_form_external(
+                self.token_url(self.realm),
+                form=TokenGrantForm(
+                    client_id=client.client_id,
+                    client_secret=client.secret,
+                    username=identity.username,
+                    password=identity.password,
+                    scope="openid email",
+                ),
+                response_type=TokenResponse,
+            ),
+            "browser-profile identity mapping",
+        )
+
+    def userinfo(self, token: str) -> UserInfo:
+        return unwrap(
+            get_external(
+                f"{self.issuer}/protocol/openid-connect/userinfo",
+                headers=AuthHeaders(authorization=f"Bearer {token}"),
+                response_type=UserInfo,
+            )
+        )
+
 
 def keycloak_from_env() -> Keycloak:
     admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip()
@@ -226,3 +349,137 @@ def keycloak_from_env() -> Keycloak:
         admin_username=admin_username,
         admin_password=admin_password,
     )
+
+
+class TokenClaims(BaseModel):
+    sub: str
+    iss: str
+    aud: str | tuple[str, ...]
+    exp: int
+    scope: str = ""
+    groups: tuple[str, ...] = ()
+
+
+class Discovery(BaseModel):
+    issuer: str
+    authorization_endpoint: str
+    token_endpoint: str
+    userinfo_endpoint: str
+    jwks_uri: str
+
+
+class UserInfo(BaseModel):
+    sub: str
+    email: str
+
+
+class BrowserAttributes(BaseModel):
+    pkce: str = Field(default="S256", alias="pkce.code.challenge.method")
+
+
+class AudienceConfig(BaseModel):
+    audience: str = Field(default="litellm-e2e", alias="included.custom.audience")
+    access_token: str = Field(default="true", alias="access.token.claim")
+    id_token: str = Field(default="false", alias="id.token.claim")
+
+
+class AudienceMapper(BaseModel):
+    name: str = "litellm-audience"
+    protocol: str = "openid-connect"
+    mapper: str = Field(default="oidc-audience-mapper", alias="protocolMapper")
+    config: AudienceConfig = Field(default_factory=AudienceConfig)
+
+
+class BrowserClientBody(BaseModel):
+    client_id: str = Field(alias="clientId")
+    secret: str = Field(repr=False)
+    redirect_uris: tuple[str, ...] = Field(alias="redirectUris")
+    enabled: bool = True
+    public_client: bool = Field(default=False, alias="publicClient")
+    standard_flow_enabled: bool = Field(default=True, alias="standardFlowEnabled")
+    direct_access_grants_enabled: bool = Field(default=True, alias="directAccessGrantsEnabled")
+    default_client_scopes: tuple[str, ...] = Field(default=("email", "basic"), alias="defaultClientScopes")
+    attributes: BrowserAttributes = Field(default_factory=BrowserAttributes)
+    protocol_mappers: tuple[AudienceMapper, ...] = Field(default=(AudienceMapper(),), alias="protocolMappers")
+
+
+@dataclass(frozen=True, slots=True)
+class BrowserClient:
+    client_id: str
+    secret: str = field(repr=False)
+    callback_url: str
+
+    def environment(self, discovery: Discovery) -> dict[str, str]:
+        return {
+            "GENERIC_CLIENT_ID": self.client_id,
+            "GENERIC_CLIENT_SECRET": self.secret,
+            "GENERIC_USER_ID_ATTRIBUTE": "sub",
+            "GENERIC_AUTHORIZATION_ENDPOINT": discovery.authorization_endpoint,
+            "GENERIC_TOKEN_ENDPOINT": discovery.token_endpoint,
+            "GENERIC_USERINFO_ENDPOINT": discovery.userinfo_endpoint,
+            "GENERIC_CLIENT_USE_PKCE": "true",
+            "GENERIC_SCOPE": "openid email",
+        }
+
+
+def token_claims(token: str) -> TokenClaims:
+    payload: Final = token.split(".")[1]
+    return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)))
+
+
+def _signal_process_group(process_id: int, signum: int) -> bool:
+    try:
+        os.killpg(process_id, signum)
+    except ProcessLookupError:
+        return False
+    return True
+
+
+def _stop_process_group(child: subprocess.Popen[bytes]) -> None:
+    _signal_process_group(child.pid, signal.SIGTERM)
+    deadline: Final = time.monotonic() + 5
+    while _process_group_exists(child.pid):
+        child.poll()
+        if time.monotonic() >= deadline:
+            _signal_process_group(child.pid, signal.SIGKILL)
+            break
+        time.sleep(0.05)
+    child.wait()
+
+
+def _process_group_exists(process_id: int) -> bool:
+    try:
+        os.killpg(process_id, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
+    idp: Final = keycloak_from_env().with_strict_cleanup()
+    with ExitStack() as cleanup:
+
+        def terminate(signum: int, frame: FrameType | None) -> None:
+            raise SystemExit(128 + signum)
+
+        previous: Final = signal.signal(signal.SIGTERM, terminate)
+        cleanup.callback(signal.signal, signal.SIGTERM, previous)
+
+        def defer(callback: Callable[[], object]) -> None:
+            cleanup.callback(callback)
+
+        client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer)
+        environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url}
+        with subprocess.Popen(command, env=environment, start_new_session=True) as child:
+            try:
+                return child.wait()
+            finally:
+                _stop_process_group(child)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        raise SystemExit("Usage: idp.py PROXY_URL COMMAND [ARG ...]; requires a running test IdP")
+    raise SystemExit(run_oidc_profile(sys.argv[1], sys.argv[2:]))

tests/e2e/junit_properties.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 from collections.abc import Iterable
 
 import pytest
+from coverage_registry.management_cases import case_properties
 
 # Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
 # at runtime names this suite's place in the repo. test_junit_properties.py
@@ -94,7 +95,7 @@ def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]:
         ("package", package_from_nodeid(item.nodeid)),
         ("covers", ",".join(covers_from_item(item))),
         ("source", source_from_item(item)),
-    )
+    ) + case_properties(item.nodeid)
 
 
 def attach_result_properties(item: pytest.Item) -> None:

tests/e2e/management/conftest.py

@@ -5,8 +5,14 @@ holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams,
 users, and orgs this suite creates.
 """
 
-import pytest
+from collections.abc import Generator
+from typing import Final
 
+import pytest
+from e2e_http import without_retries
+from idp import Keycloak
+from lifecycle import ResourceManager
+from management.jwt_actors import ActorFactory
 from management_client import ManagementClient, build_client
 from proxy_client import ProxyClient
 
@@ -21,3 +27,14 @@ def pytest_configure(config: pytest.Config) -> None:
 @pytest.fixture(scope="session")
 def client(proxy: ProxyClient) -> ManagementClient:
     return build_client(proxy)
+
+
+@pytest.fixture
+def actor_factory(proxy: ProxyClient, idp: Keycloak) -> Generator[ActorFactory]:
+    bootstrap: Final = build_client(proxy)
+    resources: Final = ResourceManager(client=proxy, strict_cleanup=True)
+    with without_retries():
+        try:
+            yield ActorFactory(bootstrap=bootstrap, idp=idp, resources=resources)
+        finally:
+            resources.teardown()

tests/e2e/management/jwt_actors.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Final, Literal
+
+from e2e_config import unique_marker
+from e2e_http import NoBody, unwrap
+from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak
+from lifecycle import ResourceManager
+from management.management_client import ManagementClient
+from models import (
+    KeyGenerateBody,
+    KeyGenerateResponse,
+    OrgDeleteBody,
+    OrgDeleteResponse,
+    OrgMemberAddBody,
+    OrgMemberEntry,
+    OrgNewBody,
+    TeamDeleteBody,
+    TeamMemberAddBody,
+    TeamMemberEntry,
+    TeamNewBody,
+    UserNewBody,
+    UserRole,
+)
+from proxy_client import Caller
+
+ActorRole = Literal[
+    "proxy_admin",
+    "proxy_admin_viewer",
+    "organization_admin",
+    "team_admin",
+    "team_member",
+    "internal_user",
+    "internal_user_viewer",
+    "unrelated_user",
+]
+ActorProfile = Literal["database_role", "group_scoped"]
+
+
+@dataclass(frozen=True, slots=True)
+class Tenant:
+    organization_id: str
+    team_id: str
+    group_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class Actor:
+    identity: Identity
+    role: ActorRole
+    global_role: UserRole
+    profile: ActorProfile
+    tenants: tuple[Tenant, ...]
+
+    def mint_caller(self, idp: Keycloak) -> Caller:
+        return Caller(
+            credential=idp.access_token(
+                self.identity, client_id=ADMIN_CLIENT_ID if self.role == "proxy_admin" else TESTS_CLIENT_ID
+            ),
+            kind="direct_jwt",
+            role=self.role,
+            tenant=self.tenants[0].organization_id if self.tenants else None,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ActorFactory:
+    bootstrap: ManagementClient
+    idp: Keycloak
+    resources: ResourceManager
+
+    def __post_init__(self) -> None:
+        if self.bootstrap.proxy.caller is not None:
+            raise ValueError("Actor bootstrap requires a separately held master client")
+
+    def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse:
+        created: Final = unwrap(
+            self.bootstrap.generate_key(
+                KeyGenerateBody(
+                    team_id=tenant.team_id if tenant is not None else None,
+                    user_id=user_id,
+                    key_alias=f"e2e-actor-key-{unique_marker()}",
+                )
+            )
+        )
+        self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True))
+        return created
+
+    def tenant(self) -> Tenant:
+        marker: Final = unique_marker()
+        organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}"))
+        self.resources.defer(
+            lambda: unwrap(
+                self.bootstrap.proxy.transport.delete(
+                    "/organization/delete",
+                    headers=self.bootstrap.proxy.management_headers(),
+                    json=OrgDeleteBody(organization_ids=[organization_id]),
+                    response_type=OrgDeleteResponse,
+                )
+            )
+        )
+        team_id: Final = self.bootstrap.proxy.create_team(
+            TeamNewBody(team_alias=f"e2e-team-{marker}", organization_id=organization_id)
+        )
+        self.resources.defer(
+            lambda: unwrap(
+                self.bootstrap.proxy.transport.post(
+                    "/team/delete",
+                    headers=self.bootstrap.proxy.management_headers(),
+                    json=TeamDeleteBody(team_ids=[team_id]),
+                    response_type=NoBody,
+                )
+            )
+        )
+        self.bootstrap.delete_team_member(team_id, self.bootstrap.user_info().user_id)
+        group_id: Final = self.idp.create_group(team_id)
+        self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id))
+        return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id)
+
+    def create(
+        self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role"
+    ) -> Actor:
+        if role in ("team_admin", "team_member", "organization_admin") and not tenants:
+            raise ValueError("A membership actor requires a tenant")
+        identity: Final = self.idp.with_strict_cleanup().provision_user(
+            marker=unique_marker(),
+            groups=tuple(tenant.team_id for tenant in tenants) if profile == "group_scoped" else (),
+            group_ids=tuple(tenant.group_id for tenant in tenants) if profile == "group_scoped" else (),
+            defer=self.resources.defer,
+        )
+        global_role: Final[UserRole] = (
+            role
+            if role in ("proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer")
+            else "internal_user"
+        )
+        self.bootstrap.create_user(
+            UserNewBody(
+                user_id=identity.user_id,
+                user_email=f"{identity.username}@example.com",
+                user_role=global_role,
+                auto_create_key=False,
+            )
+        )
+        self.resources.defer(lambda: self.bootstrap.delete_user_strict(identity.user_id))
+        for tenant in tenants:
+            unwrap(
+                self.bootstrap.proxy.transport.post(
+                    "/organization/member_add",
+                    headers=self.bootstrap.proxy.management_headers(),
+                    json=OrgMemberAddBody(
+                        organization_id=tenant.organization_id,
+                        member=OrgMemberEntry(
+                            user_id=identity.user_id,
+                            role="org_admin" if role == "organization_admin" else "internal_user",
+                        ),
+                    ),
+                    response_type=NoBody,
+                )
+            )
+            unwrap(
+                self.bootstrap.proxy.transport.post(
+                    "/team/member_add",
+                    headers=self.bootstrap.proxy.management_headers(),
+                    json=TeamMemberAddBody(
+                        team_id=tenant.team_id,
+                        member=TeamMemberEntry(
+                            user_id=identity.user_id,
+                            role="admin" if role == "team_admin" else "user",
+                        ),
+                    ),
+                    response_type=NoBody,
+                )
+            )
+        return Actor(identity=identity, role=role, global_role=global_role, profile=profile, tenants=tenants)

tests/e2e/management/management_client.py

@@ -7,7 +7,8 @@ llm-only key hitting a management route).
 from __future__ import annotations
 
 import time
-from dataclasses import dataclass
+import warnings
+from dataclasses import dataclass, field, replace
 
 import jwt
 from e2e_config import MASTER_KEY
@@ -20,6 +21,7 @@ from e2e_http import (
     StreamingResponse,
     Success,
     UnknownApiError,
+    retry_attempts,
     unwrap,
 )
 from models import (
@@ -81,7 +83,7 @@ from models import (
     UserNewResponse,
     UserUpdateBody,
 )
-from proxy_client import ProxyClient
+from proxy_client import Caller, ProxyClient
 
 MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
 ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
@@ -98,7 +100,7 @@ class DashboardSession:
     its bearer on every subsequent call, the claims it renders the signed-in user
     from, and where it lands the browser."""
 
-    session_key: str
+    session_key: str = field(repr=False)
     claims: UiSessionClaims
     redirect_url: str
 
@@ -106,7 +108,10 @@ class DashboardSession:
 @dataclass(frozen=True, slots=True)
 class ManagementClient:
     proxy: ProxyClient
-    master_key: str
+    master_key: str = field(repr=False)
+
+    def with_caller(self, caller: Caller) -> ManagementClient:
+        return replace(self, proxy=self.proxy.with_caller(caller))
 
     def llm_only_key(self) -> str:
         return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
@@ -117,7 +122,7 @@ class ManagementClient:
         dashboard creates it under the session key their sign-in minted). Returns
         the outcome rather than unwrapping it, so a caller can poll a route that is
         only transiently refusing."""
-        headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
+        headers = self.proxy.management_headers(caller_key)
         return self.proxy.transport.post(
             "/key/generate",
             headers=headers,
@@ -131,9 +136,9 @@ class ManagementClient:
         sign-in minted, never the master key). Returns the outcome rather than
         unwrapping it, so a caller can poll a route that is only transiently
         refusing; `update_key_models` is the unwrapping shorthand."""
-        headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
+        headers = self.proxy.management_headers(caller_key)
         last: Result[NoBody] = NetworkError(message="/key/update was never attempted")
-        for attempt in range(_KEY_WRITE_ATTEMPTS):
+        for attempt in range(retry_attempts(_KEY_WRITE_ATTEMPTS)):
             last = self.proxy.transport.post(
                 "/key/update",
                 headers=headers,
@@ -144,6 +149,7 @@ class ManagementClient:
                 case UnknownApiError(body=error_body) if any(
                     marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS
                 ):
+                    warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
                     time.sleep(0.5 * (attempt + 1))
                     continue
                 case _:
@@ -153,25 +159,26 @@ class ManagementClient:
     def update_key_models(self, key: str, models: list[str]) -> None:
         _ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
 
-    def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]:
+    def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]:
         return self.proxy.transport.get(
             "/key/info",
-            headers=self.proxy.transport.bearer(caller_key),
+            headers=self.proxy.management_headers(caller_key),
             params=KeyInfoParams(key=key),
             response_type=KeyInfoResponse,
         )
 
-    def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None:
+    def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None:
         """Strict delete for the act phase of a test: a failed delete is a hard
         failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
-        _ = unwrap(
-            self.proxy.transport.post(
-                "/key/delete",
-                headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key),
-                json=KeyDeleteBody(keys=[key]),
-                response_type=NoBody,
-            )
+        result = self.proxy.transport.post(
+            "/key/delete",
+            headers=self.proxy.management_headers(caller_key),
+            json=KeyDeleteBody(keys=[key]),
+            response_type=NoBody,
         )
+        if missing_ok and isinstance(result, UnknownApiError) and result.status_code == 404:
+            return
+        _ = unwrap(result)
 
     def delete_model_strict(self, model_id: str) -> None:
         """Strict delete for the act phase of a test: a failed delete is a hard
@@ -179,7 +186,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/model/delete",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=ModelDeleteBody(id=model_id),
                 response_type=NoBody,
             )
@@ -190,7 +197,7 @@ class ManagementClient:
         Connection button, probing the live provider with the supplied params."""
         return self.proxy.transport.post(
             "/health/test_connection",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=body,
             response_type=ConnectionTestResponse,
             timeout=120.0,
@@ -200,7 +207,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/key/block",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=KeyBlockBody(key=key),
                 response_type=NoBody,
             )
@@ -209,7 +216,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.post(
                 "/key/regenerate",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=KeyRegenerateBody(key=key, grace_period=grace_period),
                 response_type=KeyGenerateResponse,
             )
@@ -219,7 +226,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.post(
                 f"/key/{key}/reset_spend",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=KeyResetSpendBody(reset_to=reset_to),
                 response_type=KeyResetSpendResponse,
             )
@@ -228,7 +235,7 @@ class ManagementClient:
     def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]:
         """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is
         who is asking: the master key by default, or a virtual key."""
-        headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
+        headers = self.proxy.management_headers(caller_key)
         return self.proxy.transport.get(
             "/key/list",
             headers=headers,
@@ -266,7 +273,7 @@ class ManagementClient:
         team_id = unwrap(
             self.proxy.transport.post(
                 "/team/new",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=TeamNewResponse,
             )
@@ -276,10 +283,10 @@ class ManagementClient:
 
     def update_team(self, body: TeamUpdateBody) -> None:
         last: Result[NoBody] | None = None
-        for attempt in range(5):
+        for attempt in range(retry_attempts(5)):
             last = self.proxy.transport.post(
                 "/team/update",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=NoBody,
             )
@@ -289,6 +296,7 @@ class ManagementClient:
                 case UnknownApiError(body=body_text) if (
                     "connecting to redis" in body_text.lower() or "name resolution" in body_text.lower()
                 ):
+                    warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
                     time.sleep(0.5 * (attempt + 1))
                     continue
                 case _:
@@ -299,7 +307,7 @@ class ManagementClient:
     def delete_team(self, team_id: str) -> None:
         _ = self.proxy.transport.post(
             "/team/delete",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=TeamDeleteBody(team_ids=[team_id]),
             response_type=NoBody,
         )
@@ -308,7 +316,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.get(
                 "/team/info",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=TeamInfoParams(team_id=team_id),
                 response_type=TeamInfoResponse,
             )
@@ -320,7 +328,7 @@ class ManagementClient:
             for entry in unwrap(
                 self.proxy.transport.get(
                     "/team/list",
-                    headers=self.proxy.transport.master,
+                    headers=self.proxy.management_headers(),
                     params=NoBody(),
                     response_type=TeamListResponse,
                 )
@@ -328,14 +336,16 @@ class ManagementClient:
         )
 
     def team_info_status(self, team_id: str) -> ProbeResult:
-        return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
+        return self.proxy.transport.probe(
+            "/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers()
+        )
 
     def _wait_for_team(self, team_id: str) -> None:
         last: Result[TeamInfoResponse] | None = None
-        for _ in range(_TEAM_READY_ATTEMPTS):
+        for _ in range(retry_attempts(_TEAM_READY_ATTEMPTS)):
             last = self.proxy.transport.get(
                 "/team/info",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=TeamInfoParams(team_id=team_id),
                 response_type=TeamInfoResponse,
             )
@@ -343,25 +353,29 @@ class ManagementClient:
                 case Success():
                     return
                 case _:
+                    warnings.warn("Repeating team read while the team becomes available", RuntimeWarning, stacklevel=2)
                     time.sleep(_TEAM_READY_SLEEP_SECONDS)
         assert last is not None
         raise AssertionError(last)
 
     def add_team_member(self, team_id: str, user_id: str) -> None:
         last: Result[NoBody] | None = None
-        for attempt in range(_TEAM_READY_ATTEMPTS):
+        for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)):
             last = self.proxy.transport.post(
                 "/team/member_add",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)),
                 response_type=NoBody,
             )
             match last:
                 case Success():
                     return
-                case UnknownApiError(body=body) if (
-                    "doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS
+                case UnknownApiError(body=body) if "doesn't exist" in body and attempt + 1 < retry_attempts(
+                    _TEAM_READY_ATTEMPTS
                 ):
+                    warnings.warn(
+                        "Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2
+                    )
                     time.sleep(_TEAM_READY_SLEEP_SECONDS)
                     continue
                 case _:
@@ -373,7 +387,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/team/member_delete",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id),
                 response_type=NoBody,
             )
@@ -383,7 +397,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.post(
                 "/user/new",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=UserNewResponse,
             )
@@ -393,7 +407,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/customer/new",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=CustomerNewBody(user_id=user_id),
                 response_type=CustomerResponse,
             )
@@ -404,7 +418,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.get(
                 "/customer/info",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=CustomerInfoParams(end_user_id=end_user_id),
                 response_type=CustomerResponse,
             )
@@ -413,7 +427,7 @@ class ManagementClient:
     def delete_customer(self, user_id: str) -> None:
         _ = self.proxy.transport.post(
             "/customer/delete",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=CustomerDeleteBody(user_ids=[user_id]),
             response_type=NoBody,
         )
@@ -422,7 +436,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/user/update",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=NoBody,
             )
@@ -431,7 +445,7 @@ class ManagementClient:
     def delete_user(self, user_id: str) -> None:
         _ = self.proxy.transport.post(
             "/user/delete",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=UserDeleteBody(user_ids=[user_id]),
             response_type=NoBody,
         )
@@ -442,17 +456,17 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.post(
                 "/user/delete",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=UserDeleteBody(user_ids=[user_id]),
                 response_type=UserDeleteResponse,
             )
         )
 
-    def user_info(self, user_id: str) -> UserInfoResponse:
+    def user_info(self, user_id: str | None = None) -> UserInfoResponse:
         return unwrap(
             self.proxy.transport.get(
                 "/user/info",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=UserInfoParams(user_id=user_id),
                 response_type=UserInfoResponse,
             )
@@ -462,7 +476,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.get(
                 "/user/list",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=UserListParams(user_ids=user_id),
                 response_type=UserListResponse,
             )
@@ -472,7 +486,7 @@ class ManagementClient:
         listing = unwrap(
             self.proxy.transport.get(
                 "/user/list",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=UserListParams(user_ids=user_id),
                 response_type=UserListResponse,
             )
@@ -483,7 +497,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.post(
                 "/organization/new",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=OrgNewResponse,
             )
@@ -493,7 +507,7 @@ class ManagementClient:
         _ = unwrap(
             self.proxy.transport.patch(
                 "/organization/update",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=NoBody,
             )
@@ -502,7 +516,7 @@ class ManagementClient:
     def delete_org(self, organization_id: str) -> None:
         _ = self.proxy.transport.delete(
             "/organization/delete",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=OrgDeleteBody(organization_ids=[organization_id]),
             response_type=NoBody,
         )
@@ -511,19 +525,24 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.get(
                 "/organization/info",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 params=OrgInfoParams(organization_id=organization_id),
                 response_type=OrgInfoResponse,
             )
         )
 
     def org_info_status(self, organization_id: str) -> ProbeResult:
-        return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id))
+        return self.proxy.transport.probe(
+            "/organization/info",
+            params=OrgInfoParams(organization_id=organization_id),
+            headers=self.proxy.management_headers(),
+        )
+
     def create_tag(self, body: TagNewBody) -> None:
         _ = unwrap(
             self.proxy.transport.post(
                 "/tag/new",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=NoBody,
             )
@@ -532,7 +551,7 @@ class ManagementClient:
     def delete_tag(self, name: str) -> None:
         _ = self.proxy.transport.post(
             "/tag/delete",
-            headers=self.proxy.transport.master,
+            headers=self.proxy.management_headers(),
             json=TagDeleteBody(name=name),
             response_type=NoBody,
         )
@@ -542,7 +561,7 @@ class ManagementClient:
             unwrap(
                 self.proxy.transport.get(
                     "/tag/list",
-                    headers=self.proxy.transport.master,
+                    headers=self.proxy.management_headers(),
                     params=NoBody(),
                     response_type=TagListResponse,
                 )
@@ -553,7 +572,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.post(
                 "/v1/mcp/server",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=McpServerRow,
             )
@@ -565,7 +584,7 @@ class ManagementClient:
         return unwrap(
             self.proxy.transport.put(
                 "/v1/mcp/server",
-                headers=self.proxy.transport.master,
+                headers=self.proxy.management_headers(),
                 json=body,
                 response_type=McpServerRow,
             )
@@ -576,7 +595,7 @@ class ManagementClient:
         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,
+            headers=self.proxy.management_headers(),
             json=NoBody(),
             response_type=NoBody,
         )

tests/e2e/management/test_jwt_management_e2e.py

@@ -2,60 +2,247 @@
 
 from __future__ import annotations
 
-from typing import Final
+from typing import Final, Literal
 
 import pytest
-from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
+from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker
 from e2e_http import UnauthorizedError, UnknownApiError, unwrap
-from idp import ADMIN_CLIENT_ID, Identity, Keycloak
+from idp import ADMIN_CLIENT_ID, Identity, Keycloak, token_claims
 from lifecycle import ResourceManager
+from management.jwt_actors import ActorFactory, ActorRole
 from management_client import ManagementClient
-from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody
+from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserInfoParams, UserInfoResponse, UserNewBody
+from proxy_client import Caller
 
 pytestmark = pytest.mark.e2e
 
 
 class TestJwtManagement:
-    @pytest.mark.covers("mgmt.key.jwt.lifecycle")
-    def test_admin_creates_reads_updates_clears_and_deletes_a_key(
-        self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
-    ) -> None:
-        admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
-        alias: Final = f"e2e-jwt-key-{unique_marker()}"
-        created: Final = unwrap(
-            client.generate_key(
-                KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]),
-                caller_key=admin,
+    @pytest.mark.parametrize(
+        "role",
+        (
+            "proxy_admin",
+            "proxy_admin_viewer",
+            "organization_admin",
+            "team_admin",
+            "team_member",
+            "internal_user",
+            "internal_user_viewer",
+            "unrelated_user",
+        ),
+    )
+    @pytest.mark.covers("mgmt.user.jwt.database_roles")
+    def test_actor_subject_and_database_role(self, actor_factory: ActorFactory, role: ActorRole) -> None:
+        tenants: Final = (
+            (actor_factory.tenant(),) if role in ("organization_admin", "team_admin", "team_member") else ()
+        )
+        actor: Final = actor_factory.create(role, tenants=tenants)
+        caller: Final = actor.mint_caller(actor_factory.idp)
+        claims: Final = token_claims(caller.credential)
+        assert claims.sub == actor.identity.user_id
+        assert claims.iss == actor_factory.idp.issuer
+        assert claims.aud == "litellm-e2e" or "litellm-e2e" in claims.aud
+        assert actor.identity.groups == ()
+        assert ("litellm_proxy_admin" in claims.scope.split()) == (role == "proxy_admin")
+        stored: Final = actor_factory.bootstrap.user_info(actor.identity.user_id)
+        assert stored.user_id == actor.identity.user_id
+        assert stored.user_info.user_role == actor.global_role
+        bound: Final = actor_factory.bootstrap.with_caller(caller)
+        own: Final = unwrap(
+            bound.proxy.transport.get(
+                "/user/info",
+                headers=bound.proxy.management_headers(),
+                params=UserInfoParams(),
+                response_type=UserInfoResponse,
             )
         )
-        resources.defer(lambda: client.proxy.delete_key(created.key))
+        assert own.user_id == actor.identity.user_id
+        assert own.user_info.user_role == actor.global_role
+        for tenant in tenants:
+            info = actor_factory.bootstrap.team_info(tenant.team_id)
+            assert info.organization_id == tenant.organization_id
+            assert {(member.user_id, member.role) for member in info.members_with_roles} == {
+                (actor.identity.user_id, "admin" if role == "team_admin" else "user")
+            }
+            assert {
+                (member.user_id, member.user_role)
+                for member in actor_factory.bootstrap.org_info(tenant.organization_id).members
+            } == {(actor.identity.user_id, "org_admin" if role == "organization_admin" else "internal_user")}
 
-        original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
-        assert original.key_alias == alias and original.team_id == jwt_identity.group
+    @pytest.mark.covers("mgmt.key.jwt.viewer_denied")
+    def test_admin_viewer_reads_but_cannot_update(self, actor_factory: ActorFactory) -> None:
+        actor: Final = actor_factory.create("proxy_admin_viewer")
+        viewer: Final = actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp))
+        alias: Final = f"e2e-viewer-{unique_marker()}"
+        key: Final = actor_factory.key().key
+        unwrap(actor_factory.bootstrap.update_key(KeyUpdateBody(key=key, key_alias=alias)))
+        assert viewer.proxy.key_info(key).key_alias == alias
+        denied: Final = viewer.update_key(KeyUpdateBody(key=key, key_alias="forbidden"))
+        assert isinstance(denied, UnknownApiError) and denied.status_code == 403, f"viewer write was accepted: {denied}"
+        assert "proxy_admin_viewer" in denied.body and "/key/update" in denied.body
+        assert actor_factory.bootstrap.proxy.key_info(key).key_alias == alias
+
+    @pytest.mark.covers("mgmt.user.oidc.identity_mapping")
+    def test_oidc_browser_profile_identity_mapping(self, actor_factory: ActorFactory) -> None:
+        actor: Final = actor_factory.create("internal_user")
+        idp: Final = actor_factory.idp.with_strict_cleanup()
+        discovery: Final = idp.discovery()
+        assert discovery.issuer == idp.issuer
+        assert discovery.jwks_uri == idp.jwks_url
+        callback: Final = f"{PROXY_BASE_URL}/sso/callback"
+        browser: Final = idp.browser_client(callback_url=callback, defer=actor_factory.resources.defer)
+        token: Final = idp.browser_token(actor.identity, browser)
+        assert token_claims(token).sub == actor.identity.user_id
+        userinfo: Final = idp.userinfo(token)
+        assert userinfo.sub == actor.identity.user_id
+        assert userinfo.email == f"{actor.identity.username}@example.com"
+        assert browser.environment(discovery)["GENERIC_USER_ID_ATTRIBUTE"] == "sub"
+
+    @pytest.mark.covers("mgmt.key.jwt.lifecycle")
+    @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key"))
+    def test_admin_creates_reads_updates_clears_and_deletes_a_key(
+        self,
+        actor_factory: ActorFactory,
+        credential_kind: Literal["direct_jwt", "virtual_key"],
+    ) -> None:
+        tenant: Final = actor_factory.tenant()
+        actor: Final = actor_factory.create("proxy_admin", tenants=(tenant,), profile="group_scoped")
+        virtual_key: Final = (
+            actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None
+        )
+        admin: Final = virtual_key if virtual_key is not None else actor.mint_caller(actor_factory.idp).credential
+        bound: Final = actor_factory.bootstrap.with_caller(
+            Caller(credential=admin, kind=credential_kind, role="proxy_admin")
+        )
+        assert bound.user_info().user_id == actor.identity.user_id
+        alias: Final = f"e2e-jwt-key-{unique_marker()}"
+        created: Final = unwrap(
+            bound.generate_key(
+                KeyGenerateBody(key_alias=alias, team_id=tenant.team_id, models=[CHEAP_OPENAI_MODEL]),
+            )
+        )
+        actor_factory.resources.defer(lambda: actor_factory.bootstrap.delete_key_strict(created.key, missing_ok=True))
+
+        original: Final = unwrap(bound.key_info_as(created.key)).info
+        assert original.key_alias == alias and original.team_id == tenant.team_id
         assert original.models == [CHEAP_OPENAI_MODEL]
 
         updated_alias: Final = f"{alias}-updated"
-        unwrap(
-            client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin)
-        )
-        updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
+        unwrap(bound.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120)))
+        updated: Final = unwrap(bound.key_info_as(created.key)).info
         assert updated.key_alias == updated_alias and updated.rpm_limit == 120
         assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction"
 
-        unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin))
-        cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
+        unwrap(bound.update_key(KeyUpdateBody(key=created.key, models=[])))
+        cleared: Final = unwrap(bound.key_info_as(created.key)).info
         assert cleared.models == [] and cleared.rpm_limit == 120
 
-        assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1
-        client.delete_key_strict(created.key, caller_key=admin)
-        assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0
+        assert unwrap(bound.key_list(updated_alias)).total_count == 1
+        bound.delete_key_strict(created.key)
+        assert unwrap(bound.key_list(updated_alias)).total_count == 0
+
+    @pytest.mark.covers("mgmt.team.jwt.tenant_isolation")
+    def test_two_actor_sets_keep_tenants_and_keys_isolated(self, actor_factory: ActorFactory) -> None:
+        first: Final = actor_factory.tenant()
+        second: Final = actor_factory.tenant()
+        assert first.organization_id != second.organization_id and first.team_id != second.team_id
+        actors: Final = tuple(
+            actor_factory.create("team_member", tenants=(tenant,), profile="group_scoped") for tenant in (first, second)
+        )
+        assert actors[0].identity.user_id != actors[1].identity.user_id
+        callers: Final = tuple(
+            actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) for actor in actors
+        )
+        keys: Final = tuple(actor_factory.key(tenant) for tenant in (first, second))
+        assert keys[0].key != keys[1].key
+        assert callers[0].proxy.key_info(keys[0].key).team_id == first.team_id
+        assert callers[1].proxy.key_info(keys[1].key).team_id == second.team_id
+        for caller, other_key in ((callers[0], keys[1].key), (callers[1], keys[0].key)):
+            hidden = caller.key_info_as(other_key)
+            assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403
+        assert tuple(actor.identity.groups for actor in actors) == ((first.team_id,), (second.team_id,))
+
+    @pytest.mark.covers("mgmt.team.jwt.multiple_memberships")
+    def test_multi_group_actor_keeps_exact_memberships(self, actor_factory: ActorFactory) -> None:
+        tenants: Final = (actor_factory.tenant(), actor_factory.tenant())
+        actor: Final = actor_factory.create("team_member", tenants=tenants, profile="group_scoped")
+        claims: Final = token_claims(actor.mint_caller(actor_factory.idp).credential)
+        assert set(claims.groups) == {tenant.team_id for tenant in tenants}
+        assert "litellm_proxy_admin" not in claims.scope.split()
+        assert actor.identity.groups == tuple(tenant.team_id for tenant in tenants)
+        for tenant in tenants:
+            assert {
+                (entry.user_id, entry.role)
+                for entry in actor_factory.bootstrap.team_info(tenant.team_id).members_with_roles
+            } == {(actor.identity.user_id, "user")}
+
+    @pytest.mark.covers("mgmt.user.jwt.cleanup")
+    def test_successful_actor_cleanup_removes_owned_state(self, actor_factory: ActorFactory) -> None:
+        resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True)
+        factory: Final = ActorFactory(bootstrap=actor_factory.bootstrap, idp=actor_factory.idp, resources=resources)
+        try:
+            tenant: Final = factory.tenant()
+            actor: Final = factory.create("team_member", tenants=(tenant,), profile="group_scoped")
+            key: Final = factory.key(tenant)
+            alias: Final = factory.bootstrap.proxy.key_info(key.key).key_alias
+            assert alias is not None
+        finally:
+            resources.teardown()
+        assert factory.bootstrap.user_count(actor.identity.user_id) == 0
+        assert factory.bootstrap.key_alias_count(alias) == 0
+        assert factory.bootstrap.team_info_status(tenant.team_id).status_code == 404
+        assert factory.bootstrap.org_info_status(tenant.organization_id).status_code == 404
+        factory.idp.assert_absent("users", actor.identity.user_id)
+        factory.idp.assert_absent("groups", tenant.group_id)
+
+    @pytest.mark.parametrize("stage", ("group", "user"))
+    @pytest.mark.covers("mgmt.user.jwt.partial_cleanup")
+    def test_partial_setup_removes_previously_created_identities(
+        self,
+        actor_factory: ActorFactory,
+        stage: Literal["group", "user"],
+    ) -> None:
+        idp: Final = actor_factory.idp.with_strict_cleanup()
+        resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True)
+        marker: Final = unique_marker()
+        group_id: Final = idp.create_group(f"e2e-partial-{marker}")
+        resources.defer(lambda: idp.delete_group(group_id))
+        try:
+            identity: Final = (
+                idp.provision_user(
+                    marker=marker,
+                    groups=(f"e2e-partial-{marker}",),
+                    group_ids=(group_id,),
+                    defer=resources.defer,
+                )
+                if stage == "user"
+                else None
+            )
+            if identity is None:
+                with pytest.raises(pytest.fail.Exception, match="HTTP 409"):
+                    idp.create_group(f"e2e-partial-{marker}")
+            else:
+                with pytest.raises(pytest.fail.Exception, match="HTTP 409"):
+                    idp.create_user(
+                        username=identity.username,
+                        email=f"{identity.username}@example.com",
+                        password=identity.password,
+                        groups=identity.groups,
+                    )
+        finally:
+            resources.teardown()
+        idp.assert_absent("groups", group_id)
+        if identity is not None:
+            idp.assert_absent("users", identity.user_id)
 
     @pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied")
     def test_member_cannot_write_and_another_team_cannot_read_the_key(
         self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
     ) -> None:
         admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
+        bound: Final = client.with_caller(Caller(credential=admin, kind="direct_jwt", role="proxy_admin"))
         member: Final = idp.access_token(jwt_identity)
+        member_client: Final = client.with_caller(Caller(credential=member, kind="direct_jwt", role="team_member"))
         alias: Final = f"e2e-jwt-owned-{unique_marker()}"
         created: Final = unwrap(
             client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin)
@@ -63,14 +250,14 @@ class TestJwtManagement:
         resources.defer(lambda: client.proxy.delete_key(created.key))
 
         client.add_team_member(jwt_identity.group, jwt_identity.user_id)
-        assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias
+        assert unwrap(member_client.key_info_as(created.key)).info.key_alias == alias
 
-        refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member)
+        refused: Final = member_client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"))
         assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}"
         assert "does not have permissions for endpoint" in refused.body.lower(), (
             f"expected a permission denial: {refused}"
         )
-        assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias
+        assert unwrap(bound.key_info_as(created.key)).info.key_alias == alias
 
         marker: Final = unique_marker()
         outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
@@ -88,4 +275,4 @@ class TestJwtManagement:
         assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, (
             f"another team must not read this key: {hidden}"
         )
-        assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group
+        assert unwrap(bound.key_info_as(created.key)).info.team_id == jwt_identity.group

tests/e2e/models.py

@@ -1091,13 +1091,13 @@ class UiLoginBody(BaseModel):
 
 
 class UiLoginResponse(BaseModel):
-    token: str
+    token: str = Field(repr=False)
     redirect_url: str
 
 
 class UiSessionClaims(BaseModel):
     user_id: str
-    key: str
+    key: str = Field(repr=False)
     user_role: str
     login_method: Literal["sso", "username_password"]
     exp: int
@@ -1135,6 +1135,7 @@ class TeamInfoParams(BaseModel):
 
 
 class TeamData(BaseModel):
+    organization_id: str | None = None
     team_alias: str | None = None
     models: list[str] = []
     members_with_roles: list[TeamMemberEntry] = []
@@ -1175,6 +1176,7 @@ class UserNewBody(BaseModel):
     user_email: str
     user_role: UserRole
     user_id: str | None = None
+    auto_create_key: bool | None = None
 
 
 class UserNewResponse(BaseModel):
@@ -1187,7 +1189,7 @@ class UserUpdateBody(BaseModel):
 
 
 class UserInfoParams(BaseModel):
-    user_id: str
+    user_id: str | None = None
 
 
 class UserData(BaseModel):
@@ -1240,16 +1242,36 @@ class OrgInfoParams(BaseModel):
     organization_id: str
 
 
+class OrgMembership(BaseModel):
+    user_id: str
+    user_role: str
+
+
 class OrgInfoResponse(BaseModel):
     organization_id: str
     organization_alias: str | None = None
     models: list[str] = []
+    members: tuple[OrgMembership, ...] = ()
+
+
+class OrgMemberEntry(BaseModel):
+    user_id: str
+    role: Literal["org_admin", "internal_user"]
+
+
+class OrgMemberAddBody(BaseModel):
+    organization_id: str
+    member: OrgMemberEntry
 
 
 class OrgDeleteBody(BaseModel):
     organization_ids: list[str]
 
 
+class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]):
+    pass
+
+
 # ---------- tags (management) ----------
 
 

tests/e2e/proxy_client.py

@@ -11,14 +11,23 @@ from __future__ import annotations
 import time
 import warnings
 from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import reduce
+from dataclasses import dataclass, field, replace
 from datetime import datetime
+from functools import reduce
 from types import MappingProxyType
-from typing import Final
-
-from pydantic import BaseModel
+from typing import Final, Literal
 
+from e2e_config import (
+    CONTROL_PLANE_BASE_URL,
+    MASTER_KEY,
+    POLL_INTERVAL,
+    POLL_TIMEOUT,
+    PROXY_BASE_URL,
+    PROXY_REPLICA_URLS,
+    REQUEST_TIMEOUT,
+    SLOW_PROVIDER_TIMEOUT_SECONDS,
+    settle_propagation,
+)
 from e2e_http import (
     AnthropicHeaders,
     AuthHeaders,
@@ -55,6 +64,7 @@ from models import (
     KeyInfoParams,
     KeyInfoResponse,
     LiteLLMParamsBody,
+    MemorySummaryResponse,
     ModelDeleteBody,
     ModelInfoBody,
     ModelInfoEntry,
@@ -63,7 +73,6 @@ from models import (
     ModelNewBody,
     ModelNewResponse,
     ModelsListParams,
-    MemorySummaryResponse,
     ModelsListResponse,
     ModelUpdateBody,
     OcrBody,
@@ -76,23 +85,13 @@ from models import (
     TeamDeleteBody,
     TeamNewBody,
     TeamNewResponse,
-    UserDeleteBody,
-    UserDeleteResponse,
     ToolsetCreateBody,
     ToolsetRow,
     ToolsetUpdateBody,
+    UserDeleteBody,
+    UserDeleteResponse,
 )
-from e2e_config import (
-    CONTROL_PLANE_BASE_URL,
-    MASTER_KEY,
-    POLL_INTERVAL,
-    POLL_TIMEOUT,
-    PROXY_BASE_URL,
-    PROXY_REPLICA_URLS,
-    REQUEST_TIMEOUT,
-    SLOW_PROVIDER_TIMEOUT_SECONDS,
-    settle_propagation,
-)
+from pydantic import BaseModel
 from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
 
 RowsPredicate = Callable[[list[SpendLogRow]], bool]
@@ -421,11 +420,23 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re
     )
 
 
+CredentialKind = Literal["master", "direct_jwt", "virtual_key", "dashboard_session"]
+
+
+@dataclass(frozen=True, slots=True)
+class Caller:
+    credential: str = field(repr=False)
+    kind: CredentialKind
+    role: str
+    tenant: str | None = None
+
+
 @dataclass(frozen=True, slots=True)
 class ProxyClient:
     transport: Transport
     replicas: Mapping[str, Transport]
     control_replicas: Mapping[str, Transport]
+    caller: Caller | None = None
     poll_timeout: float = 120.0
     poll_interval: float = 5.0
     model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT
@@ -433,13 +444,24 @@ class ProxyClient:
     model_servable_interval: float = MODEL_SERVABLE_INTERVAL
     model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
 
+    def with_caller(self, caller: Caller) -> ProxyClient:
+        return replace(self, caller=caller)
+
+    def management_headers(self, caller_key: str | None = None, *, transport: Transport | None = None) -> AuthHeaders:
+        selected: Final = self.transport if transport is None else transport
+        if caller_key is not None:
+            return selected.bearer(caller_key)
+        if self.caller is not None:
+            return selected.bearer(self.caller.credential)
+        return selected.master
+
     # ---- keys / customers (satisfies lifecycle.ResourceClient) ----------
 
     def generate_key(self, body: KeyGenerateBody) -> str:
         return unwrap(
             self.transport.post(
                 "/key/generate",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=KeyGenerateResponse,
             )
@@ -448,7 +470,7 @@ class ProxyClient:
     def delete_key(self, key: str) -> None:
         _ = self.transport.post(
             "/key/delete",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=KeyDeleteBody(keys=[key]),
             response_type=NoBody,
         )
@@ -458,7 +480,7 @@ class ProxyClient:
             return
         _ = self.transport.post(
             "/customer/delete",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=CustomerDeleteBody(user_ids=user_ids),
             response_type=NoBody,
         )
@@ -467,7 +489,7 @@ class ProxyClient:
         return unwrap(
             self.transport.get(
                 "/key/info",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 params=KeyInfoParams(key=key),
                 response_type=KeyInfoResponse,
             )
@@ -477,7 +499,7 @@ class ProxyClient:
         return {
             url: transport.get(
                 "/debug/memory/summary",
-                headers=transport.master,
+                headers=self.management_headers(transport=transport),
                 params=NoBody(),
                 response_type=MemorySummaryResponse,
             )
@@ -524,11 +546,12 @@ class ProxyClient:
             {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)}
         )
 
-    @staticmethod
     def _body_poller[R: BaseModel](
-        transport: Transport, path: str, params: BaseModel, response_type: type[R]
+        self, transport: Transport, path: str, params: BaseModel, response_type: type[R]
     ) -> Poller[Result[R]]:
-        return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type)
+        return lambda: transport.get(
+            path, headers=self.management_headers(transport=transport), params=params, response_type=response_type
+        )
 
     def model_info(self) -> list[ModelInfoEntry]:
         """Every configured deployment with the price the proxy resolved for it
@@ -536,7 +559,7 @@ class ProxyClient:
         return unwrap(
             self.transport.get(
                 "/model/info",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 params=NoBody(),
                 response_type=ModelInfoResponse,
             )
@@ -546,7 +569,7 @@ class ProxyClient:
         return unwrap(
             self.transport.get(
                 "/public/litellm_model_cost_map",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 params=NoBody(),
                 response_type=CostMap,
             )
@@ -607,7 +630,7 @@ class ProxyClient:
         model_id = unwrap(
             self.transport.post(
                 "/model/new",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=ModelNewResponse,
             )
@@ -623,7 +646,7 @@ class ProxyClient:
 
     def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
         """Block until every replica lists `model_name`, or fail at model_servable_timeout."""
-        headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for)
+        headers: Final = self.management_headers(listed_for)
         outcome: Final = await_servable_everywhere(
             {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
             model_name=model_name,
@@ -666,7 +689,7 @@ class ProxyClient:
         unwrap(
             self.transport.post(
                 "/model/update",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=ModelUpdateBody(
                     litellm_params=litellm_params,
                     model_info=ModelInfoBody(id=model_id),
@@ -678,7 +701,7 @@ class ProxyClient:
     def delete_model(self, model_id: str) -> None:
         result = self.transport.post(
             "/model/delete",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=ModelDeleteBody(id=model_id),
             response_type=NoBody,
         )
@@ -747,11 +770,10 @@ class ProxyClient:
                     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]]:
+    def _reader[R: BaseModel](self, transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]:
         return lambda request_timeout: transport.get(
             path,
-            headers=transport.master,
+            headers=self.management_headers(transport=transport),
             params=NoBody(),
             response_type=response_type,
             timeout=request_timeout,
@@ -763,7 +785,7 @@ class ProxyClient:
         return unwrap(
             self.transport.post(
                 "/v1/mcp/toolset",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=ToolsetRow,
             )
@@ -775,7 +797,7 @@ class ProxyClient:
         return unwrap(
             self.transport.put(
                 "/v1/mcp/toolset",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=ToolsetRow,
             )
@@ -786,7 +808,7 @@ class ProxyClient:
         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,
+            headers=self.management_headers(),
             json=NoBody(),
             response_type=NoBody,
         )
@@ -795,7 +817,7 @@ class ProxyClient:
         unwrap(
             self.transport.post(
                 "/credentials",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=CredentialCreateResponse,
             )
@@ -804,7 +826,7 @@ class ProxyClient:
     def delete_credential(self, credential_name: str) -> None:
         result = self.transport.delete(
             f"/credentials/{credential_name}",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=NoBody(),
             response_type=NoBody,
         )
@@ -815,7 +837,7 @@ class ProxyClient:
         return unwrap(
             self.transport.post(
                 "/team/new",
-                headers=self.transport.master,
+                headers=self.management_headers(),
                 json=body,
                 response_type=TeamNewResponse,
             )
@@ -824,7 +846,7 @@ class ProxyClient:
     def delete_team(self, team_id: str) -> None:
         result = self.transport.post(
             "/team/delete",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=TeamDeleteBody(team_ids=[team_id]),
             response_type=NoBody,
         )
@@ -836,7 +858,7 @@ class ProxyClient:
         a user the proxy only upserts after a successful auth."""
         result = self.transport.post(
             "/user/delete",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             json=UserDeleteBody(user_ids=[user_id]),
             response_type=UserDeleteResponse,
         )
@@ -909,7 +931,7 @@ class ProxyClient:
     def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
         result = self.transport.get(
             "/spend/logs",
-            headers=self.transport.master,
+            headers=self.management_headers(),
             params=params,
             response_type=SpendLogs,
         )
@@ -924,7 +946,7 @@ class ProxyClient:
             return unwrap(
                 self.transport.get(
                     "/spend/logs/v2",
-                    headers=self.transport.master,
+                    headers=self.management_headers(),
                     params=SpendLogsPageParams(
                         start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
                         end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
@@ -977,7 +999,7 @@ class ProxyClient:
     # ---- route probe ----------------------------------------------------
 
     def probe(self, path: str, *, params: NoBody) -> ProbeResult:
-        return self.transport.probe(path, params=params)
+        return self.transport.probe(path, params=params, headers=self.management_headers())
 
 
 def build_proxy_client(

tests/e2e/test_e2e_http.py

@@ -29,6 +29,7 @@ from e2e_http import (
     request_with_retry,
     streaming_outcome,
     wire_body,
+    without_retries,
 )
 from pydantic import BaseModel, TypeAdapter
 
@@ -56,6 +57,15 @@ def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]
 
 
 class TestTransientRetryPolicy:
+    def test_qualification_disables_retries_and_restores_the_default(self) -> None:
+        responses: Final = (FakeResponse(529), FakeResponse(200))
+        sleep: Final = SleepRecorder()
+        with without_retries():
+            assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[0]
+        assert sleep.delays == ()
+        assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[1]
+        assert sleep.delays == (0.5,)
+
     def test_transient_set_is_only_statuses_the_proxy_cannot_emit(self) -> None:
         assert TRANSIENT_STATUSES == frozenset({529})
         assert 429 not in TRANSIENT_STATUSES

tests/e2e/test_idp.py

@@ -4,12 +4,20 @@ these carry no `e2e` marker and run everywhere."""
 
 from __future__ import annotations
 
+import os
+import signal
+import socket
+import subprocess
+import sys
+import time
+from builtins import ExceptionGroup
 from collections.abc import Callable, Generator
 from contextlib import ExitStack, contextmanager
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
 from queue import SimpleQueue
 from threading import Thread
-from typing import Final
+from typing import Final, Literal
 
 import pytest
 from e2e_http import ExternalWrite
@@ -18,6 +26,8 @@ from idp import (
     KEYCLOAK_ADMIN_USER_ENV,
     KEYCLOAK_REALM_ENV,
     KEYCLOAK_URL_ENV,
+    BrowserClientBody,
+    Discovery,
     Keycloak,
     PasswordCredential,
     UserCreateBody,
@@ -60,24 +70,48 @@ def _idp_server(
 ) -> Generator[tuple[Keycloak, SimpleQueue[str]]]:
     """Exercise provisioning failures through the same HTTP transport as live tests."""
     deletions: SimpleQueue[str] = SimpleQueue()
+    clients: SimpleQueue[BrowserClientBody] = SimpleQueue()
 
     class Handler(BaseHTTPRequestHandler):
         def log_message(self, format: str, *args: object) -> None:
             pass
 
         def do_POST(self) -> None:
-            self.rfile.read(int(self.headers.get("Content-Length", "0")))
+            body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
             if self.path.endswith("/token"):
                 self.send_response(admin_status)
                 self.end_headers()
                 self.wfile.write(b'{"access_token":"synthetic-harness-token"}')
             else:
+                if self.path.endswith("/clients"):
+                    clients.put(BrowserClientBody.model_validate_json(body))
                 self.send_response(user_status if self.path.endswith("/users") else 201)
                 self.send_header("Location", f"{self.path}/resource-1")
                 self.end_headers()
                 if user_status != 201 and self.path.endswith("/users"):
                     self.wfile.write(b"injected create failure")
 
+        def do_GET(self) -> None:
+            self.send_response(200)
+            self.end_headers()
+            if "/clients/" in self.path:
+                client: Final = clients.get_nowait()
+                clients.put(client)
+                self.wfile.write(client.model_dump_json(by_alias=True).encode())
+            else:
+                issuer: Final = f"http://127.0.0.1:{server.server_port}/realms/test"
+                self.wfile.write(
+                    Discovery(
+                        issuer=issuer,
+                        authorization_endpoint=f"{issuer}/auth",
+                        token_endpoint=f"{issuer}/token",
+                        userinfo_endpoint=f"{issuer}/userinfo",
+                        jwks_uri=f"{issuer}/certs",
+                    )
+                    .model_dump_json()
+                    .encode()
+                )
+
         def do_DELETE(self) -> None:
             deletions.put(self.path)
             self.send_response(delete_status)
@@ -115,6 +149,68 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No
         assert deletions.empty()
 
 
+@pytest.mark.parametrize(
+    ("exit_mode", "ignore_termination"), (("normal", False), ("parent", False), ("group", False), ("parent", True))
+)
+def test_oidc_launcher_removes_client_on_exit_and_termination(
+    tmp_path: Path, exit_mode: Literal["normal", "parent", "group"], ignore_termination: bool
+) -> None:
+    ready: Final = tmp_path / "ready"
+    descendant_command: Final = (
+        "import signal,socket,time; from pathlib import Path; "
+        + ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore_termination else "")
+        + "listener=socket.socket(); listener.bind(('127.0.0.1',0)); listener.listen(); "
+        f"Path({str(ready)!r}).write_text(str(listener.getsockname()[1])); time.sleep(120)"
+    )
+    child_command: Final = (
+        "import os,subprocess,sys,time; from pathlib import Path; "
+        'assert os.environ["GENERIC_CLIENT_SECRET"]; '
+        'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; '
+        f"subprocess.Popen([sys.executable, '-c', {descendant_command!r}]); "
+        f"ready=Path({str(ready)!r})\n"
+        "while not ready.exists(): time.sleep(0.05)\n"
+        + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)")
+    )
+    with _idp_server() as (idp, deletions):
+        with subprocess.Popen(
+            [
+                sys.executable,
+                str(Path(__file__).with_name("idp.py")),
+                "http://127.0.0.1:9999",
+                sys.executable,
+                "-c",
+                child_command,
+            ],
+            env={
+                **os.environ,
+                KEYCLOAK_URL_ENV: idp.base_url,
+                KEYCLOAK_REALM_ENV: idp.realm,
+                KEYCLOAK_ADMIN_USER_ENV: idp.admin_username,
+                KEYCLOAK_ADMIN_PASSWORD_ENV: idp.admin_password,
+            },
+            start_new_session=True,
+        ) as process:
+            try:
+                deadline: Final = time.monotonic() + 15
+                while not ready.exists() and time.monotonic() < deadline and process.poll() is None:
+                    time.sleep(0.05)
+                assert ready.exists(), "OIDC child did not start"
+                if exit_mode == "parent":
+                    process.terminate()
+                elif exit_mode == "group":
+                    os.killpg(process.pid, signal.SIGTERM)
+                assert process.wait(timeout=15) == (7 if exit_mode == "normal" else 143)
+                with socket.socket() as connection:
+                    connection.settimeout(1)
+                    assert connection.connect_ex(("127.0.0.1", int(ready.read_text()))) != 0
+            finally:
+                if process.poll() is None:
+                    os.killpg(process.pid, signal.SIGKILL)
+                    process.wait(timeout=5)
+        assert deletions.get(timeout=5) == "/admin/realms/test/clients/resource-1"
+        assert deletions.empty()
+
+
 def test_successful_provisioning_cleans_up_user_before_group() -> None:
     with _idp_server() as (idp, deletions):
         with ExitStack() as cleanup:
@@ -134,6 +230,43 @@ def test_cleanup_failure_is_visible() -> None:
             idp.delete_group("group")
 
 
+def test_strict_cleanup_reports_each_failure_and_continues() -> None:
+    from lifecycle import ResourceManager
+    from proxy_client import build_proxy_client
+
+    with _idp_server(delete_status=500) as (idp, deletions):
+        resources: Final = ResourceManager(client=build_proxy_client(), strict_cleanup=True)
+        strict: Final = idp.with_strict_cleanup()
+        resources.defer(lambda: strict.delete_group("group"))
+        resources.defer(lambda: strict.delete_user("user"))
+        with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as error:
+            resources.teardown()
+        assert len(error.value.exceptions) == 2
+        assert deletions.get_nowait() == "/admin/realms/test/users/user"
+        assert deletions.get_nowait() == "/admin/realms/test/groups/group"
+
+
+@pytest.mark.parametrize("groups", ((), ("one",), ("one", "two")))
+def test_provisioning_records_zero_one_or_multiple_groups(groups: tuple[str, ...]) -> None:
+    with _idp_server() as (idp, deletions):
+        with ExitStack() as cleanup:
+
+            def defer(callback: Callable[[], object]) -> None:
+                cleanup.callback(callback)
+
+            identity: Final = idp.provision_groups(
+                marker="memberships",
+                groups=groups,
+                defer=defer,
+            )
+            assert identity.groups == groups
+            assert len(identity.group_ids) == len(groups)
+        assert deletions.get_nowait() == "/admin/realms/test/users/resource-1"
+        for _ in groups:
+            assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1"
+        assert deletions.empty()
+
+
 def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None:
     with _idp_server(admin_status=401) as (idp, _):
         cleanup: Final = ExitStack()

tests/e2e/test_proxy_client.py

@@ -11,19 +11,53 @@ injected clock, so nothing here monkeypatches anything.
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Mapping
+import json
+from builtins import ExceptionGroup
+from collections.abc import Callable, Generator, Iterable, Mapping
+from contextlib import contextmanager
 from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from itertools import chain, repeat
+from queue import SimpleQueue
+from threading import Thread
 from types import MappingProxyType
 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 e2e_http import NoBody, Result, Success, without_retries
+from idp import Keycloak
+from lifecycle import ResourceManager
+from management.jwt_actors import ActorFactory
+from management.management_client import ManagementClient
+from models import (
+    ConnectionTestBody,
+    CredentialCreateBody,
+    KeyGenerateBody,
+    KeyInfo,
+    KeyInfoResponse,
+    KeyUpdateBody,
+    LiteLLMParamsBody,
+    McpServerCreateBody,
+    McpServerUpdateBody,
+    ModelListEntry,
+    ModelsListResponse,
+    OrgNewBody,
+    OrgUpdateBody,
+    SpendLogsParams,
+    TagNewBody,
+    TeamNewBody,
+    TeamUpdateBody,
+    ToolsetCreateBody,
+    ToolsetUpdateBody,
+    UserNewBody,
+    UserUpdateBody,
+)
 from proxy_client import (
-    ConvergeOutcome,
+    Caller,
     Converged,
+    ConvergeOutcome,
+    CredentialKind,
     EverywhereConverged,
     ModelsPoller,
     NeverConvergedOn,
@@ -42,6 +76,115 @@ from proxy_client import (
 )
 from transport import Transport
 
+
+@contextmanager
+def caller_boundary(
+    status: int = 200, bodies: SimpleQueue[bytes] | None = None, *, delete_status: int | None = None
+) -> Generator[tuple[ManagementClient, SimpleQueue[str]]]:
+    received: Final[SimpleQueue[str]] = SimpleQueue()
+
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, format: str, *args: object) -> None:
+            pass
+
+        def do_GET(self) -> None:
+            received.put(self.headers.get("Authorization", ""))
+            self.send_response(delete_status if self.path == "/key/delete" and delete_status is not None else status)
+            self.end_headers()
+            self.wfile.write(
+                b'{"key":"owned","info":{"key_alias":"owned"},"data":[{"id":"owned"}],"team_id":"owned","team_info":{},"model_id":"owned"}'
+            )
+
+        def do_POST(self) -> None:
+            body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
+            if bodies is not None:
+                bodies.put(body)
+            self.do_GET()
+
+        do_PATCH = do_POST
+        do_PUT = do_POST
+        do_DELETE = do_POST
+
+    server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
+    thread.start()
+    url: Final = f"http://127.0.0.1:{server.server_port}"
+    proxy: Final = build_proxy_client(
+        base_url=url, control_plane_base_url=url, replica_urls=(url,), master_key="bootstrap"
+    )
+    try:
+        yield ManagementClient(proxy=proxy, master_key="bootstrap"), received
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(timeout=5)
+
+
+class TestBoundManagementCaller:
+    def test_strict_key_cleanup_accepts_missing_only_when_requested(self) -> None:
+        with caller_boundary(delete_status=404) as (bootstrap, received), without_retries():
+            with pytest.raises(AssertionError):
+                bootstrap.delete_key_strict("owned")
+            bootstrap.delete_key_strict("owned", missing_ok=True)
+            assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap")
+
+    def test_actor_key_cleanup_reports_failure_and_continues(self) -> None:
+        with caller_boundary(delete_status=500) as (bootstrap, received), without_retries():
+            resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True)
+            remaining: SimpleQueue[str] = SimpleQueue()
+            resources.defer(lambda: remaining.put("cleaned"))
+            factory: Final = ActorFactory(
+                bootstrap=bootstrap,
+                idp=Keycloak(base_url="http://unused.test", realm="test", admin_username="test", admin_password="test"),
+                resources=resources,
+            )
+            assert factory.key().key == "owned"
+            with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as failure:
+                resources.teardown()
+            assert len(failure.value.exceptions) == 1
+            assert remaining.get_nowait() == "cleaned"
+            assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap")
+
+    @pytest.mark.parametrize("kind", ("direct_jwt", "virtual_key", "dashboard_session"))
+    def test_direct_delegated_and_replica_reads_keep_the_bound_caller(self, kind: CredentialKind) -> None:
+        with caller_boundary() as (bootstrap, received):
+            caller: Final = Caller(credential="synthetic-caller", kind=kind, role="internal_user", tenant="tenant-a")
+            bound: Final = bootstrap.with_caller(caller)
+            bound.update_key(KeyUpdateBody(key="owned", key_alias="updated"))
+            bound.proxy.key_info("owned")
+            bound.proxy.read_back_everywhere(
+                "/key/info",
+                params=KeyUpdateBody(key="owned"),
+                response_type=KeyInfoResponse,
+                converged=lambda result: isinstance(result, Success),
+            )
+            bound.proxy.read_body_back_everywhere(
+                "/key/info", KeyInfoResponse, settled=lambda result: result.info.key_alias == "owned"
+            )
+            assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer synthetic-caller",) * 4
+            assert received.empty()
+            bootstrap.proxy.key_info("owned")
+            assert received.get_nowait() == "Bearer bootstrap"
+
+    def test_explicit_override_wins_without_rebinding_or_changing_master(self) -> None:
+        with caller_boundary() as (bootstrap, received):
+            bound: Final = bootstrap.with_caller(Caller(credential="bound", kind="direct_jwt", role="internal_user"))
+            bound.update_key(KeyUpdateBody(key="owned"), caller_key="override")
+            bound.proxy.key_info("owned")
+            assert received.get_nowait() == "Bearer override"
+            assert received.get_nowait() == "Bearer bound"
+            assert bound.master_key == "bootstrap"
+
+    def test_credentials_are_absent_from_binding_and_header_diagnostics(self) -> None:
+        with caller_boundary() as (bootstrap, _):
+            caller: Final = Caller(credential="private-value", kind="direct_jwt", role="internal_user")
+            bound: Final = bootstrap.with_caller(caller)
+            assert "private-value" not in repr(caller)
+            assert "private-value" not in repr(bound)
+            assert "private-value" not in repr(bound.proxy.management_headers())
+            assert "bootstrap" not in repr(bound)
+
+
 MODEL: Final = "gpt-under-test"
 _NO_TRANSPORTS: Final = cast(Transport, None)
 TIMEOUT: Final = 10.0
@@ -275,3 +418,166 @@ class TestReplicasFor:
         client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={})
         with pytest.raises(AssertionError, match="no replica is configured"):
             _ = client.replicas_for("/v1/models")
+
+
+MANAGEMENT_OPERATIONS: Final[tuple[tuple[str, Callable[[ManagementClient], object]], ...]] = (
+    ("generate_key", lambda c: c.generate_key(KeyGenerateBody())),
+    ("llm_only_key", lambda c: c.llm_only_key()),
+    ("update_key", lambda c: c.update_key(KeyUpdateBody(key="owned"))),
+    ("update_key_models", lambda c: c.update_key_models("owned", [])),
+    ("key_info", lambda c: c.key_info_as("owned")),
+    ("delete_key_strict", lambda c: c.delete_key_strict("owned")),
+    ("delete_model_strict", lambda c: c.delete_model_strict("owned")),
+    (
+        "connection_test",
+        lambda c: c.connection_test(
+            ConnectionTestBody(litellm_params=LiteLLMParamsBody(model="synthetic"), mode="chat")
+        ),
+    ),
+    ("block_key", lambda c: c.block_key("owned")),
+    ("regenerate_key", lambda c: c.regenerate_key("owned")),
+    ("reset_key_spend", lambda c: c.reset_key_spend("owned", 0)),
+    ("key_list", lambda c: c.key_list("owned")),
+    ("key_alias_count", lambda c: c.key_alias_count("owned")),
+    ("create_team", lambda c: c.create_team(TeamNewBody(team_alias="owned"))),
+    ("update_team", lambda c: c.update_team(TeamUpdateBody(team_id="owned", team_alias="updated"))),
+    ("delete_team", lambda c: c.delete_team("owned")),
+    ("team_info", lambda c: c.team_info("owned")),
+    ("team_list_ids", lambda c: c.team_list_ids()),
+    ("team_info_status", lambda c: c.team_info_status("owned")),
+    ("add_team_member", lambda c: c.add_team_member("owned", "user")),
+    ("delete_team_member", lambda c: c.delete_team_member("owned", "user")),
+    ("create_user", lambda c: c.create_user(UserNewBody(user_email="actor@example.com", user_role="internal_user"))),
+    ("create_customer", lambda c: c.create_customer("owned")),
+    ("customer_info", lambda c: c.customer_info("owned")),
+    ("delete_customer", lambda c: c.delete_customer("owned")),
+    ("update_user", lambda c: c.update_user(UserUpdateBody(user_id="owned", user_role="internal_user"))),
+    ("delete_user", lambda c: c.delete_user("owned")),
+    ("delete_user_strict", lambda c: c.delete_user_strict("owned")),
+    ("user_info", lambda c: c.user_info("owned")),
+    ("user_count", lambda c: c.user_count("owned")),
+    ("user_list_ids", lambda c: c.user_list_ids("owned")),
+    ("create_org", lambda c: c.create_org(OrgNewBody(organization_alias="owned"))),
+    ("update_org", lambda c: c.update_org(OrgUpdateBody(organization_id="owned", organization_alias="updated"))),
+    ("delete_org", lambda c: c.delete_org("owned")),
+    ("org_info", lambda c: c.org_info("owned")),
+    ("org_info_status", lambda c: c.org_info_status("owned")),
+    ("create_tag", lambda c: c.create_tag(TagNewBody(name="owned"))),
+    ("delete_tag", lambda c: c.delete_tag("owned")),
+    ("tag_list", lambda c: c.tag_list()),
+    ("create_mcp_server", lambda c: c.create_mcp_server(McpServerCreateBody(alias="owned", url="http://example.test"))),
+    ("update_mcp_server", lambda c: c.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))),
+    ("delete_mcp_server", lambda c: c.delete_mcp_server("owned")),
+    ("proxy.generate_key", lambda c: c.proxy.generate_key(KeyGenerateBody())),
+    ("proxy.delete_key", lambda c: c.proxy.delete_key("owned")),
+    ("proxy.delete_customers", lambda c: c.proxy.delete_customers(["owned"])),
+    ("proxy.key_info", lambda c: c.proxy.key_info("owned")),
+    ("proxy.memory_summary", lambda c: c.proxy.memory_summary_everywhere()),
+    ("proxy.model_info", lambda c: c.proxy.model_info()),
+    ("proxy.model_cost_map", lambda c: c.proxy.model_cost_map()),
+    ("proxy.create_model", lambda c: c.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))),
+    ("proxy.update_model", lambda c: c.proxy.update_model("owned", LiteLLMParamsBody(model="synthetic"))),
+    ("proxy.delete_model", lambda c: c.proxy.delete_model("owned")),
+    ("proxy.create_toolset", lambda c: c.proxy.create_toolset(ToolsetCreateBody(toolset_name="owned", tools=[]))),
+    ("proxy.update_toolset", lambda c: c.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))),
+    ("proxy.delete_toolset", lambda c: c.proxy.delete_toolset("owned")),
+    (
+        "proxy.create_credential",
+        lambda c: c.proxy.create_credential(CredentialCreateBody(credential_name="owned", credential_values={})),
+    ),
+    ("proxy.delete_credential", lambda c: c.proxy.delete_credential("owned")),
+    ("proxy.create_team", lambda c: c.proxy.create_team(TeamNewBody(team_alias="owned"))),
+    ("proxy.delete_team", lambda c: c.proxy.delete_team("owned")),
+    ("proxy.delete_user", lambda c: c.proxy.delete_user("owned")),
+    ("proxy.spend_logs", lambda c: c.proxy.spend_logs(SpendLogsParams(api_key="owned"))),
+    ("proxy.probe", lambda c: c.proxy.probe("/user/info", params=NoBody())),
+)
+
+
+@pytest.mark.parametrize(
+    ("name", "operation"), MANAGEMENT_OPERATIONS, ids=tuple(name for name, _ in MANAGEMENT_OPERATIONS)
+)
+@pytest.mark.parametrize("kind", ("master", "direct_jwt", "virtual_key", "dashboard_session"))
+def test_management_operations_send_the_selected_credential(
+    name: str,
+    operation: Callable[[ManagementClient], object],
+    kind: CredentialKind,
+) -> None:
+    with caller_boundary(status=401) as (bootstrap, received), without_retries():
+        client: Final = (
+            bootstrap
+            if kind == "master"
+            else bootstrap.with_caller(Caller(credential=f"synthetic-{kind}", kind=kind, role="internal_user"))
+        )
+        try:
+            operation(client)
+        except AssertionError:
+            pass
+        expected: Final = "Bearer bootstrap" if kind == "master" else f"Bearer synthetic-{kind}"
+        assert received.get_nowait() == expected, name
+        assert received.empty(), "an unauthorized request must not be retried"
+
+
+class TestSplitCallerPropagation:
+    def test_control_and_data_replica_readers_keep_the_caller(self) -> None:
+        with caller_boundary() as (data, data_headers), caller_boundary() as (control, control_headers):
+            data_url: Final = next(iter(data.proxy.replicas))
+            control_url: Final = next(iter(control.proxy.replicas))
+            proxy: Final = build_proxy_client(
+                base_url=data_url,
+                control_plane_base_url=control_url,
+                replica_urls=(data_url,),
+                master_key="bootstrap",
+            ).with_caller(Caller(credential="tenant-token", kind="direct_jwt", role="team_member"))
+            proxy.key_info("owned")
+            proxy.read_body_back_everywhere(
+                "/key/info", KeyInfoResponse, settled=lambda info: info.info.key_alias == "owned"
+            )
+            proxy.read_back_everywhere(
+                "/key/info",
+                params=NoBody(),
+                response_type=KeyInfoResponse,
+                converged=lambda result: isinstance(result, Success),
+            )
+            assert control_headers.get_nowait() == "Bearer tenant-token"
+            assert control_headers.get_nowait() == "Bearer tenant-token"
+            assert data_headers.get_nowait() == "Bearer tenant-token"
+            assert control_headers.empty() and data_headers.empty()
+
+    def test_successful_team_and_model_polling_uses_the_bound_caller(self) -> None:
+        with caller_boundary() as (bootstrap, received):
+            bound: Final = bootstrap.with_caller(Caller(credential="caller", kind="direct_jwt", role="proxy_admin"))
+            bound.create_team(TeamNewBody(team_alias="owned"))
+            bound.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))
+            assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer caller",) * 4
+            assert received.empty()
+
+    def test_expired_shaped_token_is_sent_once_without_renewal(self) -> None:
+        with caller_boundary(status=401) as (bootstrap, received):
+            bound: Final = bootstrap.with_caller(
+                Caller(credential="expired.payload.signature", kind="direct_jwt", role="internal_user")
+            )
+            result: Final = bound.key_info_as("owned")
+            assert not isinstance(result, Success)
+            assert received.get_nowait() == "Bearer expired.payload.signature"
+            assert received.empty()
+
+
+@pytest.mark.parametrize("operation", ("server", "toolset"))
+def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: str) -> None:
+    bodies: Final[SimpleQueue[bytes]] = SimpleQueue()
+    with caller_boundary(status=401, bodies=bodies) as (bootstrap, _):
+        try:
+            if operation == "server":
+                bootstrap.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))
+            else:
+                bootstrap.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))
+        except AssertionError:
+            pass
+        expected: Final = (
+            {"server_id": "owned", "alias": None}
+            if operation == "server"
+            else {"toolset_id": "owned", "description": None}
+        )
+        assert json.loads(bodies.get_nowait()) == expected
+        assert bodies.empty()

tests/e2e/transport.py

@@ -7,11 +7,9 @@ client touches requests.* or builds raw dicts; they pass pydantic models here.
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Protocol
 
-from pydantic import BaseModel
-
 import e2e_http
 from e2e_http import (
     URL,
@@ -21,6 +19,7 @@ from e2e_http import (
     Result,
     StreamingResponse,
 )
+from pydantic import BaseModel
 
 
 class Transport(Protocol):
@@ -85,7 +84,7 @@ class Transport(Protocol):
         self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
     ) -> Result[R]: ...
 
-    def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ...
+    def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: ...
 
     def upload[R: BaseModel](
         self,
@@ -113,7 +112,7 @@ class Transport(Protocol):
 @dataclass(frozen=True, slots=True)
 class HttpTransport:
     base_url: str
-    master_key: str
+    master_key: str = field(repr=False)
     request_timeout: float = 60.0
 
     def _url(self, path: str) -> URL:
@@ -245,10 +244,10 @@ class HttpTransport:
             timeout=self.request_timeout,
         )
 
-    def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
+    def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult:
         return e2e_http.probe(
             self._url(path),
-            headers=self.master,
+            headers=self.master if headers is None else headers,
             params=params,
             timeout=self.request_timeout,
         )
@@ -434,8 +433,8 @@ class SplitTransport:
             path, headers=headers, json=json, params=params, stream=stream
         )
 
-    def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
-        return self._route(path).probe(path, params=params)
+    def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult:
+        return self._route(path).probe(path, params=params, headers=headers)
 
     def upload[R: BaseModel](
         self,

tests/e2e/ui/oidcSetup.ts

@@ -0,0 +1,30 @@
+import { chromium, expect } from "@playwright/test";
+import * as fs from "fs";
+import * as path from "path";
+
+export default async function oidcSetup() {
+  const baseURL = process.env.E2E_OIDC_UI_URL;
+  const issuer = process.env.JWT_ISSUER;
+  const username = process.env.E2E_OIDC_USERNAME;
+  const password = process.env.E2E_OIDC_PASSWORD;
+  if (!baseURL || !issuer || !username || !password) {
+    throw new Error("The OIDC setup requires a running stack, issuer, and provisioned actor credentials");
+  }
+  const artifactDir = process.env.E2E_UI_ARTIFACT_DIR || ".";
+  fs.mkdirSync(artifactDir, { recursive: true });
+  const browser = await chromium.launch();
+  try {
+    const page = await browser.newPage();
+    await page.goto(`${baseURL.replace(/\/$/, "")}/sso/key/generate`);
+    await expect(page).toHaveURL(new RegExp(`^${issuer.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}/`));
+    await page.getByLabel("Username or email").fill(username);
+    await page.getByLabel("Password", { exact: true }).fill(password);
+    await page.getByRole("button", { name: "Sign In", exact: true }).click();
+    await page.waitForURL((url) => url.origin === new URL(baseURL).origin && url.pathname.startsWith("/ui"));
+    const statePath = path.join(artifactDir, "oidc.storageState.json");
+    await page.context().storageState({ path: statePath });
+    fs.chmodSync(statePath, 0o600);
+  } finally {
+    await browser.close();
+  }
+}

tests/e2e/ui/playwright.oidc.config.ts

@@ -0,0 +1,22 @@
+import { defineConfig, devices } from "@playwright/test";
+import * as path from "path";
+
+const baseURL = process.env.E2E_OIDC_UI_URL;
+if (!baseURL) throw new Error("E2E_OIDC_UI_URL must point to the running OIDC stack");
+
+export default defineConfig({
+  testDir: ".",
+  testMatch: "oidc/**/*.spec.ts",
+  retries: 0,
+  workers: 1,
+  outputDir: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc", "test-results"),
+  globalSetup: require.resolve("./oidcSetup"),
+  use: {
+    ...devices["Desktop Chrome"],
+    baseURL,
+    storageState: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc.storageState.json"),
+    trace: "off",
+    screenshot: "off",
+    video: "off",
+  },
+});

tests/llm_responses_api_testing/test_openai_responses_api.py

@@ -1627,9 +1627,10 @@ async def test_openai_responses_api_token_limit_error():
 
     Parsing the in-stream ErrorEvent must not raise
     "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent".
-    The iterator now surfaces the event as litellm.APIError with status 400
-    (invalid_request_error is a non-retriable client error, so no
-    MidStreamFallbackError wrapping) carrying the provider's message.
+    The iterator routes the event through litellm.exception_type, so it surfaces as
+    the typed 400 client error the non-streaming path raises (litellm.BadRequestError)
+    carrying the provider's message. invalid_request_error is a non-retriable client
+    error, so there is no MidStreamFallbackError wrapping.
     """
     litellm._turn_on_debug()
 
@@ -1644,7 +1645,7 @@ async def test_openai_responses_api_token_limit_error():
         async for event in response:
             print(event)
 
-    with pytest.raises(litellm.APIError) as exc_info:
+    with pytest.raises(litellm.BadRequestError) as exc_info:
         await _drain()
 
     assert exc_info.value.status_code == 400

tests/router_unit_tests/test_router_aresponses_streaming_fallback.py

@@ -372,7 +372,7 @@ async def test_aresponses_fallback_on_in_stream_error_event():
     raised = mock_fallback.await_args.kwargs["e"]
     assert isinstance(raised, MidStreamFallbackError)
     assert raised.status_code == 429
-    assert isinstance(raised.original_exception, litellm.APIError)
+    assert isinstance(raised.original_exception, litellm.RateLimitError)
     assert raised.original_exception.status_code == 429
     assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question"
 

tests/test_litellm/batches/test_batch_utils.py

@@ -695,6 +695,38 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch):
     assert result.failed_requests == 0
 
 
+def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch):
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "vertex_ai/gemini-embedding-2",
+        {
+            "input_cost_per_token_batches": 1e-7,
+            "input_cost_per_audio_token_batches": 3.25e-6,
+            "input_cost_per_image_token_batches": 2.25e-7,
+            "input_cost_per_video_token_batches": 6e-6,
+        },
+    )
+    responses = [
+        {
+            "response": {
+                "usageMetadata": {
+                    "promptTokenCount": 84,
+                    "candidatesTokenCount": 0,
+                    "totalTokenCount": 84,
+                    "promptTokensDetails": [
+                        {"modality": "AUDIO", "tokenCount": 64},
+                        {"modality": "TEXT", "tokenCount": 20},
+                    ],
+                }
+            }
+        }
+    ]
+
+    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2")
+
+    assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7)
+
+
 def test_vertex_cost_skips_none_response_body(monkeypatch):
     import litellm.cost_calculator as cc
 

tests/test_litellm/integrations/test_custom_guardrail.py

@@ -2625,7 +2625,7 @@ class TestLoggingOnlyApplyGuardrail:
         assert [e["guardrail_status"] for e in entries] == ["success"]
 
     @pytest.mark.asyncio
-    async def test_native_lifecycle_hook_guardrail_is_left_alone(self):
+    async def test_native_lifecycle_hook_guardrail_scans_in_logging_only(self):
         class _NativeHooks(_ApplyOnlyObserver):
             use_native_lifecycle_hooks = True
 
@@ -2634,9 +2634,9 @@ class TestLoggingOnlyApplyGuardrail:
 
         out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
 
-        assert guardrail.calls == []
-        assert out_kwargs is kwargs
+        assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
         assert out_response is response
+        assert out_kwargs["standard_logging_object"]["guardrail_information"]
 
     @pytest.mark.asyncio
     async def test_aresponses_scans_logged_messages_when_input_is_cleared(self):
@@ -2890,6 +2890,61 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
         assert len(_guardrail_entries(request_data)) == 1
 
 
+class _NativeLifecycleLoggingGuardrail(CustomGuardrail):
+    """Native lifecycle guardrail that also implements apply_guardrail, like the azure guards."""
+
+    use_native_lifecycle_hooks: ClassVar[bool] = True
+
+    def __init__(self):
+        from litellm.types.guardrails import GuardrailEventHooks
+
+        super().__init__(
+            guardrail_name="native-logging-guardrail",
+            event_hook=GuardrailEventHooks.logging_only,
+        )
+        self.calls: list[tuple[Literal["request", "response"], list[str]]] = []
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict[str, object],
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.calls.append((input_type, list(inputs.get("texts") or [])))
+        return inputs
+
+
+@pytest.mark.asyncio
+async def test_native_lifecycle_guardrail_logging_only_scans_assembled_response():
+    """A use_native_lifecycle_hooks guardrail accepts mode logging_only and its
+    async_logging_hook scans kwargs["async_complete_streaming_response"], not the raw result."""
+    from litellm.types.utils import Choices, Message, ModelResponse
+
+    guardrail = _NativeLifecycleLoggingGuardrail()
+    assembled = ModelResponse(
+        choices=[Choices(message=Message(role="assistant", content="assembled stream text"))]
+    )
+    sentinel_result = object()
+    kwargs = {
+        "model": "gpt-5.4-mini",
+        "messages": [{"role": "user", "content": "hi"}],
+        "litellm_call_id": "call-1",
+        "litellm_params": {"metadata": {}},
+        "optional_params": {},
+        "standard_logging_object": {"guardrail_information": None},
+        "async_complete_streaming_response": assembled,
+    }
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=sentinel_result, call_type=CallTypes.acompletion.value
+    )
+
+    assert out_result is sentinel_result
+    assert ("response", ["assembled stream text"]) in guardrail.calls
+    assert out_kwargs["standard_logging_object"]["guardrail_information"]
+
+
 class TestPreCallHookResponseIsNotLoggedVerbatim:
     """Regression for LIT-6935: a pre_call hook returning the request payload leaked the prompt
     into ``guardrail_response`` and from there onto OTEL guardrail spans."""

tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py

@@ -74,6 +74,108 @@ def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, s
     assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4])
 
 
+def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None:
+    model_info: ModelInfo = {
+        "key": "gemini-embedding-2",
+        "max_tokens": None,
+        "max_input_tokens": None,
+        "max_output_tokens": None,
+        "input_cost_per_token": 2e-7,
+        "input_cost_per_audio_token": 6.5e-6,
+        "input_cost_per_audio_per_second": 0.00016,
+        "output_cost_per_token": 0.0,
+        "litellm_provider": "vertex_ai",
+        "mode": "embedding",
+        "supported_openai_params": None,
+    }
+    usage = Usage(
+        prompt_tokens=64,
+        completion_tokens=0,
+        total_tokens=64,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            audio_tokens=64,
+            audio_length_seconds=2,
+        ),
+    )
+
+    prompt_cost, _ = generic_cost_per_token(
+        model="gemini-embedding-2",
+        usage=usage,
+        custom_llm_provider="vertex_ai",
+        model_info=model_info,
+    )
+
+    assert prompt_cost == pytest.approx(2 * 0.00016)
+
+
+def test_generic_cost_per_token_prefers_image_per_image_rate() -> None:
+    model_info: ModelInfo = {
+        "key": "gemini-embedding-2",
+        "max_tokens": None,
+        "max_input_tokens": None,
+        "max_output_tokens": None,
+        "input_cost_per_token": 2e-7,
+        "input_cost_per_image_token": 4.5e-7,
+        "input_cost_per_image": 0.00012,
+        "output_cost_per_token": 0.0,
+        "litellm_provider": "vertex_ai",
+        "mode": "embedding",
+        "supported_openai_params": None,
+    }
+    usage = Usage(
+        prompt_tokens=258,
+        completion_tokens=0,
+        total_tokens=258,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            image_tokens=258,
+            image_count=1,
+        ),
+    )
+
+    prompt_cost, _ = generic_cost_per_token(
+        model="gemini-embedding-2",
+        usage=usage,
+        custom_llm_provider="vertex_ai",
+        model_info=model_info,
+    )
+
+    assert prompt_cost == pytest.approx(0.00012)
+
+
+def test_generic_cost_per_token_prefers_video_per_second_rate() -> None:
+    model_info: ModelInfo = {
+        "key": "gemini-embedding-2",
+        "max_tokens": None,
+        "max_input_tokens": None,
+        "max_output_tokens": None,
+        "input_cost_per_token": 2e-7,
+        "input_cost_per_video_token": 1.2e-5,
+        "input_cost_per_video_per_second": 0.00079,
+        "output_cost_per_token": 0.0,
+        "litellm_provider": "vertex_ai",
+        "mode": "embedding",
+        "supported_openai_params": None,
+    }
+    usage = Usage(
+        prompt_tokens=516,
+        completion_tokens=0,
+        total_tokens=516,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            video_tokens=516,
+            video_length_seconds=2,
+        ),
+    )
+
+    prompt_cost, _ = generic_cost_per_token(
+        model="gemini-embedding-2",
+        usage=usage,
+        custom_llm_provider="vertex_ai",
+        model_info=model_info,
+    )
+
+    assert prompt_cost == pytest.approx(2 * 0.00079)
+
+
 def test_missing_cache_read_uses_off_peak_input_rate():
     from datetime import datetime, timezone
 

tests/test_litellm/llms/vertex_ai/batches/test_transformation.py

@@ -17,9 +17,9 @@ from unittest.mock import patch
 
 import pytest
 
-
 from litellm.llms.vertex_ai.batches.transformation import (  # noqa: E402
     VertexAIBatchTransformation,
+    vertex_prompt_tokens_details,
 )
 from litellm.llms.vertex_ai.common_utils import (  # noqa: E402
     VertexAIError,
@@ -41,6 +41,22 @@ ENDPOINT_INPUT_FILE = (
 )
 
 
+def test_vertex_prompt_tokens_details_rejects_malformed_details():
+    assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None
+    assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None
+    assert (
+        vertex_prompt_tokens_details(
+            {
+                "promptTokensDetails": [
+                    {"modality": "AUDIO", "tokenCount": 1},
+                    "malformed",
+                ]
+            }
+        )
+        is None
+    )
+
+
 # =========================================================================== #
 # transform_openai_batch_request_to_vertex_ai_batch_request
 # =========================================================================== #

tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py

@@ -10,6 +10,7 @@ Covers:
 
 import pytest
 
+import litellm
 from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
 from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
     _build_part_for_input,
@@ -22,11 +23,19 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation
 from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject
 from litellm.types.utils import EmbeddingResponse
 
-
 IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
 GCS_URL = "gs://my-bucket/image.png"
 
 
+@pytest.fixture(autouse=True)
+def _local_model_cost_map(monkeypatch):
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+    litellm.get_model_info.cache_clear()
+    yield
+    litellm.get_model_info.cache_clear()
+
+
 class TestIsMultimodalInput:
     def test_text_only_string(self):
         assert _is_multimodal_input("hello world") is False
@@ -324,7 +333,7 @@ class TestProcessEmbedContentResponseUsage:
         )
         assert result.usage.prompt_tokens == 258
         assert result.usage.total_tokens == 258
-        assert result.usage.prompt_tokens_details.image_count == 1
+        assert result.usage.prompt_tokens_details.image_tokens == 258
 
         prompt_cost, _ = generic_cost_per_token(
             model=self.MODEL,
@@ -358,7 +367,7 @@ class TestProcessEmbedContentResponseUsage:
         )
         assert prompt_cost > 0
 
-    def test_video_modality_derives_seconds_and_text_floor(self):
+    def test_video_modality_preserves_token_count(self):
         response_json = {
             "embedding": {"values": [0.1]},
             "usageMetadata": {
@@ -374,10 +383,8 @@ class TestProcessEmbedContentResponseUsage:
             response_json=response_json,
         )
         assert result.usage.prompt_tokens == 516
-        assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
-            2.0
-        )
-        assert result.usage.prompt_tokens_details.text_tokens == 1
+        assert result.usage.prompt_tokens_details.video_tokens == 516
+        assert result.usage.prompt_tokens_details.text_tokens == 0
 
     def test_missing_usage_metadata_does_not_estimate_from_base64(self):
         response_json = {"embedding": {"values": [0.1, 0.2]}}
@@ -400,8 +407,7 @@ class TestProcessEmbedContentResponseUsage:
         )
         assert result.usage.prompt_tokens > 0
 
-    def test_file_reference_image_billed_per_image_not_text(self):
-        """files/... image refs must bill per-image, not at the text token rate."""
+    def test_file_reference_image_billed_per_image_token_rate(self):
         response_json = {
             "embedding": {"values": [0.1, 0.2, 0.3]},
             "usageMetadata": {
@@ -422,7 +428,7 @@ class TestProcessEmbedContentResponseUsage:
                 }
             },
         )
-        assert result.usage.prompt_tokens_details.image_count == 1
+        assert result.usage.prompt_tokens_details.image_tokens == 258
         assert result.usage.prompt_tokens_details.text_tokens == 0
 
         prompt_cost, _ = generic_cost_per_token(
@@ -430,10 +436,10 @@ class TestProcessEmbedContentResponseUsage:
             usage=result.usage,
             custom_llm_provider="vertex_ai",
         )
-        assert prompt_cost == pytest.approx(0.00012)
+        assert prompt_cost == pytest.approx(258 * 4.5e-7)
 
     def test_file_reference_non_image_not_counted_as_image(self):
-        """A files/... ref resolving to a non-image mime must not be image-counted."""
+        """A files/... ref resolving to a non-image mime keeps audio token billing."""
         response_json = {
             "embedding": {"values": [0.1, 0.2]},
             "usageMetadata": {
@@ -454,21 +460,18 @@ class TestProcessEmbedContentResponseUsage:
                 }
             },
         )
-        assert result.usage.prompt_tokens_details.image_count == 0
         assert result.usage.prompt_tokens_details.audio_tokens == 64
-        assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
-            2.0
-        )
+        assert result.usage.prompt_tokens_details.image_tokens == 0
 
         prompt_cost, _ = generic_cost_per_token(
             model=self.MODEL,
             usage=result.usage,
             custom_llm_provider="vertex_ai",
         )
-        assert prompt_cost == pytest.approx(2.0 * 0.00016)
+        assert prompt_cost == pytest.approx(64 * 6.5e-6)
 
     def test_video_plus_audio_does_not_double_bill_text(self):
-        """Video+audio responses must not get video tokens reassigned to text."""
+        """Video and audio responses are billed from their respective token counts."""
         response_json = {
             "embedding": {"values": [0.1]},
             "usageMetadata": {
@@ -486,18 +489,145 @@ class TestProcessEmbedContentResponseUsage:
             model=self.MODEL,
             response_json=response_json,
         )
-        assert result.usage.prompt_tokens_details.text_tokens == 1
-        assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
-            2.0
-        )
-        assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
-            2.0
-        )
+        assert result.usage.prompt_tokens_details.text_tokens == 0
+        assert result.usage.prompt_tokens_details.video_tokens == 516
+        assert result.usage.prompt_tokens_details.audio_tokens == 64
 
         prompt_cost, _ = generic_cost_per_token(
             model=self.MODEL,
             usage=result.usage,
             custom_llm_provider="vertex_ai",
         )
-        # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4
-        assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016)
+        assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6)
+
+    def test_preview_alias_bills_audio_per_token(self):
+        response_json = {
+            "embedding": {"values": [0.1]},
+            "usageMetadata": {
+                "promptTokenCount": 64,
+                "totalTokenCount": 64,
+                "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}],
+            },
+        }
+        result = process_embed_content_response(
+            input="audio",
+            model_response=EmbeddingResponse(),
+            model="gemini-embedding-2-preview",
+            response_json=response_json,
+        )
+        prompt_cost, _ = generic_cost_per_token(
+            model="gemini-embedding-2-preview",
+            usage=result.usage,
+            custom_llm_provider="vertex_ai",
+        )
+        assert prompt_cost == pytest.approx(64 * 6.5e-6)
+
+    def test_image_without_modality_details_uses_image_rate(self):
+        response_json = {
+            "embedding": {"values": [0.1]},
+            "usageMetadata": {
+                "promptTokenCount": 258,
+                "totalTokenCount": 258,
+            },
+        }
+        result = process_embed_content_response(
+            input=IMAGE_DATA_URI,
+            model_response=EmbeddingResponse(),
+            model=self.MODEL,
+            response_json=response_json,
+        )
+        assert result.usage.prompt_tokens_details.image_tokens == 258
+        assert result.usage.prompt_tokens_details.text_tokens == 0
+
+        prompt_cost, _ = generic_cost_per_token(
+            model=self.MODEL,
+            usage=result.usage,
+            custom_llm_provider="vertex_ai",
+        )
+        assert prompt_cost == pytest.approx(258 * 4.5e-7)
+
+    @pytest.mark.parametrize(
+        "input_value,resolved_files,expected_image_tokens",
+        [
+            (GCS_URL, {}, 258),
+            ("gs://my-bucket/clip.mp4", {}, 0),
+            ("gs://my-bucket/unknown.bin", {}, 0),
+            ("files/image-123", {"files/image-123": {"mime_type": "image/jpeg"}}, 258),
+            ("files/missing", {}, 0),
+            ("data:application/octet-stream;base64,abc", {}, 0),
+            ([[IMAGE_DATA_URI]], {}, 258),
+            ([], {}, 0),
+        ],
+    )
+    def test_missing_modality_details_classifies_image_inputs(self, input_value, resolved_files, expected_image_tokens):
+        response_json = {
+            "embedding": {"values": [0.1]},
+            "usageMetadata": {
+                "promptTokenCount": 258,
+                "totalTokenCount": 258,
+            },
+        }
+        result = process_embed_content_response(
+            input=input_value,
+            model_response=EmbeddingResponse(),
+            model=self.MODEL,
+            response_json=response_json,
+            resolved_files=resolved_files,
+        )
+        assert result.usage.prompt_tokens_details.image_tokens == expected_image_tokens
+        assert result.usage.prompt_tokens_details.text_tokens == 0
+
+        prompt_cost, _ = generic_cost_per_token(
+            model=self.MODEL,
+            usage=result.usage,
+            custom_llm_provider="vertex_ai",
+        )
+        expected_rate = 4.5e-7 if expected_image_tokens else 2e-7
+        assert prompt_cost == pytest.approx(258 * expected_rate)
+
+    def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self):
+        response_json = {
+            "embedding": {"values": [0.1]},
+            "usageMetadata": {
+                "promptTokenCount": 270,
+                "totalTokenCount": 270,
+            },
+        }
+        result = process_embed_content_response(
+            input=["a short caption", IMAGE_DATA_URI],
+            model_response=EmbeddingResponse(),
+            model=self.MODEL,
+            response_json=response_json,
+        )
+        assert result.usage.prompt_tokens_details.image_tokens == 0
+
+        prompt_cost, _ = generic_cost_per_token(
+            model=self.MODEL,
+            usage=result.usage,
+            custom_llm_provider="vertex_ai",
+        )
+        assert prompt_cost == pytest.approx(270 * 2e-7)
+
+    def test_text_without_modality_details_uses_text_rate(self):
+        response_json = {
+            "embedding": {"values": [0.1]},
+            "usageMetadata": {
+                "promptTokenCount": 12,
+                "totalTokenCount": 12,
+            },
+        }
+        result = process_embed_content_response(
+            input="a short caption",
+            model_response=EmbeddingResponse(),
+            model=self.MODEL,
+            response_json=response_json,
+        )
+        assert result.usage.prompt_tokens_details.text_tokens == 0
+        assert result.usage.prompt_tokens_details.image_tokens == 0
+
+        prompt_cost, _ = generic_cost_per_token(
+            model=self.MODEL,
+            usage=result.usage,
+            custom_llm_provider="vertex_ai",
+        )
+        assert prompt_cost == pytest.approx(12 * 2e-7)

tests/test_litellm/proxy/client/cli/test_pi.py

@@ -20,6 +20,13 @@ from litellm.proxy.client.cli.commands.pi import (
 )
 
 
+def test_listing_failure_is_str_enum():
+    assert issubclass(ListingFailure, str)
+    assert ListingFailure.REJECTED.value == "rejected"
+    assert ListingFailure("rejected") is ListingFailure.REJECTED
+    assert str(ListingFailure.REJECTED.value) == "rejected"
+
+
 class _FakeResponse:
     def __init__(self, status_code, payload=None):
         self.status_code = status_code

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -2375,8 +2375,12 @@ _GROUNDING_QUERY_TEXT = "What is the capital of Japan?"
 _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo."
 
 
-def _grounding_guardrail() -> BedrockGuardrail:
-    return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT")
+def _grounding_guardrail(from_messages: bool = False) -> BedrockGuardrail:
+    return BedrockGuardrail(
+        guardrailIdentifier="test-guardrail",
+        guardrailVersion="DRAFT",
+        contextual_grounding_from_messages=from_messages,
+    )
 
 
 def _grounding_messages() -> list:
@@ -2418,9 +2422,11 @@ def _input_request(messages: list) -> dict:
     return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
 
 
-def _output_request(messages: list, response=None) -> dict:
+def _output_request(messages: list, response=None, from_messages: bool = False) -> dict:
     """Arrange a guardrail and act: build the Bedrock OUTPUT payload."""
-    return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
+    return _grounding_guardrail(from_messages).convert_to_bedrock_format(
+        source="OUTPUT", response=response, messages=messages
+    )
 
 
 def test_grounding_input_strips_grounding_and_query_qualifiers():
@@ -2474,6 +2480,131 @@ def test_grounding_output_keeps_legacy_payload_without_tags():
     assert actual_request == expected_request
 
 
+def test_grounding_output_derives_source_and_query_from_plain_messages():
+    """Flag on: untagged system + user text is sent as grounding_source + query."""
+    messages = [
+        {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+        {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+    ]
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],
+    }
+
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
+
+    assert actual_request == expected_request
+
+
+def test_grounding_output_plain_messages_stay_legacy_when_flag_is_off():
+    """Default config: plain system + user text is never sent as grounding context."""
+    messages = [
+        {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+        {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+    ]
+    expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]}
+
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT))
+
+    assert actual_request == expected_request
+
+
+def test_grounding_output_derived_query_is_latest_user_turn_only():
+    """Only the latest user turn is the query; system and developer turns are the source."""
+    developer_text = "Answer in one sentence."
+    messages = [
+        {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+        {"role": "developer", "content": [{"type": "text", "text": developer_text}]},
+        {"role": "user", "content": "Hi"},
+        {"role": "assistant", "content": "Hello, how can I help?"},
+        {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+    ]
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [
+            _GROUNDING_SOURCE_BLOCK,
+            {"text": {"text": developer_text, "qualifiers": ["grounding_source"]}},
+            _QUERY_BLOCK,
+            _GUARD_BLOCK,
+        ],
+    }
+
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
+
+    assert actual_request == expected_request
+
+
+@pytest.mark.parametrize(
+    "messages",
+    [
+        pytest.param([{"role": "system", "content": _GROUNDING_SOURCE_TEXT}], id="system-without-user"),
+        pytest.param(
+            [
+                {"role": "tool", "content": _GROUNDING_SOURCE_TEXT, "tool_call_id": "c1"},
+                {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+            ],
+            id="tool-result-is-not-a-source",
+        ),
+        pytest.param(
+            [
+                {"role": "system", "content": ""},
+                {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+            ],
+            id="empty-system-prompt",
+        ),
+        pytest.param(
+            [
+                {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+                {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/a.png"}}]},
+            ],
+            id="image-only-user-turn",
+        ),
+    ],
+)
+def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(messages):
+    """Bedrock rejects a source without a query and vice versa, so send neither."""
+    expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]}
+
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
+
+    assert actual_request == expected_request
+
+
+def test_grounding_output_explicit_tags_take_precedence_over_plain_messages():
+    """Tagged blocks win: untagged text around them is not added as source or query."""
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        *_grounding_messages(),
+        {"role": "user", "content": "Please be brief."},
+    ]
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],
+    }
+
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
+
+    assert actual_request == expected_request
+
+
+def test_grounding_input_ignores_plain_message_derivation():
+    """INPUT scans never derive grounding qualifiers from plain messages."""
+    messages = [
+        {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+        {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+    ]
+    expected_request = {
+        "source": "INPUT",
+        "content": [{"text": {"text": _GROUNDING_SOURCE_TEXT}}, {"text": {"text": _GROUNDING_QUERY_TEXT}}],
+    }
+
+    actual_request = _grounding_guardrail(from_messages=True).convert_to_bedrock_format(
+        source="INPUT", messages=messages
+    )
+
+    assert actual_request == expected_request
+
+
 def test_grounding_output_combines_multiple_sources():
     """Every grounding_source block is emitted; Bedrock combines them into one corpus."""
     uk_source_text = "London is the capital of UK."
@@ -2597,6 +2728,56 @@ async def test_grounding_output_blocked_raises_400():
     assert exc_info.value.status_code == 400
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "from_messages, request_messages",
+    [
+        (
+            True,
+            [
+                {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+                {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+            ],
+        ),
+        (
+            False,
+            [
+                {"role": "system", "content": [{"type": "grounding_source", "text": _GROUNDING_SOURCE_TEXT}]},
+                {"role": "user", "content": [{"type": "query", "text": _GROUNDING_QUERY_TEXT}]},
+            ],
+        ),
+    ],
+    ids=["plain-messages-flag-on", "tagged-messages-flag-off"],
+)
+async def test_apply_guardrail_response_forwards_request_messages_for_grounding(from_messages, request_messages):
+    guardrail = _grounding_guardrail(from_messages=from_messages)
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],
+    }
+
+    mock_credentials = MagicMock()
+    mock_credentials.access_key = "test-access-key"
+    mock_credentials.secret_key = "test-secret-key"
+    mock_credentials.token = None
+
+    with (
+        patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post,
+        patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")),
+        patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare,
+    ):
+        mock_post.return_value = _passing_bedrock_httpx_response(_GROUNDING_RESPONSE_TEXT)
+
+        await guardrail.apply_guardrail(
+            inputs={"texts": [_GROUNDING_RESPONSE_TEXT]},
+            request_data={"messages": request_messages},
+            input_type="response",
+        )
+
+    assert mock_prepare.call_count == 1
+    assert json.loads(json.dumps(mock_prepare.call_args.kwargs["data"])) == expected_request
+
+
 ###############################################################################
 # LIT-4186: disable_exception_on_block regression tests
 #

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py

@@ -2,6 +2,8 @@ import asyncio
 import base64
 import io
 import json
+from collections.abc import Iterator, Sequence
+from typing import cast
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
@@ -14,7 +16,7 @@ import litellm
 import litellm.types.utils
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache
-from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError
 from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail
@@ -4929,3 +4931,390 @@ def test_every_responses_delta_event_is_in_the_scanned_set():
     }
     assert not missing
     assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES
+
+
+def _clean_armor_response() -> dict[str, object]:
+    return {
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND",
+            "filterResults": {},
+        }
+    }
+
+
+def _flagged_armor_response() -> dict[str, object]:
+    return {
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}},
+        }
+    }
+
+
+class _FakeArmorHandler(AsyncHTTPHandler):
+    def __init__(self, responses: Sequence[dict[str, object] | Exception]):
+        self.responses: Iterator[dict[str, object] | Exception] = iter(responses)
+        self.calls: list[dict[str, object]] = []
+        self.raise_on_call: Exception | None = None
+
+    async def post(
+        self,
+        url: str,
+        json: dict[str, object] | None = None,
+        headers: dict[str, str] | None = None,
+        **kwargs: object,
+    ) -> httpx.Response:
+        if self.raise_on_call is not None:
+            raise self.raise_on_call
+        if json is not None:
+            self.calls.append(json)
+        response: dict[str, object] | Exception = next(self.responses)
+        if isinstance(response, Exception):
+            raise response
+        return httpx.Response(200, json=response, request=httpx.Request("POST", url))
+
+
+async def _async_token_provider() -> tuple[str, str]:
+    return ("test-token", "test-project")
+
+
+def _logging_only_guardrail(
+    responses: Sequence[dict[str, object] | Exception] = (_clean_armor_response(), _clean_armor_response()),
+) -> ModelArmorGuardrail:
+    handler = _FakeArmorHandler(responses)
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-logging",
+        event_hook=GuardrailEventHooks.logging_only,
+        async_handler=handler,
+        access_token_provider=_async_token_provider,
+    )
+    return guardrail
+
+
+def _logged_kwargs() -> dict[str, object]:
+    return {
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "hi"}],
+        "litellm_call_id": "call-1",
+        "litellm_params": {"metadata": {}},
+        "optional_params": {},
+        "standard_logging_object": {"guardrail_information": None},
+    }
+
+
+def _chat_response(text: str) -> litellm.ModelResponse:
+    return litellm.ModelResponse(
+        choices=[
+            litellm.types.utils.Choices(
+                message=litellm.types.utils.Message(role="assistant", content=text)
+            )
+        ]
+    )
+
+
+def _stream_chunk(text: str) -> litellm.ModelResponseStream:
+    return litellm.ModelResponseStream(
+        choices=[
+            litellm.types.utils.StreamingChoices(
+                delta=litellm.types.utils.Delta(content=text)
+            )
+        ]
+    )
+
+
+def _metadata_entries(kwargs: dict[str, object]) -> list[dict[str, object]]:
+    standard_logging_object = cast(dict[str, object], kwargs["standard_logging_object"])
+    entries = standard_logging_object.get("guardrail_information") or []
+    return cast(list[dict[str, object]], entries)
+
+
+def test_logging_only_mode_is_accepted_and_keeps_native_hooks():
+    guardrail = _logging_only_guardrail()
+    assert guardrail.event_hook == GuardrailEventHooks.logging_only
+    assert guardrail.use_native_lifecycle_hooks is True
+    assert GuardrailEventHooks.logging_only in ModelArmorGuardrail.get_supported_event_hooks()
+
+    post_call_guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-post",
+        event_hook=GuardrailEventHooks.post_call,
+    )
+    assert post_call_guardrail._deployment_hook_target() is post_call_guardrail
+
+
+@pytest.mark.asyncio
+async def test_logging_only_stream_yields_chunks_without_waiting_for_scan():
+    """A logging_only guardrail must pass stream chunks straight through; the scan happens
+    afterwards on the assembled response via async_logging_hook."""
+    guardrail = _logging_only_guardrail(
+        [_clean_armor_response(), _clean_armor_response()]
+    )
+    handler = cast(_FakeArmorHandler, guardrail.async_handler)
+    handler.raise_on_call = AssertionError("logging_only must not scan the stream")
+
+    produced = 0
+
+    async def gen():
+        nonlocal produced
+        for i in range(3):
+            produced += 1
+            yield _stream_chunk(f"chunk-{i} ")
+
+    hook_iter = guardrail.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        response=gen(),
+        request_data={"metadata": {}, "guardrails": ["model-armor-logging"]},
+    )
+    first = await hook_iter.__anext__()
+    assert produced == 1
+    chunks = [first]
+    async for chunk in hook_iter:
+        chunks.append(chunk)
+    assert len(chunks) == 3
+    assert handler.calls == []
+    handler.raise_on_call = None
+
+    response = _chat_response("all clear")
+    kwargs = _logged_kwargs()
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert len(entries) >= 1
+    entry = entries[-1]
+    assert entry["guardrail_status"] == "success"
+    assert entry["guardrail_mode"] == "logging_only"
+    assert entry["guardrail_provider"] == "model_armor"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_flagged_verdict_without_altering_response():
+    guardrail = _logging_only_guardrail(
+        [_flagged_armor_response(), _flagged_armor_response()]
+    )
+    response = _chat_response("flagged output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert entries[-1]["guardrail_status"] == "guardrail_flagged"
+    assert entries[-1]["guardrail_mode"] == "logging_only"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_model_armor_api_error():
+    guardrail = _logging_only_guardrail(
+        [
+            ModelArmorAPIError("Model Armor API error (upstream 500)"),
+            ModelArmorAPIError("Model Armor API error (upstream 500)"),
+        ]
+    )
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert entries[-1]["guardrail_status"] == "guardrail_failed_to_respond"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scans_assembled_responses_api_stream():
+    """The terminal ResponseCompletedEvent is an envelope; the scan must run on the
+    assembled ResponsesAPIResponse kept in kwargs."""
+    from openai.types.responses import ResponseOutputMessage, ResponseOutputText
+
+    from litellm.types.llms.openai import (
+        ResponseCompletedEvent,
+        ResponsesAPIResponse,
+        ResponsesAPIStreamEvents,
+    )
+
+    assembled = ResponsesAPIResponse(
+        id="resp-1",
+        created_at=1700000000,
+        output=[
+            ResponseOutputMessage(
+                id="msg-1",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        annotations=[], text="assembled output text", type="output_text"
+                    )
+                ],
+            )
+        ],
+    )
+    event = ResponseCompletedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=assembled
+    )
+
+    guardrail = _logging_only_guardrail()
+    kwargs = _logged_kwargs()
+    del kwargs["messages"]
+    kwargs["input"] = "hello"
+    kwargs["async_complete_streaming_response"] = assembled
+
+    out_kwargs, _ = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=event, call_type="aresponses"
+    )
+
+    handler = cast(_FakeArmorHandler, guardrail.async_handler)
+    response_scans = [call for call in handler.calls if "modelResponseData" in call]
+    assert response_scans, "expected a model_response scan of the assembled response"
+    assert "assembled output text" in response_scans[0]["modelResponseData"]["text"]
+    assert _metadata_entries(out_kwargs)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scans_anthropic_messages_model_response():
+    """/v1/messages logs a ModelResponse; the output scan must extract the assistant text."""
+    guardrail = _logging_only_guardrail()
+    kwargs = _logged_kwargs()
+    kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
+    response = _chat_response("anthropic assembled text")
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="anthropic_messages"
+    )
+
+    assert out_result is response
+    handler = cast(_FakeArmorHandler, guardrail.async_handler)
+    response_scans = [call for call in handler.calls if "modelResponseData" in call]
+    assert response_scans
+    assert "anthropic assembled text" in response_scans[0]["modelResponseData"]["text"]
+    assert _metadata_entries(out_kwargs)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_skips_output_scan_when_no_assembled_response():
+    guardrail = _logging_only_guardrail()
+    kwargs = _logged_kwargs()
+
+    await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion")
+
+    handler = cast(_FakeArmorHandler, guardrail.async_handler)
+    assert all("modelResponseData" not in call for call in handler.calls)
+
+
+@pytest.mark.asyncio
+async def test_native_post_call_mode_ignores_logging_hook():
+    handler = _FakeArmorHandler([_clean_armor_response()])
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-post",
+        event_hook=GuardrailEventHooks.post_call,
+        async_handler=handler,
+        access_token_provider=_async_token_provider,
+    )
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_kwargs is kwargs
+    assert out_result is response
+    assert handler.calls == []
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_records_flagged_without_raising():
+    guardrail = _logging_only_guardrail([_flagged_armor_response()])
+    request_data = {"metadata": {}}
+    inputs = {"texts": ["forbidden output"]}
+
+    result = await guardrail.apply_guardrail(
+        inputs=inputs,
+        request_data=request_data,
+        input_type="response",
+    )
+
+    assert result == inputs
+    entries = request_data["metadata"]["standard_logging_guardrail_information"]
+    assert entries[-1]["guardrail_status"] == "guardrail_flagged"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_transport_error():
+    guardrail = _logging_only_guardrail([httpx.ConnectError("boom"), httpx.ConnectError("boom")])
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    failed = [e for e in entries if e["guardrail_status"] == "guardrail_failed_to_respond"]
+    assert failed
+    assert all(e["guardrail_provider"] == "model_armor" for e in failed)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_flagged_prompt_still_scans_response():
+    """A flagged input scan must not abort the output scan; both verdicts are recorded."""
+    guardrail = _logging_only_guardrail(
+        [_flagged_armor_response(), _flagged_armor_response()]
+    )
+    response = _chat_response("flagged output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, _ = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    handler = cast(_FakeArmorHandler, guardrail.async_handler)
+    sources = ["user_prompt" if "userPromptData" in call else "model_response" for call in handler.calls]
+    assert sources == ["user_prompt", "model_response"]
+    entries = _metadata_entries(out_kwargs)
+    flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
+    assert len(flagged) == 2
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
+    """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a
+    non-logging_only instance must signal the block so flagged text is not returned as clean."""
+    handler = _FakeArmorHandler([_flagged_armor_response()])
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-pre",
+        event_hook=GuardrailEventHooks.pre_call,
+        async_handler=handler,
+        access_token_provider=_async_token_provider,
+    )
+    request_data = {"metadata": {}}
+
+    with pytest.raises(HTTPException) as exc_info:
+        await guardrail.apply_guardrail(
+            inputs={"texts": ["forbidden prompt"]},
+            request_data=request_data,
+            input_type="request",
+        )
+
+    assert exc_info.value.status_code == 400
+    entries = request_data["metadata"]["standard_logging_guardrail_information"]
+    flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
+    assert len(flagged) == 1

tests/test_litellm/proxy/guardrails/test_init_guardrails.py

@@ -71,6 +71,52 @@ def test_initialize_bedrock_forwards_chunk_budget_chars():
     assert initialized[-1].chunk_budget_chars == 60_000
 
 
+def test_initialize_bedrock_forwards_contextual_grounding_from_messages():
+    """`contextual_grounding_from_messages: true` in config.yaml must make the post-call
+    payload carry the plain system prompt and user turn as grounding_source and query."""
+    import litellm
+    from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
+    from litellm.types.utils import Choices, Message, ModelResponse
+
+    test_guardrail = {
+        "guardrail_name": "test_bedrock_grounding_from_messages",
+        "litellm_params": {
+            "guardrail": SupportedGuardrailIntegrations.BEDROCK.value,
+            "mode": "post_call",
+            "guardrailIdentifier": "test-guardrail",
+            "guardrailVersion": "DRAFT",
+            "contextual_grounding_from_messages": True,
+        },
+    }
+    messages = [
+        {"role": "system", "content": "Returns are accepted for 30 days."},
+        {"role": "user", "content": "How long is the return window?"},
+    ]
+    response = ModelResponse(
+        choices=[Choices(index=0, message=Message(role="assistant", content="30 days."), finish_reason="stop")]
+    )
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [
+            {"text": {"text": "Returns are accepted for 30 days.", "qualifiers": ["grounding_source"]}},
+            {"text": {"text": "How long is the return window?", "qualifiers": ["query"]}},
+            {"text": {"text": "30 days.", "qualifiers": ["guard_content"]}},
+        ],
+    }
+
+    guardrail_handler = InMemoryGuardrailHandler()
+    guardrail_handler.initialize_guardrail(guardrail=test_guardrail)
+
+    initialized = [
+        callback
+        for callback in litellm.callbacks
+        if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_grounding_from_messages"
+    ]
+    assert initialized, "bedrock guardrail was not registered as a callback"
+    actual_request = initialized[-1].convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
+    assert json.loads(json.dumps(actual_request)) == expected_request
+
+
 def test_initialize_guardrail_preserves_guardrail_info():
     """
     Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the

tests/test_litellm/responses/test_streaming_iterator_error_events.py

@@ -1,9 +1,14 @@
 """
 Regression: in-stream error events (type="error", type="response.failed") must
 raise instead of being returned as benign chunks, mirroring chat streaming
-semantics (_handle_stream_fallback_error): non-retriable 4xx (except 429)
-raise litellm.APIError directly; 429 and 5xx are wrapped in
-MidStreamFallbackError so the Router's mid-stream fallback machinery fires.
+semantics (_handle_stream_fallback_error). The event's code, type and status go
+through litellm.exception_type, so each event raises the same typed exception
+the non-streaming path raises for that provider error: non-retriable 4xx
+(except 429) raise that typed exception directly, so a context-length event
+surfaces as ContextWindowExceededError(400) with no MidStreamFallbackError
+wrapping, while 429, 5xx and ContentPolicyViolationError are wrapped in
+MidStreamFallbackError so the Router's mid-stream fallback machinery fires and
+its content_policy_fallbacks dispatch sees the trigger it matches on.
 
 Status mapping must consider both the OpenAI error `type` (e.g.
 "invalid_request_error") and `code` (e.g. "invalid_prompt",
@@ -66,12 +71,12 @@ def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(
     with pytest.raises(MidStreamFallbackError) as exc_info:
         iterator._maybe_raise_for_error_event(chunk)
     assert exc_info.value.status_code == 500
-    assert isinstance(exc_info.value.original_exception, litellm.APIError)
+    assert isinstance(exc_info.value.original_exception, litellm.InternalServerError)
     assert exc_info.value.original_exception.status_code == 500
 
 
 def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fallback():
-    """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped APIError."""
+    """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped RateLimitError."""
     iterator = _make_iterator()
     chunk = _make_error_chunk("tokens", "rate_limit_exceeded", "Too many requests")
     with pytest.raises(MidStreamFallbackError) as exc_info:
@@ -79,15 +84,15 @@ def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fall
     assert exc_info.value.status_code == 429
     assert exc_info.value.generated_content == ""
     assert exc_info.value.is_pre_first_chunk is True
-    assert isinstance(exc_info.value.original_exception, litellm.APIError)
+    assert isinstance(exc_info.value.original_exception, litellm.RateLimitError)
     assert exc_info.value.original_exception.status_code == 429
 
 
 def test_maybe_raise_for_error_event_maps_invalid_request_type_to_400():
-    """Client errors classified via the `type` field must raise APIError directly (no fallback)."""
+    """Client errors classified via the `type` field must raise BadRequestError directly (no fallback)."""
     iterator = _make_iterator()
     chunk = _make_error_chunk("invalid_request_error", "invalid_prompt", "bad request")
-    with pytest.raises(litellm.APIError) as exc_info:
+    with pytest.raises(litellm.BadRequestError) as exc_info:
         iterator._maybe_raise_for_error_event(chunk)
     assert exc_info.value.status_code == 400
     assert not isinstance(exc_info.value, MidStreamFallbackError)
@@ -99,12 +104,86 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400():
     chunk = Mock()
     chunk.type = "error"
     chunk.error = {"code": "context_length_exceeded", "message": "too long"}
-    with pytest.raises(litellm.APIError) as exc_info:
+    with pytest.raises(litellm.BadRequestError) as exc_info:
         iterator._maybe_raise_for_error_event(chunk)
     assert exc_info.value.status_code == 400
     assert not isinstance(exc_info.value, MidStreamFallbackError)
 
 
+def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly():
+    """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming
+    path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream
+    fallback, preserving the direct-SDK 400 contract from issue #15785."""
+    iterator = _make_iterator()
+    chunk = _make_error_chunk(
+        "invalid_request_error",
+        "context_length_exceeded",
+        "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.",
+    )
+    with pytest.raises(litellm.ContextWindowExceededError) as exc_info:
+        iterator._maybe_raise_for_error_event(chunk)
+    assert exc_info.value.status_code == 400
+    assert not isinstance(exc_info.value, MidStreamFallbackError)
+    assert "maximum context length" in str(exc_info.value)
+
+
+CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream."
+
+
+@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"])
+def test_maybe_raise_for_error_event_wraps_content_policy_violation_for_content_policy_fallbacks(
+    custom_llm_provider: str,
+):
+    """Regression: a content_policy_violation error event used to raise a bare APIError, so the Router's
+    content_policy_fallbacks never fired. It must map to ContentPolicyViolationError (the same exception the
+    non-streaming path raises) and be wrapped so the Router's mid-stream fallback catches it."""
+    iterator = _make_iterator()
+    iterator.custom_llm_provider = custom_llm_provider
+    chunk = _make_error_chunk("invalid_request_error", "content_policy_violation", CONTENT_POLICY_MESSAGE)
+    with pytest.raises(MidStreamFallbackError) as exc_info:
+        iterator._maybe_raise_for_error_event(chunk)
+    assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError)
+    assert exc_info.value.original_exception.status_code == 400
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.is_pre_first_chunk is True
+    assert CONTENT_POLICY_MESSAGE in str(exc_info.value.original_exception)
+
+
+def test_maybe_raise_for_response_failed_event_wraps_content_policy_violation():
+    iterator = _make_iterator()
+    chunk = _make_failed_chunk(
+        {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE}
+    )
+    with pytest.raises(MidStreamFallbackError) as exc_info:
+        iterator._maybe_raise_for_error_event(chunk)
+    assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError)
+
+
+@pytest.mark.parametrize(
+    "error_type,error_code,expected_exception",
+    [
+        ("invalid_request_error", "content_policy_violation", litellm.ContentPolicyViolationError),
+        ("tokens", "rate_limit_exceeded", litellm.RateLimitError),
+        ("invalid_request_error", "insufficient_quota", litellm.RateLimitError),
+        ("server_error", "internal_error", litellm.InternalServerError),
+        ("invalid_request_error", "invalid_prompt", litellm.BadRequestError),
+        ("invalid_request_error", "model_not_found", litellm.NotFoundError),
+        ("server_error", "vector_store_timeout", litellm.Timeout),
+    ],
+)
+def test_error_event_raises_the_same_typed_exception_as_the_non_streaming_path(
+    error_type: str, error_code: str, expected_exception: type[Exception]
+):
+    iterator = _make_iterator()
+    chunk = _make_error_chunk(error_type, error_code, "provider message")
+    with pytest.raises((MidStreamFallbackError, expected_exception)) as exc_info:
+        iterator._maybe_raise_for_error_event(chunk)
+    raised = exc_info.value
+    typed_exception = raised.original_exception if isinstance(raised, MidStreamFallbackError) else raised
+    assert type(typed_exception) is expected_exception
+    assert "provider message" in str(typed_exception)
+
+
 def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429():
     """OpenAI returns HTTP 429 for insufficient_quota; it must not map to 400 even though its type
     is invalid_request_error-adjacent, and it must be wrapped for fallback."""
@@ -113,6 +192,7 @@ def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429():
     with pytest.raises(MidStreamFallbackError) as exc_info:
         iterator._maybe_raise_for_error_event(chunk)
     assert exc_info.value.status_code == 429
+    assert isinstance(exc_info.value.original_exception, litellm.RateLimitError)
 
 
 def test_maybe_raise_for_error_event_passes_through_normal_chunk():
@@ -186,10 +266,40 @@ async def test_async_iterator_raises_mid_stream_fallback_on_rate_limit_error_eve
     assert exc_info.value.status_code == 429
     assert exc_info.value.is_pre_first_chunk is True
     assert exc_info.value.generated_content == ""
-    assert isinstance(exc_info.value.original_exception, litellm.APIError)
+    assert isinstance(exc_info.value.original_exception, litellm.RateLimitError)
     assert exc_info.value.original_exception.status_code == 429
 
 
+@pytest.mark.asyncio
+async def test_async_iterator_content_policy_violation_after_first_chunk_carries_generated_content():
+    """The customer's case: text streams, then the provider halts the stream with a
+    content_policy_violation error event. The iterator must surface ContentPolicyViolationError
+    inside MidStreamFallbackError, together with the text already streamed."""
+    iterator = _make_async_iterator_with_events(
+        [
+            {"type": "response.output_text.delta", "delta": "partial "},
+            {
+                "type": "error",
+                "error": {
+                    "type": "invalid_request_error",
+                    "code": "content_policy_violation",
+                    "message": CONTENT_POLICY_MESSAGE,
+                },
+            },
+        ]
+    )
+
+    stream = aiter(iterator)
+    first_chunk = await anext(stream)
+    assert first_chunk is not None
+
+    with pytest.raises(MidStreamFallbackError) as exc_info:
+        await anext(stream)
+    assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError)
+    assert exc_info.value.is_pre_first_chunk is False
+    assert exc_info.value.generated_content == "partial "
+
+
 @pytest.mark.asyncio
 async def test_async_iterator_error_after_first_chunk_carries_generated_content():
     """An error after streamed output must expose the accumulated text so the router's
@@ -205,14 +315,13 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content(
         ]
     )
 
-    chunks = []
-    async def _drain():
-        async for chunk in iterator:
-            chunks.append(chunk)
+    stream = aiter(iterator)
+    first_chunk = await anext(stream)
+    second_chunk = await anext(stream)
+    assert first_chunk is not None and second_chunk is not None
 
     with pytest.raises(MidStreamFallbackError) as exc_info:
-        await _drain()
-    assert len(chunks) == 2
+        await anext(stream)
     assert exc_info.value.status_code == 500
     assert exc_info.value.is_pre_first_chunk is False
     assert exc_info.value.generated_content == "hello world"
@@ -265,7 +374,7 @@ def test_handle_logging_failed_response_maps_rate_limit_to_429():
     ):
         iterator._handle_logging_failed_response()
     logged_exception = mock_run_async.call_args.kwargs["exception"]
-    assert isinstance(logged_exception, litellm.APIError)
+    assert isinstance(logged_exception, litellm.RateLimitError)
     assert logged_exception.status_code == 429
     assert "throttled" in str(logged_exception)
 
@@ -282,10 +391,28 @@ def test_handle_logging_failed_response_maps_type_field_to_400():
     ):
         iterator._handle_logging_failed_response()
     logged_exception = mock_run_async.call_args.kwargs["exception"]
-    assert isinstance(logged_exception, litellm.APIError)
+    assert isinstance(logged_exception, litellm.BadRequestError)
     assert logged_exception.status_code == 400
 
 
+def test_handle_logging_failed_response_logs_content_policy_violation():
+    """Failure logging must record the same typed exception the stream raises, so logging
+    integrations see a content policy violation instead of a generic APIError."""
+    iterator = _make_iterator()
+    iterator.completed_response = _make_failed_chunk(
+        {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE}
+    )
+    with (
+        patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async,
+        patch.object(import_module("litellm.responses.streaming_iterator"), "executor"),
+    ):
+        iterator._handle_logging_failed_response()
+    logged_exception = mock_run_async.call_args.kwargs["exception"]
+    assert isinstance(logged_exception, litellm.ContentPolicyViolationError)
+    assert logged_exception.status_code == 400
+    assert CONTENT_POLICY_MESSAGE in str(logged_exception)
+
+
 def test_handle_logging_failed_response_records_usage_and_cost():
     """Usage on a response.failed event must reach failure spend accounting via combined_usage_object."""
     iterator = _make_iterator()
@@ -357,7 +484,7 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event():
         for _ in iterator:
             pass
     assert exc_info.value.status_code == 429
-    assert isinstance(exc_info.value.original_exception, litellm.APIError)
+    assert isinstance(exc_info.value.original_exception, litellm.RateLimitError)
 
 
 def test_every_openai_sdk_response_error_code_has_explicit_status_mapping():
@@ -413,7 +540,7 @@ def test_maybe_raise_for_response_failed_event_maps_image_code_to_400():
     chunk = Mock()
     chunk.type = "response.failed"
     chunk.response = mock_response_obj
-    with pytest.raises(litellm.APIError) as exc_info:
+    with pytest.raises(litellm.BadRequestError) as exc_info:
         iterator._maybe_raise_for_error_event(chunk)
     assert exc_info.value.status_code == 400
     assert not isinstance(exc_info.value, MidStreamFallbackError)

tests/test_litellm/test_cost_calculator.py

@@ -3909,6 +3909,57 @@ def _batch_cache_usage() -> Usage:
     )
 
 
+def test_batch_cost_calculator_prices_multimodal_tokens_at_modality_rates():
+    from litellm.cost_calculator import batch_cost_calculator
+
+    model_info: ModelInfo = {
+        "input_cost_per_token_batches": 1e-7,
+        "input_cost_per_audio_token_batches": 3.25e-6,
+        "input_cost_per_image_token_batches": 2.25e-7,
+        "input_cost_per_video_token_batches": 6e-6,
+    }
+    usage = Usage(
+        prompt_tokens=100,
+        completion_tokens=0,
+        total_tokens=100,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            audio_tokens=64,
+            image_tokens=10,
+            video_tokens=6,
+        ),
+    )
+
+    prompt_cost, _ = batch_cost_calculator(
+        usage=usage,
+        model="gemini-embedding-2",
+        custom_llm_provider="vertex_ai",
+        model_info=model_info,
+    )
+
+    assert prompt_cost == pytest.approx(20 * 1e-7 + 64 * 3.25e-6 + 10 * 2.25e-7 + 6 * 6e-6)
+
+
+def test_batch_cost_calculator_falls_back_to_text_batch_rate_for_modalities():
+    from litellm.cost_calculator import batch_cost_calculator
+
+    model_info: ModelInfo = {"input_cost_per_token_batches": 1e-7}
+    usage = Usage(
+        prompt_tokens=100,
+        completion_tokens=0,
+        total_tokens=100,
+        prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=64),
+    )
+
+    prompt_cost, _ = batch_cost_calculator(
+        usage=usage,
+        model="gemini-embedding-2",
+        custom_llm_provider="vertex_ai",
+        model_info=model_info,
+    )
+
+    assert prompt_cost == pytest.approx(100 * 1e-7)
+
+
 def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate():
     """
     LIT-4008 regression: anthropic batch usage is dominated by cache tokens.

tests/test_litellm/test_router.py

@@ -3694,6 +3694,111 @@ async def test_aresponses_streaming_iterator_fallback():
     assert call_kwargs["disable_fallbacks"] is False
 
 
+@pytest.mark.asyncio
+async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback():
+    """Regression: a mid-stream content_policy_violation error event never reached
+    content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not
+    catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the
+    fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched.
+    The stream below is the customer's shape: a raw OpenAI error event with code
+    content_policy_violation, transformed by the real OpenAI config, and the router must
+    call the content_policy_fallbacks target, not the general fallbacks one."""
+    from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+    from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
+
+    router = litellm.Router(
+        model_list=[
+            {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}},
+            {
+                "model_name": "content-fallback",
+                "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"},
+            },
+            {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}},
+        ],
+        fallbacks=[{"primary": ["general-fallback"]}],
+        content_policy_fallbacks=[{"primary": ["content-fallback"]}],
+    )
+    error_event = {
+        "type": "error",
+        "sequence_number": 2,
+        "error": {
+            "type": "invalid_request_error",
+            "code": "content_policy_violation",
+            "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.",
+            "param": None,
+        },
+    }
+
+    async def aiter_bytes():
+        yield f"data: {json.dumps(error_event)}\n\n".encode()
+
+    raw_response = MagicMock()
+    raw_response.headers = {}
+    raw_response.aiter_bytes = aiter_bytes
+    logging_obj = MagicMock(spec=LiteLLMLogging)
+    logging_obj.model_call_details = {"litellm_params": {}}
+    logging_obj.completion_start_time = None
+    source = ResponsesAPIStreamingIterator(
+        response=raw_response,
+        model="gpt-5.4",
+        responses_api_provider_config=OpenAIResponsesAPIConfig(),
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+    fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")]
+    fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks))
+
+    wrapped = await router._aresponses_streaming_iterator(
+        response=source,
+        initial_kwargs={
+            "model": "primary",
+            "stream": True,
+            "input": "Hi",
+            "original_generic_function": fallback_call,
+        },
+    )
+    collected = [chunk async for chunk in wrapped]
+
+    assert collected == fallback_chunks
+    fallback_call.assert_awaited_once()
+    assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash"
+
+
+@pytest.mark.asyncio
+async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch():
+    """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the
+    ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope."""
+    router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini")
+    content_policy_error = litellm.ContentPolicyViolationError(
+        message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4"
+    )
+    src = _make_responses_iterator(
+        chunks=[MagicMock(type="response.created")],
+        error=MidStreamFallbackError(
+            message=str(content_policy_error),
+            model="openai/gpt-5.4",
+            llm_provider="openai",
+            original_exception=content_policy_error,
+            is_pre_first_chunk=True,
+        ),
+        model="openai/gpt-5.4",
+    )
+
+    with patch.object(
+        router,
+        "async_function_with_fallbacks_common_utils",
+        new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])),
+    ) as mock_fallback_utils:
+        wrapped = await router._aresponses_streaming_iterator(
+            response=src,
+            initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"},
+        )
+        [chunk async for chunk in wrapped]
+
+    mock_fallback_utils.assert_awaited_once()
+    assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "fallback_headers",

tests/test_litellm/test_utils.py

@@ -892,7 +892,10 @@ def validate_model_cost_values(model_data, exceptions=None):
         "input_cost_per_video_per_second_above_8s_interval",
         "input_cost_per_video_per_second_above_15s_interval",
         "input_cost_per_video_per_second_above_128k_tokens",
+        "input_cost_per_audio_token_batches",
+        "input_cost_per_image_token_batches",
         "input_cost_per_token_batches",
+        "input_cost_per_video_token_batches",
         "output_cost_per_token_batches",
         "input_cost_per_token_cache_hit",
         "cache_creation_input_token_cost",
@@ -1041,7 +1044,10 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
                 "input_cost_per_second": {"type": "number"},
                 "input_cost_per_token": {"type": "number"},
                 "input_cost_per_token_above_128k_tokens": {"type": "number"},
+                "input_cost_per_audio_token_batches": {"type": "number"},
+                "input_cost_per_image_token_batches": {"type": "number"},
                 "input_cost_per_token_batches": {"type": "number"},
+                "input_cost_per_video_token_batches": {"type": "number"},
                 "input_cost_per_token_cache_hit": {"type": "number"},
                 "input_cost_per_video_per_second": {"type": "number"},
                 "input_cost_per_video_per_second_above_8s_interval": {"type": "number"},
@@ -2946,7 +2952,7 @@ def test_model_info_for_openrouter_kimi_k2_5():
 
 
 def test_gemini_embedding_2_ga_in_cost_map():
-    """GA and Vertex preview gemini-embedding-2 entries align with multimodal unit pricing."""
+    """GA and Vertex preview gemini-embedding-2 entries align with multimodal token pricing."""
     import json
     from pathlib import Path
 
@@ -2968,9 +2974,15 @@ def test_gemini_embedding_2_ga_in_cost_map():
         assert info.get("mode") == "embedding"
         assert info.get("supports_multimodal") is True
         assert info.get("input_cost_per_token") == 2e-07
-        assert info.get("input_cost_per_image") == 0.00012
-        assert info.get("input_cost_per_audio_per_second") == 0.00016
-        assert info.get("input_cost_per_video_per_second") == 0.00079
+        assert info.get("input_cost_per_audio_token") == 6.5e-06
+        assert info.get("input_cost_per_image_token") == 4.5e-07
+        assert info.get("input_cost_per_video_token") == 1.2e-05
+        assert info.get("input_cost_per_audio_token_batches") == 3.25e-06
+        assert info.get("input_cost_per_image_token_batches") == 2.25e-07
+        assert info.get("input_cost_per_video_token_batches") == 6e-06
+        assert "input_cost_per_image" not in info
+        assert "input_cost_per_audio_per_second" not in info
+        assert "input_cost_per_video_per_second" not in info
         if provider in ("vertex_ai-embedding-models", "vertex_ai"):
             assert (
                 info.get("uses_embed_content") is True

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -29857,6 +29857,8 @@ export interface components {
             input_cost_per_audio_per_second_above_128k_tokens?: number | null;
             /** Input Cost Per Audio Token */
             input_cost_per_audio_token?: number | null;
+            /** Input Cost Per Audio Token Batches */
+            input_cost_per_audio_token_batches?: number | null;
             /** Input Cost Per Character */
             input_cost_per_character?: number | null;
             /** Input Cost Per Character Above 128K Tokens */
@@ -29867,6 +29869,8 @@ export interface components {
             input_cost_per_image_above_128k_tokens?: number | null;
             /** Input Cost Per Image Token */
             input_cost_per_image_token?: number | null;
+            /** Input Cost Per Image Token Batches */
+            input_cost_per_image_token_batches?: number | null;
             /** Input Cost Per Pixel */
             input_cost_per_pixel?: number | null;
             /** Input Cost Per Query */
@@ -29909,6 +29913,8 @@ export interface components {
             input_cost_per_video_per_second_above_8s_interval?: number | null;
             /** Input Cost Per Video Token */
             input_cost_per_video_token?: number | null;
+            /** Input Cost Per Video Token Batches */
+            input_cost_per_video_token_batches?: number | null;
             /** Itpm */
             itpm?: number | null;
             /** Keepalive Seconds */
@@ -31006,6 +31012,12 @@ export interface components {
              * @description Enable content moderation to check for harmful content (harassment, hate speech, etc.).
              */
             content_moderation_check?: boolean | null;
+            /**
+             * Contextual Grounding From Messages
+             * @description ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.
+             * @default false
+             */
+            contextual_grounding_from_messages: boolean;
             /**
              * Credentials
              * @description Path to Google Cloud credentials JSON file or JSON string
@@ -40071,6 +40083,8 @@ export interface components {
             input_cost_per_audio_per_second_above_128k_tokens?: number | null;
             /** Input Cost Per Audio Token */
             input_cost_per_audio_token?: number | null;
+            /** Input Cost Per Audio Token Batches */
+            input_cost_per_audio_token_batches?: number | null;
             /** Input Cost Per Character */
             input_cost_per_character?: number | null;
             /** Input Cost Per Character Above 128K Tokens */
@@ -40081,6 +40095,8 @@ export interface components {
             input_cost_per_image_above_128k_tokens?: number | null;
             /** Input Cost Per Image Token */
             input_cost_per_image_token?: number | null;
+            /** Input Cost Per Image Token Batches */
+            input_cost_per_image_token_batches?: number | null;
             /** Input Cost Per Pixel */
             input_cost_per_pixel?: number | null;
             /** Input Cost Per Query */
@@ -40123,6 +40139,8 @@ export interface components {
             input_cost_per_video_per_second_above_8s_interval?: number | null;
             /** Input Cost Per Video Token */
             input_cost_per_video_token?: number | null;
+            /** Input Cost Per Video Token Batches */
+            input_cost_per_video_token_batches?: number | null;
             /** Itpm */
             itpm?: number | null;
             /** Keepalive Seconds */