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

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

@@ -128,6 +128,9 @@ jobs:
       - name: check_fastuuid_usage
         run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
 
+      - name: check_py310_typing_imports
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py
+
       - name: check_e2e_no_raw_requests
         run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
 
@@ -145,3 +148,33 @@ jobs:
 
       - name: documentation_test_api_docs
         run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
+
+  python-310-import-smoke:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.10"
+
+      - name: Set up uv
+        uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - name: Install dependencies
+        run: uv sync --frozen --extra proxy --python 3.10
+
+      - run: uv run --no-sync python --version
+
+      - name: Import litellm
+        run: uv run --no-sync python -c "import litellm"
+
+      - name: Check litellm CLI
+        run: uv run --no-sync litellm --version

basedpyright-code-budget.json

@@ -57,7 +57,7 @@
     "limit": 5601
   },
   "reportMissingTypeArgument": {
-    "limit": 15290
+    "limit": 15288
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -105,7 +105,7 @@
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38332
+    "limit": 38324
   },
   "reportUnknownParameterType": {
     "limit": 19625
@@ -123,7 +123,7 @@
     "limit": 4
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 823
+    "limit": 819
   },
   "reportUntypedBaseClass": {
     "limit": 0

enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py

@@ -11,17 +11,35 @@ import sys
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
+import functools
 import tempfile
-from typing import Optional
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, ClassVar, Literal, Optional
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails._content_utils import walk_user_text
+from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 GUARDRAIL_NAME = "hide_secrets"
 
+GUARDRAIL_PROVIDER = "hide-secrets"
+
+# Per-invocation tally of redacted secrets by detect-secrets plugin type; None
+# means the guardrail did not run, so _process_response records nothing.
+_masked_entity_count: ContextVar[Optional[dict]] = ContextVar(
+    "hide_secrets_masked_entity_count", default=None
+)
+
 _custom_plugins_path = "file://" + os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "secrets_plugins"
 )
@@ -422,6 +440,10 @@ _default_detect_secrets_config = {
 
 
 class _ENTERPRISE_SecretDetection(CustomGuardrail):
+    # Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
+    # path skips should_run_check and never sees data["prompt"]).
+    use_native_lifecycle_hooks: ClassVar[bool] = True
+
     def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs):
         self.user_defined_detect_secrets_config = detect_secrets_config
         super().__init__(**kwargs)
@@ -455,6 +477,26 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
 
         return detected_secrets
 
+    def redact_text(self, text: str, source: str = "message") -> str:
+        """Replace every detected secret in ``text`` with ``[REDACTED]`` and
+        tally the detected types into the per-invocation masked-entity count."""
+        detected_secrets = self.scan_message_for_secrets(text)
+        if not detected_secrets:
+            return text
+        counts = _masked_entity_count.get()
+        if counts is not None:
+            for secret in detected_secrets:
+                counts[secret["type"]] = counts.get(secret["type"], 0) + 1
+        secret_types = [secret["type"] for secret in detected_secrets]
+        verbose_proxy_logger.warning(
+            f"Detected and redacted secrets in {source}: {secret_types}"
+        )
+        return functools.reduce(
+            lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
+            detected_secrets,
+            text,
+        )
+
     async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
         if user_api_key_dict.permissions is not None:
             if GUARDRAIL_NAME in user_api_key_dict.permissions:
@@ -463,7 +505,45 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
 
         return True
 
+    @log_guardrail_information
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional["LiteLLMLoggingObj"] = None,
+    ) -> GenericGuardrailAPIInputs:
+        """Unified-interface entrypoint, used by /guardrails/apply_guardrail
+        (the UI test playground). Proxied traffic keeps using
+        ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``."""
+        texts = inputs.get("texts")
+        if not texts or not any(texts):
+            return inputs
+        _masked_entity_count.set({})
+        return {**inputs, "texts": [self.redact_text(text) for text in texts]}
+
+    def _redact_prompt(self, data: dict) -> int:
+        """Redact ``data["prompt"]`` (the text-completion shape, which
+        ``walk_user_text`` does not cover) and return how many non-empty
+        strings were inspected."""
+        prompt = data.get("prompt")
+        if isinstance(prompt, str):
+            if not prompt:
+                return 0
+            data["prompt"] = self.redact_text(prompt, source="prompt")
+            return 1
+        if isinstance(prompt, list):
+            data["prompt"] = [  # mutable-ok: data["prompt"] is a list on the wire
+                self.redact_text(item, source="prompt")
+                if isinstance(item, str) and item
+                else item
+                for item in prompt
+            ]
+            return sum(1 for item in prompt if isinstance(item, str) and item)
+        return 0
+
     #### CALL HOOKS - proxy only ####
+    @log_guardrail_information
     async def async_pre_call_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -471,53 +551,84 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
         data: dict,
         call_type: str,  # "completion", "embeddings", "image_generation", "moderation"
     ):
+        _masked_entity_count.set(None)
         if await self.should_run_check(user_api_key_dict) is False:
             return
 
+        _masked_entity_count.set({})
+
         # Covers multimodal list content + Responses-API input.
-        def _redact_message_text(text: str) -> str:
-            detected_secrets = self.scan_message_for_secrets(text)
-            for secret in detected_secrets:
-                text = text.replace(secret["value"], "[REDACTED]")
-            if detected_secrets:
-                secret_types = [secret["type"] for secret in detected_secrets]
-                verbose_proxy_logger.warning(
-                    f"Detected and redacted secrets in message: {secret_types}"
-                )
-            return text
+        inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data)
 
-        walk_user_text(data, _redact_message_text)
+        if inspected == 0:
+            # Image-only, empty-text, and unsupported payloads inspected
+            # nothing, so recording "allow" would count a run that never
+            # looked at any content.
+            _masked_entity_count.set(None)
 
-        if "prompt" in data:
-            if isinstance(data["prompt"], str):
-                detected_secrets = self.scan_message_for_secrets(data["prompt"])
-                for secret in detected_secrets:
-                    data["prompt"] = data["prompt"].replace(
-                        secret["value"], "[REDACTED]"
-                    )
-                if len(detected_secrets) > 0:
-                    secret_types = [secret["type"] for secret in detected_secrets]
-                    verbose_proxy_logger.warning(
-                        f"Detected and redacted secrets in prompt: {secret_types}"
-                    )
-            elif isinstance(data["prompt"], list):
-                # Index back into the list — assigning to ``item`` would only
-                # rebind the loop variable and leave ``data["prompt"]``
-                # carrying the unredacted secret.
-                for idx, item in enumerate(data["prompt"]):
-                    if isinstance(item, str):
-                        detected_secrets = self.scan_message_for_secrets(item)
-                        for secret in detected_secrets:
-                            item = item.replace(secret["value"], "[REDACTED]")
-                        data["prompt"][idx] = item
-                        if len(detected_secrets) > 0:
-                            secret_types = [
-                                secret["type"] for secret in detected_secrets
-                            ]
-                            verbose_proxy_logger.warning(
-                                f"Detected and redacted secrets in prompt: {secret_types}"
-                            )
-
-        # ``data["input"]`` (Responses API and embeddings/moderation) is
-        # already covered by ``walk_user_text`` above.
         return
+
+    def _process_response(
+        self,
+        response: Optional[dict],
+        request_data: dict,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+        duration: Optional[float] = None,
+        event_type: Optional[GuardrailEventHooks] = None,
+        original_inputs: Optional[dict] = None,
+    ):
+        """Record allow/mask plus the masked-entity tally for a completed run.
+
+        Records nothing when the guardrail inspected nothing (opted-out key,
+        empty inputs) or when the instance has no guardrail_name (legacy
+        ``litellm_settings.callbacks`` deployments, which predate guardrail
+        telemetry and stay without it).
+        """
+        counts = _masked_entity_count.get()
+        _masked_entity_count.set(None)
+        if counts is None or self.guardrail_name is None:
+            return response
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response="mask" if counts else "allow",
+            request_data=request_data,
+            guardrail_status="success",
+            duration=duration,
+            start_time=start_time,
+            end_time=end_time,
+            event_type=event_type,
+            guardrail_provider=GUARDRAIL_PROVIDER,
+            masked_entity_count=counts,
+        )
+        return response
+
+    def _process_error(
+        self,
+        e: Exception,
+        request_data: dict,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+        duration: Optional[float] = None,
+        event_type: Optional[GuardrailEventHooks] = None,
+    ):
+        """Label the failed run with this guardrail's provider so error rows
+        group with the successful ones in the monitor. Nameless legacy
+        instances record nothing, matching ``_process_response``."""
+        _masked_entity_count.set(None)
+        if self.guardrail_name is None:
+            raise e
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=e,
+            request_data=request_data,
+            guardrail_status=(
+                "guardrail_intervened"
+                if self._is_guardrail_intervention(e)
+                else "guardrail_failed_to_respond"
+            ),
+            duration=duration,
+            start_time=start_time,
+            end_time=end_time,
+            event_type=event_type,
+            guardrail_provider=GUARDRAIL_PROVIDER,
+        )
+        raise e

litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs

@@ -186,68 +186,180 @@ mod tests {
     use serde_json::json;
 
     #[test]
-    fn supported_params_match_python_mistral_ocr_config() {
+    fn extract_header_is_a_supported_ocr_param() {
+        assert!(supported_ocr_params().contains(&"extract_header"));
+    }
+
+    #[test]
+    fn extract_footer_is_a_supported_ocr_param() {
+        assert!(supported_ocr_params().contains(&"extract_footer"));
+    }
+
+    #[test]
+    fn existing_ocr_params_remain_supported() {
+        for param in [
+            "pages",
+            "include_image_base64",
+            "image_limit",
+            "image_min_size",
+            "bbox_annotation_format",
+            "document_annotation_format",
+        ] {
+            assert!(supported_ocr_params().contains(&param));
+        }
+    }
+
+    #[test]
+    fn map_ocr_params_forwards_extract_header() {
+        let params = json!({"extract_header": true});
         assert_eq!(
-            supported_ocr_params(),
-            &[
-                "pages",
-                "include_image_base64",
-                "image_limit",
-                "image_min_size",
-                "bbox_annotation_format",
-                "document_annotation_format",
-                "document_annotation_prompt",
-                "extract_header",
-                "extract_footer",
-                "table_format",
-                "confidence_scores_granularity",
-                "include_blocks",
-                "id",
-            ]
+            map_ocr_params(params.as_object().unwrap()),
+            params.as_object().unwrap().clone()
+        );
+    }
+
+    #[test]
+    fn map_ocr_params_forwards_extract_footer() {
+        let params = json!({"extract_footer": true});
+        assert_eq!(
+            map_ocr_params(params.as_object().unwrap()),
+            params.as_object().unwrap().clone()
+        );
+    }
+
+    #[test]
+    fn map_ocr_params_forwards_extract_header_and_footer() {
+        let params = json!({"extract_header": true, "extract_footer": false});
+        assert_eq!(
+            map_ocr_params(params.as_object().unwrap()),
+            params.as_object().unwrap().clone()
         );
     }
 
     #[test]
     fn map_ocr_params_drops_unknown_params() {
-        let params = json!({
-            "extract_header": true,
-            "unsupported_param": "value",
-            "pages": [0, 1]
-        });
+        let params = json!({"extract_header": true, "unsupported_param": "value"});
         let mapped = map_ocr_params(params.as_object().unwrap());
-
         assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
-        assert_eq!(mapped.get("pages"), Some(&json!([0, 1])));
         assert!(!mapped.contains_key("unsupported_param"));
     }
 
     #[test]
-    fn transform_ocr_request_builds_mistral_body() {
+    fn new_ocr_params_are_supported() {
+        for param in [
+            "table_format",
+            "confidence_scores_granularity",
+            "document_annotation_prompt",
+            "include_blocks",
+            "id",
+        ] {
+            assert!(supported_ocr_params().contains(&param));
+        }
+    }
+
+    #[test]
+    fn map_ocr_params_forwards_new_ocr_params() {
+        for (param, value) in [
+            ("table_format", json!("html")),
+            ("confidence_scores_granularity", json!("word")),
+            (
+                "document_annotation_prompt",
+                json!("Extract all invoice line items"),
+            ),
+            ("include_blocks", json!(true)),
+            ("id", json!("req-123")),
+        ] {
+            let params = json!({param: value});
+            assert_eq!(
+                map_ocr_params(params.as_object().unwrap()),
+                params.as_object().unwrap().clone()
+            );
+        }
+    }
+
+    #[test]
+    fn transform_ocr_request_includes_each_optional_param() {
+        let document = json!({
+            "type": "document_url",
+            "document_url": "https://example.com/doc.pdf"
+        });
+        for (param, value) in [
+            ("table_format", json!("html")),
+            ("confidence_scores_granularity", json!("word")),
+            (
+                "document_annotation_prompt",
+                json!("Extract all invoice line items"),
+            ),
+            ("id", json!("req-123")),
+            ("extract_header", json!(true)),
+            ("include_blocks", json!(true)),
+            ("pages", json!([0, 1])),
+        ] {
+            let result = transform_ocr_request(
+                "mistral-ocr-latest",
+                document.clone(),
+                json!({param: value}).as_object().unwrap().clone(),
+            )
+            .expect("request should transform");
+            assert_eq!(result.data.get(param), Some(&value));
+            assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest")));
+            assert_eq!(result.data.get("document"), Some(&document));
+            assert_eq!(result.files, None);
+        }
+    }
+
+    #[test]
+    fn transform_ocr_request_includes_multiple_new_params() {
         let document = json!({
             "type": "document_url",
             "document_url": "https://example.com/doc.pdf"
         });
         let optional_params = json!({
-            "include_image_base64": true,
-            "table_format": "html"
+            "table_format": "html",
+            "confidence_scores_granularity": "page",
+            "extract_header": true
         })
         .as_object()
         .unwrap()
         .clone();
-
-        let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params)
+        let result = transform_ocr_request("mistral-ocr-latest", document, optional_params)
             .expect("request should transform");
-
+        assert_eq!(result.data.get("table_format"), Some(&json!("html")));
         assert_eq!(
-            result.data,
-            json!({
-                "model": "mistral-ocr-latest",
-                "document": document,
-                "include_image_base64": true,
-                "table_format": "html"
-            })
+            result.data.get("confidence_scores_granularity"),
+            Some(&json!("page"))
         );
-        assert_eq!(result.files, None);
+        assert_eq!(result.data.get("extract_header"), Some(&json!(true)));
+    }
+
+    #[test]
+    fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
+        let blocks = json!([{"type": "title", "content": "Invoice"}]);
+        let confidence_scores = json!({"page": 0.98});
+        let response = json!({
+            "pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}],
+            "model": "mistral-ocr-4-0",
+            "usage_info": {"pages_processed": 1}
+        });
+        let result =
+            transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform");
+        assert_eq!(result.pages[0].get("blocks"), Some(&blocks));
+        assert_eq!(
+            result.pages[0].get("confidence_scores"),
+            Some(&confidence_scores)
+        );
+    }
+
+    #[test]
+    fn transform_ocr_response_preserves_ocr4_page_fields() {
+        let response = json!({
+            "pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}],
+            "model": "mistral-ocr-4-0",
+            "usage_info": {"pages_processed": 1}
+        });
+        let result = transform_ocr_response("mistral-ocr-4-0", response.clone())
+            .expect("response should transform");
+        assert_eq!(result.pages[0], response["pages"][0]);
     }
 
     #[test]

litellm/constants.py

@@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model
 SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
 CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
+SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
 LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
     "Truncation is a DB storage safeguard. "

litellm/cost_calculator.py

@@ -4,6 +4,7 @@ import logging
 import time
 from collections.abc import Mapping, Sequence
 from functools import lru_cache
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 
 from httpx import Response
@@ -1164,6 +1165,12 @@ def _store_cost_breakdown_in_logging_obj(
         # Don't fail the main cost calculation if breakdown storage fails
 
 
+def _without_provider_stated_cost(usage: Usage | None) -> Usage | None:
+    if usage is None or getattr(usage, "cost", None) is None:
+        return usage
+    return usage.model_copy(update=MappingProxyType({"cost": None}))
+
+
 def completion_cost(
     completion_response: object | None = None,
     model: str | None = None,
@@ -1243,7 +1250,10 @@ def completion_cost(
         cache_creation_input_tokens: int | None = None
         cache_read_input_tokens: int | None = None
         audio_transcription_file_duration: float = 0.0
-        cost_per_token_usage_object: Final[Usage | None] = _get_usage_object(completion_response=completion_response)
+        provider_usage_object: Final = _get_usage_object(completion_response=completion_response)
+        cost_per_token_usage_object: Final[Usage | None] = (
+            _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object
+        )
         rerank_billed_units: RerankBilledUnits | None = None
 
         # Extract service_tier from optional_params if not provided directly

litellm/integrations/azure_storage/azure_storage.py

@@ -1,7 +1,9 @@
 import asyncio
 import os
 import time
+from collections.abc import Callable
 from datetime import datetime, timedelta
+from functools import cache
 from typing import Final
 
 from litellm._logging import verbose_logger
@@ -19,21 +21,40 @@ from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.secret_managers.get_azure_ad_token_provider import (
+    get_azure_ad_token_provider,
+)
+from litellm.types.secret_managers.get_azure_ad_token_provider import (
+    AzureCredentialType,
+)
 from litellm.types.utils import StandardLoggingPayload
 
+AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
+
+
+@cache
+def _cached_credential_chain_token_provider() -> Callable[[], str]:
+    return get_azure_ad_token_provider(
+        azure_scope=AZURE_STORAGE_TOKEN_SCOPE,
+        azure_credential=AzureCredentialType.DefaultAzureCredential,
+    )
+
 
 class AzureBlobStorageLogger(CustomBatchLogger):
     def __init__(
         self,
+        build_credential_chain_token_provider: Callable[
+            [], Callable[[], str]
+        ] = _cached_credential_chain_token_provider,
         **kwargs,
     ):
         try:
             verbose_logger.debug("AzureBlobStorageLogger: in init azure blob storage logger")
 
             # Env Variables used for Azure Storage Authentication
-            self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID")
-            self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID")
-            self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET")
+            self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") or None
+            self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") or None
+            self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") or None
             self.azure_storage_account_key: str | None = os.getenv("AZURE_STORAGE_ACCOUNT_KEY")
 
             # Required Env Variables for Azure Storage
@@ -55,6 +76,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
             # Internal variables used for Token based authentication
             self.azure_auth_token: str | None = None  # the Azure AD token to use for Azure Storage API requests
             self.token_expiry: datetime | None = None  # the expiry time of the currentAzure AD token
+            self._build_credential_chain_token_provider: Callable[[], Callable[[], str]] = (
+                build_credential_chain_token_provider
+            )
 
             asyncio.create_task(self.periodic_flush())
             self.flush_lock = asyncio.Lock()
@@ -231,10 +255,15 @@ class AzureBlobStorageLogger(CustomBatchLogger):
         """
         Wrapper to set self.azure_auth_token to a valid Azure AD token, refreshing if necessary
 
-        Refreshes the token when:
-        - Token is expired
-        - Token is not set
+        Without a service principal configured, the credential chain provider is read every
+        time; it caches internally and refreshes against the token's real expiry. The read runs
+        in a worker thread because the chain walk (IMDS probe, CLI subprocess) is blocking
         """
+        if self.tenant_id is None and self.client_id is None and self.client_secret is None:
+            token_provider: Final = self._build_credential_chain_token_provider()
+            self.azure_auth_token = await asyncio.to_thread(token_provider)
+            return
+
         # Check if token needs refresh
         if self._azure_ad_token_is_expired() or self.azure_auth_token is None:
             verbose_logger.debug("Azure AD token needs refresh")
@@ -273,13 +302,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
             tenant_id=tenant_id,
             client_id=client_id,
             client_secret=client_secret,
-            scope="https://storage.azure.com/.default",
+            scope=AZURE_STORAGE_TOKEN_SCOPE,
         )
-        token: Final = token_provider()
-
-        verbose_logger.debug("azure auth token %s", token)
-
-        return token
+        return token_provider()
 
     def _azure_ad_token_is_expired(self):
         """

litellm/integrations/custom_guardrail.py

@@ -1485,7 +1485,7 @@ def log_guardrail_information(func):
         if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
             original_inputs = kwargs.get("inputs")
 
-        logging_obj: Final = kwargs.get("logging_obj")
+        logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
         self_recorded_token: Final = _guardrail_self_recorded.set(False)
         try:
             response: Final = await func(*args, **kwargs)
@@ -1527,7 +1527,7 @@ def log_guardrail_information(func):
         if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
             original_inputs = kwargs.get("inputs")
 
-        logging_obj: Final = kwargs.get("logging_obj")
+        logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
         self_recorded_token: Final = _guardrail_self_recorded.set(False)
         try:
             response: Final = func(*args, **kwargs)

litellm/integrations/datadog/datadog_llm_obs.py

@@ -19,11 +19,13 @@ import httpx
 import litellm
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
+from litellm.constants import REDACTED_BY_LITELLM
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.datadog.datadog_handler import (
     get_datadog_base_url_from_env,
     get_datadog_service,
     get_datadog_tags,
+    normalize_datadog_tag_value,
 )
 from litellm.integrations.datadog.datadog_mock_client import (
     create_mock_datadog_client,
@@ -34,6 +36,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     convert_content_list_to_str,
     handle_any_messages_to_chat_completion_str_messages_conversion,
 )
+from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.llms.custom_httpx.http_handler import (
@@ -43,6 +46,7 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
 from litellm.types.integrations.datadog_llm_obs import *
 from litellm.types.utils import (
+    PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
     CallTypes,
     StandardLoggingGuardrailInformation,
     StandardLoggingPayload,
@@ -52,6 +56,120 @@ from litellm.types.utils import (
 _EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
 _EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
 _MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
+_SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
+    {"agent", "assistant", "developer", "function", "model", "system", "tool", "user"}
+)
+
+_PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset(
+    {
+        "routing_decision",
+        "requester_metadata",
+        "prompt_management_metadata",
+        "mcp_tool_call_metadata",
+        "vector_store_request_metadata",
+    }
+)
+
+_ROUTER_SPAN_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "tier": "router_tier",
+        "cause": "router_cause",
+        "score": "router_score",
+        "escalated": "router_escalated",
+        "signals": "router_signals",
+        "routed_model": "routed_model",
+    }
+)
+_ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "router_escalated", "routed_model")
+_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS)
+
+
+def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]:
+    metadata: Final = standard_logging_payload.get("metadata")
+    return metadata or _EMPTY_MAPPING
+
+
+def _router_span_fields(
+    standard_logging_payload: StandardLoggingPayload, redact_prompt_text: bool
+) -> Mapping[str, object]:
+    """Flatten the auto-router decision, omitting prompt-quoting fields when redaction is enabled."""
+    routing_decision: Final = _mapping_field(_metadata_of(standard_logging_payload), "routing_decision")
+    if not routing_decision:
+        return _EMPTY_MAPPING
+    escalated: Final = bool(routing_decision.get("escalated") or routing_decision.get("context_escalated"))
+    return MappingProxyType(
+        {
+            _ROUTER_SPAN_FIELDS[record_field]: value
+            for record_field, value in (*routing_decision.items(), ("escalated", escalated))
+            if record_field in _ROUTER_SPAN_FIELDS
+            and value is not None
+            and not (redact_prompt_text and record_field in PROMPT_QUOTING_ROUTING_DECISION_FIELDS)
+        }
+    )
+
+
+def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]:
+    """The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text."""
+    return MappingProxyType(
+        {
+            field: value
+            for field, value in standard_logging_metadata.items()
+            if field not in _PROMPT_CARRYING_METADATA_FIELDS
+        }
+    )
+
+
+def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]:
+    """Each message's shape with its content replaced and tool payloads dropped; no message is invented."""
+    return tuple(
+        {
+            "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "",
+            "content": REDACTED_BY_LITELLM,
+        }
+        for message in messages
+        for role in (message.get("role", ""),)
+    )
+
+
+def _cost_dimension_tags(
+    standard_logging_payload: StandardLoggingPayload, router_fields: Mapping[str, object]
+) -> tuple[str, ...]:
+    """The dimensions LLM Obs breaks token and cost metrics down by, as span tags."""
+    metadata: Final = _metadata_of(standard_logging_payload)
+    dimensions: Final = (
+        ("user", metadata.get("user_api_key_user_id")),
+        ("key_alias", metadata.get("user_api_key_alias")),
+        ("model_group", standard_logging_payload.get("model_group")),
+        *((dimension, router_fields.get(dimension)) for dimension in _ROUTER_DIMENSIONS),
+    )
+    return tuple(
+        f"{key}:{normalized}"
+        for key, value in dimensions
+        if value is not None and (normalized := normalize_datadog_tag_value(value)) != ""
+    )
+
+
+def _declared_cost_tags(span_tags: Sequence[str]) -> tuple[str, ...]:
+    """Declare only cost dimensions carrying a value on this span."""
+    present: Final = frozenset(key for tag in span_tags if (key := tag.partition(":")[0]) and tag.partition(":")[2])
+    return tuple(dimension for dimension in _COST_DIMENSIONS if dimension in present)
+
+
+def _reasoning_output_tokens(usage_object: Mapping[str, Any] | None) -> float:
+    """The provider's reasoning-token count, from either the chat or the responses spelling."""
+    if usage_object is None:
+        return 0.0
+    return next(
+        (
+            float(reasoning_tokens)
+            for details_field in ("completion_tokens_details", "output_tokens_details")
+            if isinstance(
+                reasoning_tokens := _mapping_field(usage_object, details_field).get("reasoning_tokens"), (int, float)
+            )
+            and not isinstance(reasoning_tokens, bool)
+        ),
+        0.0,
+    )
 
 
 def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
@@ -316,12 +434,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         dict_datadog_llm_obs_params: dict = {}
         if litellm.datadog_llm_observability_params is not None:
             if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams):
-                dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump()
+                dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump(exclude_unset=True)
             elif isinstance(litellm.datadog_llm_observability_params, dict):
                 # only allow params that are of DatadogLLMObsInitParams
                 dict_datadog_llm_obs_params = DatadogLLMObsInitParams(
                     **litellm.datadog_llm_observability_params
-                ).model_dump()
+                ).model_dump(exclude_unset=True)
         return dict_datadog_llm_obs_params
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
@@ -410,25 +528,40 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         if standard_logging_payload is None:
             raise Exception("DataDogLLMObs: standard_logging_object is not set")
 
-        metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
+        raw_metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
+        metadata: Final = raw_metadata if isinstance(raw_metadata, dict) else {}
+        redact_payload: Final = self._payload_logging_is_off(kwargs)
 
-        input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"]))
+        input_messages: Final = _to_dd_messages(standard_logging_payload.get("messages"))
+        output_messages: Final = self._get_response_messages(
+            standard_logging_payload=standard_logging_payload,
+            call_type=standard_logging_payload.get("call_type"),
+        )
+        input_meta: Final = InputMeta(messages=_redact_messages(input_messages) if redact_payload else input_messages)
         output_meta: Final = OutputMeta(
-            messages=self._get_response_messages(
-                standard_logging_payload=standard_logging_payload,
-                call_type=standard_logging_payload.get("call_type"),
-            )
+            messages=_redact_messages(output_messages) if redact_payload else output_messages
         )
 
         error_info: Final = self._assemble_error_info(standard_logging_payload)
 
-        metadata_parent_id: str | None = None
-        if isinstance(metadata, dict):
-            metadata_parent_id = metadata.get("parent_id")
+        raw_parent_id: Final = metadata.get("parent_id")
+        metadata_parent_id: Final[str | None] = str(raw_parent_id) if raw_parent_id else None
 
-        tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
+        tool_definitions: Final = (
+            () if redact_payload else _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
+        )
         span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id)
-        payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload)
+        router_fields: Final = _router_span_fields(standard_logging_payload, redact_prompt_text=redact_payload)
+        span_tags: Final = [
+            *get_datadog_tags(standard_logging_object=standard_logging_payload),
+            *_cost_dimension_tags(standard_logging_payload, router_fields),
+        ]
+        payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(
+            standard_logging_payload,
+            router_fields=router_fields,
+            cost_tags=_declared_cost_tags(span_tags),
+            redact_prompt_text=redact_payload,
+        )
 
         meta: Final[Meta] = {
             "kind": span_kind,
@@ -451,7 +584,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
             duration=int((end_time - start_time).total_seconds() * 1e9),
             metrics=metrics,
             status="error" if error_info else "ok",
-            tags=get_datadog_tags(standard_logging_object=standard_logging_payload),
+            tags=span_tags,
         )
 
         apm_trace_id: Final = self._get_apm_trace_id()
@@ -497,6 +630,13 @@ class DataDogLLMObsLogger(CustomBatchLogger):
                 )
         return error_info
 
+    def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool:
+        return (
+            bool(self.turn_off_message_logging)
+            or self.message_logging is not True
+            or should_redact_message_logging(dict(kwargs))
+        )
+
     def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics:
         """
         Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from.
@@ -513,10 +653,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         total_cost: Final = float(standard_logging_payload.get("response_cost", 0))
         time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload)
 
-        raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object")
+        raw_usage: Final = _metadata_of(standard_logging_payload).get("usage_object")
         usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None
         cache_read: Final = float(extract_cache_read_tokens(usage_object))
         cache_write: Final = float(extract_cache_creation_tokens(usage_object))
+        reasoning_output_tokens: Final = _reasoning_output_tokens(usage_object)
 
         metrics: Final[LLMMetrics] = {
             "input_tokens": prompt_tokens,
@@ -533,6 +674,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
                 if cache_read or cache_write
                 else {}
             ),
+            **({"reasoning_output_tokens": reasoning_output_tokens} if reasoning_output_tokens else {}),
         }
         return metrics
 
@@ -707,11 +849,21 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         # Default fallback for unknown or passthrough operations
         return "llm"
 
-    def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
+    def _get_dd_llm_obs_payload_metadata(
+        self,
+        standard_logging_payload: StandardLoggingPayload,
+        router_fields: Mapping[str, object] | None = None,
+        cost_tags: Sequence[str] = (),
+        redact_prompt_text: bool = False,
+    ) -> dict[str, object]:
         """
         Fields to track in DD LLM Observability metadata from litellm standard logging payload
         """
-        _metadata: Final[dict[str, object]] = {
+        raw_metadata: Final = _metadata_of(standard_logging_payload)
+        standard_logging_metadata: Final = (
+            _metadata_without_prompt_carriers(raw_metadata) if redact_prompt_text else raw_metadata
+        )
+        return {
             "model_name": standard_logging_payload.get("model", "unknown"),
             "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"),
             "id": standard_logging_payload.get("id", "unknown"),
@@ -719,26 +871,21 @@ class DataDogLLMObsLogger(CustomBatchLogger):
             "cache_hit": standard_logging_payload.get("cache_hit", "unknown"),
             "cache_key": standard_logging_payload.get("cache_key", "unknown"),
             "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0),
-            "guardrail_information": standard_logging_payload.get("guardrail_information", None),
+            "guardrail_information": (
+                None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None)
+            ),
             "is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload),
+            "latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)),
+            "spend_metrics": dict(self._get_spend_metrics(standard_logging_payload)),
+            **standard_logging_metadata,
+            **(router_fields or _EMPTY_MAPPING),
+            **(
+                {"_dd": {**_mapping_field(standard_logging_metadata, "_dd"), "cost_tags": list(cost_tags)}}
+                if cost_tags
+                else _EMPTY_MAPPING
+            ),
         }
 
-        #########################################################
-        # Add latency metrics to metadata
-        #########################################################
-        latency_metrics: Final = self._get_latency_metrics(standard_logging_payload)
-        _metadata.update({"latency_metrics": dict(latency_metrics)})
-
-        #########################################################
-        # Add spend metrics to metadata
-        #########################################################
-        spend_metrics: Final = self._get_spend_metrics(standard_logging_payload)
-        _metadata.update({"spend_metrics": dict(spend_metrics)})
-
-        _standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {}
-        _metadata.update(_standard_logging_metadata)
-        return _metadata
-
     def _get_latency_metrics(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsLatencyMetrics:
         """
         Get the latency metrics from the standard logging payload
@@ -808,7 +955,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0)
 
         # Get budget information from metadata
-        metadata: Final = standard_logging_payload.get("metadata", {})
+        metadata: Final = _metadata_of(standard_logging_payload)
 
         # API key max budget
         user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget")

litellm/integrations/websearch_interception/handler.py

@@ -10,9 +10,9 @@ import asyncio
 import math
 import uuid
 from collections.abc import AsyncIterator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast
 
-from typing_extensions import ReadOnly
+from typing_extensions import Never, ReadOnly
 
 import litellm
 from litellm._logging import verbose_logger

litellm/litellm_core_utils/streaming_handler.py

@@ -54,6 +54,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
 _SYNC_ITER_EXHAUSTED: Final = object()
 
 _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
+_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value})
 
 
 def _next_sync_or_exhausted(it: Any) -> object:
@@ -1886,8 +1887,8 @@ class CustomStreamWrapper:
     @staticmethod
     def _resolve_provider_reported_cost(usage_cost: object) -> float | None:
         """
-        Providers report usage.cost either as a number or, for Perplexity, as a
-        breakdown object whose total lives under ``total_cost``.
+        Providers report usage.cost either as a number or as a breakdown object
+        whose total lives under ``total_cost``.
         """
         if isinstance(usage_cost, bool):
             return None
@@ -1900,12 +1901,10 @@ class CustomStreamWrapper:
     @staticmethod
     def _propagate_usage_cost_to_hidden_params(
         response: "ModelResponse",
+        custom_llm_provider: str | None,
     ) -> None:
-        """
-        If the assembled response carries a provider-reported cost on
-        usage.cost, copy it into _hidden_params so litellm's cost
-        calculator uses it instead of a token-based estimate.
-        """
+        if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS:
+            return
         _usage: Final[Usage | None] = getattr(response, "usage", None)
         _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None))
         if _cost is not None:
@@ -2020,7 +2019,7 @@ class CustomStreamWrapper:
 
                 response = self.model_response_creator()
                 if complete_streaming_response is not None:
-                    self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
+                    self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
 
                     setattr(
                         response,
@@ -2270,7 +2269,7 @@ class CustomStreamWrapper:
 
             response: Final = self.model_response_creator()
             if complete_streaming_response is not None:
-                self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
+                self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider)
 
                 setattr(
                     response,

litellm/llms/a2a/chat/guardrail_translation/handler.py

@@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
-from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.llms.base_llm.guardrail_translation.base_translation import (
+    BaseTranslation,
+    StreamingScanKey,
+)
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
@@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation):
 
         return responses_so_far
 
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        _, valid_parsed = self._parse_streaming_responses(responses_so_far)
+        combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed)
+        return StreamingScanKey(texts=(combined_text,))
+
     def _parse_streaming_responses(
         self,
-        responses_so_far: list[object],
+        responses_so_far: Sequence[object],
     ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]:
         """Parse JSON-RPC items, returning aligned parsed list and valid entries."""
         parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far)

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
     LiteLLMAnthropicMessagesAdapter,
     is_provider_native_tool_dict,
 )
-from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.llms.base_llm.guardrail_translation.base_translation import (
+    BaseTranslation,
+    StreamingScanKey,
+)
 from litellm.llms.base_llm.guardrail_translation.utils import (
     anthropic_tool_name,
     anthropic_tool_names,
@@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     merge_guardrailed_scoped_messages,
     merge_returned_tools_into_request_tools,
     scoped_structured_message_indices,
+    stream_item_fingerprint,
 )
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
     AnthropicPassthroughLoggingHandler,
@@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation):
             inputs["model"] = response_model
         return inputs
 
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
+        return StreamingScanKey(
+            texts=(self.get_streaming_string_so_far(responses_so_far),),
+            tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
+            stream_ended=stream_ended,
+        )
+
+    @classmethod
+    def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]:
+        return tuple(
+            stream_item_fingerprint(block)
+            for item in responses_so_far
+            for event in cls._iter_sse_events(item)
+            if event.get("type") == "content_block_start"
+            and isinstance(block := event.get("content_block"), Mapping)
+            and block.get("type") == "tool_use"
+        )
+
     def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
         """
         Parse streaming responses and extract accumulated text content.

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -35,6 +35,22 @@ class StreamTransformSink:
     holdback_per_choice: dict[int, int] = field(default_factory=dict)
 
 
+@dataclass(frozen=True, slots=True)
+class StreamingScanKey:
+    """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
+    compare equal when the round would scan the same content again; ``stream_ended``
+    stays out of the comparison and only says whether the handler is on its
+    end-of-stream path, where an empty payload is still scanned today."""
+
+    texts: tuple[str, ...]
+    tool_calls: tuple[str, ...] = ()
+    stream_ended: bool = field(default=False, compare=False)
+
+    @property
+    def has_nothing_to_scan(self) -> bool:
+        return not self.stream_ended and not any(self.texts) and not self.tool_calls
+
+
 class BaseTranslation(ABC):
     @staticmethod
     def transform_user_api_key_dict_to_metadata(
@@ -151,6 +167,9 @@ class BaseTranslation(ABC):
         """
         return responses_so_far
 
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        return None
+
     def build_block_sse_chunks(
         self,
         exc: "ModifyResponseException",

litellm/llms/base_llm/guardrail_translation/utils.py

@@ -4,6 +4,8 @@ import json
 from collections.abc import Callable, Iterator, Sequence
 from typing import Any, Final, TypeVar
 
+from pydantic import BaseModel
+
 from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
 from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
 
@@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None:
     return getattr(item, field, None)
 
 
+def stream_item_fingerprint(item: object) -> str:
+    plain: Final = item.model_dump() if isinstance(item, BaseModel) else item
+    return json.dumps(plain, sort_keys=True, default=str)
+
+
+def stream_item_items(item: object, field: str) -> tuple[object, ...]:
+    value: Final = stream_item_field(item, field)
+    return tuple(value) if isinstance(value, (list, tuple)) else ()
+
+
 def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]:
     """
     ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked

litellm/llms/fireworks_ai/common_utils.py

@@ -2,6 +2,7 @@ from typing import Final
 
 from httpx import Headers
 
+from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 
@@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
     """
     Session id to send as `x-session-affinity`, or None when the caller gave none.
 
-    Deliberately does not fall back to `litellm_trace_id`: that is generated per
-    request (`str(uuid.uuid4())` when absent), so using it pins every request to a
-    different Fireworks node and prompt caching never hits.
+    Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the
+    proxy generated for a request that had none: both are per request, so using them
+    pins every request to a different Fireworks node and prompt caching never hits.
     """
     params: Final = litellm_params
+    metadata: Final = params.get("metadata")
+    if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
+        return None
     for key in ("litellm_session_id", "session_id"):
         value = params.get(key)
         if value:
             return str(value)
-    metadata: Final = params.get("metadata")
     if isinstance(metadata, dict):
         value = metadata.get("session_id")
         if value:

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -26,6 +26,7 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.llms.base_llm.guardrail_translation.base_translation import (
     BaseTranslation,
+    StreamingScanKey,
     StreamTransformSink,
 )
 from litellm.llms.base_llm.guardrail_translation.utils import (
@@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     role_out_of_guardrail_scope,
     scoped_structured_message_indices,
     stream_item_field,
+    stream_item_fingerprint,
+    stream_item_items,
 )
 from litellm.main import stream_chunk_builder
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
         """Block-only streaming path: run the guardrail so an in-flight BLOCK can
         terminate the stream. Text rewrites are not propagated to the client here
         (see ``_process_streaming_transform`` for the incremental_diff path)."""
-        # check if the stream has ended
-        has_stream_ended = False
-        for chunk in responses_so_far:
-            if chunk.choices and chunk.choices[0].finish_reason is not None:
-                has_stream_ended = True
-                break
+        has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
 
         if has_stream_ended:
             # convert to model response
@@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
         }
 
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
+        stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        return StreamingScanKey(
+            texts=tuple(self._combine_streaming_texts(chunks).values()),
+            tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
+            stream_ended=stream_ended,
+        )
+
+    @staticmethod
+    def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]:
+        return tuple(
+            stream_item_fingerprint(tool_call)
+            for chunk in responses_so_far
+            for choice in _stream_chunk_choices(chunk)
+            for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
+        )
+
+    @staticmethod
+    def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool:
+        first_choices: Final = tuple(
+            choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices
+        )
+        return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices)
+
     def _combine_streaming_texts(
-        self, responses_so_far: list["ModelResponseStream"]
+        self, responses_so_far: Sequence["ModelResponseStream"]
     ) -> dict[tuple[int, int | None], str]:
         """
         Combine all streaming chunks into complete text per choice.

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
     OpenAiResponsesToChatCompletionStreamIterator,
 )
-from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.llms.base_llm.guardrail_translation.base_translation import (
+    BaseTranslation,
+    StreamingScanKey,
+)
 from litellm.llms.base_llm.guardrail_translation.utils import (
     blocked_responses_stream_usage,
     stream_item_field,
+    stream_item_fingerprint,
+    stream_item_items,
 )
 from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
 from litellm.responses.litellm_completion_transformation.transformation import (
@@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation):
             )
         return responses_so_far
 
-    def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
+    def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
         """
         Check if the streaming has ended.
         """
         if not responses_so_far:
             return False
-        terminal_types: Final = {
-            ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
-            ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
-            ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
-        }
-        return responses_so_far[-1].get("type") in terminal_types
+        terminal_types: Final = frozenset(
+            (
+                ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
+                ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
+                ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
+            )
+        )
+        return stream_item_field(responses_so_far[-1], "type") in terminal_types
+
+    def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
+        if not responses_so_far or not hasattr(responses_so_far[-1], "get"):
+            return None
+        last_event: Final = responses_so_far[-1]
+        last_event_type: Final = stream_item_field(last_event, "type")
+        if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
+            return None
+        if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
+            return self._completed_response_scan_key(stream_item_field(last_event, "response"))
+        return StreamingScanKey(
+            texts=(self.get_streaming_string_so_far(responses_so_far),),
+            stream_ended=self._check_streaming_has_ended(responses_so_far),
+        )
+
+    @staticmethod
+    def _completed_response_scan_key(response: object) -> StreamingScanKey:
+        output_items: Final = stream_item_items(response, "output")
+        message_items: Final = tuple(
+            item for item in output_items if stream_item_field(item, "type") != "function_call"
+        )
+        return StreamingScanKey(
+            texts=tuple(
+                text
+                for item in message_items
+                for part in stream_item_items(item, "content")
+                if isinstance(text := stream_item_field(part, "text"), str) and text
+            ),
+            tool_calls=tuple(
+                stream_item_fingerprint(item)
+                for item in output_items
+                if stream_item_field(item, "type") == "function_call"
+            ),
+            stream_ended=True,
+        )
 
     def build_stream_error_items(
         self,
@@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             ),
         )
 
-    def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
+    def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
         """
         Get the string so far from the responses so far.
 
@@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation):
         """
         keyed_events: Final = tuple(
             (
-                (event.get("item_id"), event.get("output_index"), event.get("content_index")),
-                event.get("text"),
-                event.get("delta"),
+                (
+                    stream_item_field(event, "item_id"),
+                    stream_item_field(event, "output_index"),
+                    stream_item_field(event, "content_index"),
+                ),
+                stream_item_field(event, "text"),
+                stream_item_field(event, "delta"),
             )
             for event in responses_so_far
-            if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
+            if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str)
         )
 
         def part_text(part_key: tuple[object, object, object]) -> str:

litellm/llms/xai/chat/transformation.py

@@ -1,4 +1,5 @@
 from collections.abc import AsyncIterator, Iterator, Mapping
+from types import MappingProxyType
 from typing import Any, Final
 
 import httpx
@@ -11,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     filter_value_from_dict,
     strip_name_from_messages,
 )
-from litellm.llms.xai.common_utils import XAIModelInfo
+from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd
 from litellm.llms.xai.cost_calculator import (
     apply_server_side_tool_usage_details_to_usage,
 )
@@ -30,6 +31,13 @@ from ...openai.chat.gpt_transformation import (
 )
 
 
+def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None:
+    reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None))
+    if usage is None or reported_cost is None:
+        return None
+    return usage.model_copy(update=MappingProxyType({"cost": reported_cost}))
+
+
 class XAIChatConfig(OpenAIGPTConfig):
     @property
     def custom_llm_provider(self) -> str | None:
@@ -283,6 +291,9 @@ class XAIChatConfig(OpenAIGPTConfig):
 
         self._fold_reasoning_tokens_into_completion(response)
         self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None))
+        restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None))
+        if restated_usage is not None:
+            response.usage = restated_usage
         return response
 
     @staticmethod
@@ -411,4 +422,8 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler):
             XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"])
             XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"])
 
-        return super().chunk_parser(chunk)
+        parsed_chunk: Final = super().chunk_parser(chunk)
+        restated_usage: Final = _usage_restated_from_xai_ticks(getattr(parsed_chunk, "usage", None))
+        if restated_usage is not None:
+            parsed_chunk.usage = restated_usage
+        return parsed_chunk

litellm/llms/xai/common_utils.py

@@ -8,6 +8,17 @@ from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import ProviderSpecificModelInfo
 
+USD_TICKS_PER_DOLLAR: Final = 10_000_000_000
+
+
+def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None:
+    """xAI bills in ticks of a dollar: https://docs.x.ai/developers/cost-tracking"""
+    if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool):
+        return None
+    if cost_in_usd_ticks < 0:
+        return None
+    return cost_in_usd_ticks / USD_TICKS_PER_DOLLAR
+
 
 class XAIModelInfo(BaseLLMModelInfo):
     def get_provider_info(

litellm/llms/xai/cost_calculator.py

@@ -1,9 +1,11 @@
 """
 Helper util for handling XAI-specific cost calculation
+- Prefers the cost xAI reports on the response over recomputing it locally
 - Uses the generic cost calculator which already handles tiered pricing correctly
 - Handles XAI-specific reasoning token billing (billed as part of completion tokens)
 """
 
+import math
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final
 
@@ -36,6 +38,17 @@ def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping
     usage.prompt_tokens_details = prompt_tokens_details  # rebind-ok: write details onto caller usage
 
 
+def _cost_reported_by_xai(usage: "Usage") -> float | None:
+    reported_cost: Final[object] = getattr(usage, "cost", None)
+    if not isinstance(reported_cost, (int, float)) or isinstance(reported_cost, bool):
+        return None
+    if not math.isfinite(reported_cost):
+        return None
+    if reported_cost < 0:
+        return None
+    return float(reported_cost)
+
+
 def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
     """
     Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens.
@@ -48,6 +61,10 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
     Returns:
         Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
     """
+    reported_cost: Final = _cost_reported_by_xai(usage)
+    if reported_cost is not None:
+        return 0.0, reported_cost
+
     # XAI-specific completion cost: completion is billed as visible + reasoning
     # tokens. Detect when the transformation layer already folded them so we
     # don't double-count; fall back to raw xAI shape for callers that bypass
@@ -112,6 +129,9 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa
     Per-call rate comes from model_info.search_context_cost_per_query when set,
     otherwise the default xAI tools rate ($5 / 1k calls).
     """
+    if _cost_reported_by_xai(usage) is not None:
+        return 0.0
+
     details: Final = getattr(usage, "server_side_tool_usage_details", None)
     if not isinstance(details, Mapping):
         return 0.0

litellm/llms/xai/responses/transformation.py

@@ -1,17 +1,44 @@
-from typing import Any, Final
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final
+
+import httpx
 
 import litellm
 from litellm._logging import verbose_logger
 from litellm.constants import XAI_API_BASE
 from litellm.exceptions import AuthenticationError
 from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
-from litellm.llms.xai.common_utils import XAIModelInfo
+from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd
 from litellm.secret_managers.main import get_secret_str
-from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams
+from litellm.types.llms.openai import (
+    ResponseAPIUsage,
+    ResponseCompletedEvent,
+    ResponseFailedEvent,
+    ResponseIncompleteEvent,
+    ResponsesAPIOptionalRequestParams,
+    ResponsesAPIResponse,
+    ResponsesAPIStreamingResponse,
+)
 from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import LlmProviders
 
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import (
+        Logging as _LiteLLMLoggingObj,
+    )
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None:
+    reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None))
+    if usage is None or reported_cost is None:
+        return None
+    return usage.model_copy(update=MappingProxyType({"cost": reported_cost}))
+
 
 class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
     """
@@ -250,6 +277,41 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
 
         return f"{api_base}/responses"
 
+    def transform_response_api_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+    ) -> ResponsesAPIResponse:
+        response: Final = super().transform_response_api_response(
+            model=model,
+            raw_response=raw_response,
+            logging_obj=logging_obj,
+        )
+
+        restated_usage: Final = _usage_restated_from_xai_ticks(response.usage)
+        if restated_usage is not None:
+            response.usage = restated_usage
+        return response
+
+    def transform_streaming_response(
+        self,
+        model: str,
+        parsed_chunk: dict,  # mutable-ok: overrides the base class signature
+        logging_obj: LiteLLMLoggingObj,
+    ) -> ResponsesAPIStreamingResponse:
+        event: Final = super().transform_streaming_response(
+            model=model,
+            parsed_chunk=parsed_chunk,
+            logging_obj=logging_obj,
+        )
+        if not isinstance(event, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent)):
+            return event
+        restated_usage: Final = _usage_restated_from_xai_ticks(event.response.usage)
+        if restated_usage is not None:
+            event.response.usage = restated_usage
+        return event
+
     def supports_native_websocket(self) -> bool:
         """XAI does not support native WebSocket for Responses API"""
         return False

litellm/main.py

@@ -8595,9 +8595,19 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N
     return TextCompletionResponse(**response)
 
 
+_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value})
+
+
+def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool:
+    if logging_obj is None:
+        return False
+    provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider")
+    return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS
+
+
 def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None:
     usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
-    if isinstance(usage_cost, (int, float)):
+    if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj):
         return float(usage_cost)
     if logging_obj is not None:
         return None

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -820,20 +820,46 @@ def _should_strip_caller_authorization(
     if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate):
         return False
 
-    normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)}
-    has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None
+    has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
     if mcp_server.is_oauth_delegate:
         return not has_explicit_litellm_admission_header
-    admission_consumed_authorization_as_litellm_key: Final = (
-        user_api_key_auth is not None
-        and bool(getattr(user_api_key_auth, "api_key", None))
-        and not has_explicit_litellm_admission_header
-    )
-    return admission_consumed_authorization_as_litellm_key or (
+    return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or (
         user_api_key_auth is None and not has_explicit_litellm_admission_header
     )
 
 
+LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-"
+
+
+def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None:
+    return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None)
+
+
+def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool:
+    """Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``."""
+    return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
+
+
+def _authorization_is_litellm_admission_credential(
+    raw_headers: Mapping[str, str] | None,
+    user_api_key_auth: UserAPIKeyAuth | None,
+) -> bool:
+    """True when ``Authorization`` carries the LiteLLM key admission validated.
+
+    That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
+    same key in both headers.
+    """
+    if user_api_key_auth is None or not user_api_key_auth.api_key:
+        return False
+    admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
+    if not admission_header:
+        return True
+    authorization: Final = _raw_header_value(raw_headers, "authorization")
+    return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(
+        admission_header, "Bearer"
+    )
+
+
 def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
     """Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection.
 
@@ -3277,8 +3303,8 @@ class MCPServerManager:
     #########################################################
     @staticmethod
     def _extract_bearer_token(
-        oauth2_headers: dict[str, str] | None,
-        raw_headers: dict[str, str] | None,
+        oauth2_headers: Mapping[str, str] | None,
+        raw_headers: Mapping[str, str] | None,
     ) -> str | None:
         """Extract the bare Bearer token from oauth2_headers or raw_headers.
 
@@ -3298,10 +3324,29 @@ class MCPServerManager:
             return auth_value
         return None
 
+    @staticmethod
+    def _extract_subject_token(
+        oauth2_headers: Mapping[str, str] | None,
+        raw_headers: Mapping[str, str] | None,
+        user_api_key_auth: UserAPIKeyAuth | None,
+    ) -> str | None:
+        """The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key.
+
+        Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix,
+        any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``.
+        """
+        if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth):
+            return None
+        bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers)
+        if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX):
+            return None
+        return bearer
+
     def _obo_subject_token(
         self,
         server: MCPServer,
-        raw_headers: dict[str, str] | None,
+        raw_headers: Mapping[str, str] | None,
+        user_api_key_auth: UserAPIKeyAuth | None,
     ) -> str | None:
         """The caller's bearer as the token_exchange (OBO) subject token, for that mode only.
 
@@ -3311,7 +3356,7 @@ class MCPServerManager:
         """
         if server.auth_type != MCPAuth.oauth2_token_exchange:
             return None
-        return self._extract_bearer_token(None, raw_headers)
+        return self._extract_subject_token(None, raw_headers, user_api_key_auth)
 
     def _build_stdio_env(
         self,
@@ -3566,6 +3611,7 @@ class MCPServerManager:
         server: MCPServer,
         oauth2_headers: dict[str, str] | None,
         user_api_key_auth: UserAPIKeyAuth | None,
+        raw_headers: Mapping[str, str] | None = None,
     ) -> None:
         """Run the OBO exchange for a caller-supplied subject at the transport edge.
 
@@ -3577,13 +3623,15 @@ class MCPServerManager:
         """
         if server.auth_type != MCPAuth.oauth2_token_exchange:
             return
-        subject_token: Final = self._extract_bearer_token(oauth2_headers, None)
-        if not subject_token:
+        if not self._extract_bearer_token(oauth2_headers, None):
             return
         resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
         spec: Final = to_server_spec(resolved_server)
         if spec is None or not isinstance(spec.config, TokenExchangeConfig):
             return
+        subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
+        if subject_token is None:
+            raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
         match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
             case Ok(_):
                 return
@@ -3851,7 +3899,7 @@ class MCPServerManager:
             # token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes
             # never read the inbound bearer, so leave subject_token None to avoid forwarding it.
             subject_token: Final = (
-                self._extract_bearer_token(oauth2_headers, raw_headers)
+                self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
                 if server.auth_type == MCPAuth.oauth2_token_exchange
                 else None
             )
@@ -3931,6 +3979,7 @@ class MCPServerManager:
     async def get_prompts_from_server(
         self,
         server: MCPServer,
+        user_api_key_auth: UserAPIKeyAuth | None,
         mcp_auth_header: str | dict[str, str] | None = None,
         extra_headers: dict[str, str] | None = None,
         add_prefix: bool = True,
@@ -3959,7 +4008,7 @@ class MCPServerManager:
                 extra_headers.update(server.static_headers)
 
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
-            subject_token: Final = self._obo_subject_token(server, raw_headers)
+            subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
 
             client = await self._create_mcp_client(
                 server=server,
@@ -3982,6 +4031,7 @@ class MCPServerManager:
     async def get_resources_from_server(
         self,
         server: MCPServer,
+        user_api_key_auth: UserAPIKeyAuth | None,
         mcp_auth_header: str | dict[str, str] | None = None,
         extra_headers: dict[str, str] | None = None,
         add_prefix: bool = True,
@@ -4001,7 +4051,7 @@ class MCPServerManager:
                 extra_headers.update(server.static_headers)
 
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
-            subject_token: Final = self._obo_subject_token(server, raw_headers)
+            subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
 
             client = await self._create_mcp_client(
                 server=server,
@@ -4024,6 +4074,7 @@ class MCPServerManager:
     async def get_resource_templates_from_server(
         self,
         server: MCPServer,
+        user_api_key_auth: UserAPIKeyAuth | None,
         mcp_auth_header: str | dict[str, str] | None = None,
         extra_headers: dict[str, str] | None = None,
         add_prefix: bool = True,
@@ -4043,7 +4094,7 @@ class MCPServerManager:
                 extra_headers.update(server.static_headers)
 
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
-            subject_token: Final = self._obo_subject_token(server, raw_headers)
+            subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
 
             client = await self._create_mcp_client(
                 server=server,
@@ -4068,6 +4119,7 @@ class MCPServerManager:
     async def read_resource_from_server(
         self,
         server: MCPServer,
+        user_api_key_auth: UserAPIKeyAuth | None,
         url: AnyUrl,
         mcp_auth_header: str | dict[str, str] | None = None,
         extra_headers: dict[str, str] | None = None,
@@ -4084,7 +4136,7 @@ class MCPServerManager:
             extra_headers.update(server.static_headers)
 
         stdio_env: Final = self._build_stdio_env(server, raw_headers)
-        subject_token: Final = self._obo_subject_token(server, raw_headers)
+        subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
 
         client: Final = await self._create_mcp_client(
             server=server,
@@ -4099,6 +4151,7 @@ class MCPServerManager:
     async def get_prompt_from_server(
         self,
         server: MCPServer,
+        user_api_key_auth: UserAPIKeyAuth | None,
         prompt_name: str,
         arguments: dict[str, str] | None = None,
         mcp_auth_header: str | dict[str, str] | None = None,
@@ -4116,7 +4169,7 @@ class MCPServerManager:
             extra_headers.update(server.static_headers)
 
         stdio_env: Final = self._build_stdio_env(server, raw_headers)
-        subject_token: Final = self._obo_subject_token(server, raw_headers)
+        subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
 
         client: Final = await self._create_mcp_client(
             server=server,
@@ -5290,7 +5343,7 @@ class MCPServerManager:
             MCPAuth.oauth2_token_exchange,
             MCPAuth.oauth2_id_jag,
         ):
-            subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
+            subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
         elif mcp_server.auth_type == MCPAuth.oauth2:
             if mcp_server.has_client_credentials:
                 # For M2M OAuth servers, Authorization must come from token fetch.
@@ -5638,7 +5691,7 @@ class MCPServerManager:
 
         subject_token: str | None = None
         if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
-            subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
+            subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
         elif isinstance(spec.config, PassthroughConfig):
             inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers)
             per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header)

litellm/proxy/_experimental/mcp_server/server.py

@@ -2189,6 +2189,7 @@ if MCP_AVAILABLE:
             try:
                 prompts = await global_mcp_server_manager.get_prompts_from_server(
                     server=server,
+                    user_api_key_auth=user_api_key_auth,
                     mcp_auth_header=server_auth_header,
                     extra_headers=extra_headers,
                     add_prefix=True,  # Always add server prefix
@@ -2242,6 +2243,7 @@ if MCP_AVAILABLE:
             try:
                 resources = await global_mcp_server_manager.get_resources_from_server(
                     server=server,
+                    user_api_key_auth=user_api_key_auth,
                     mcp_auth_header=server_auth_header,
                     extra_headers=extra_headers,
                     add_prefix=True,  # Always add server prefix
@@ -2293,6 +2295,7 @@ if MCP_AVAILABLE:
             try:
                 resource_templates = await global_mcp_server_manager.get_resource_templates_from_server(
                     server=server,
+                    user_api_key_auth=user_api_key_auth,
                     mcp_auth_header=server_auth_header,
                     extra_headers=extra_headers,
                     add_prefix=True,  # Always add server prefix
@@ -3211,6 +3214,7 @@ if MCP_AVAILABLE:
 
         return await global_mcp_server_manager.get_prompt_from_server(
             server=server,
+            user_api_key_auth=user_api_key_auth,
             prompt_name=original_prompt_name,
             arguments=arguments,
             mcp_auth_header=server_auth_header,
@@ -3261,6 +3265,7 @@ if MCP_AVAILABLE:
 
         return await global_mcp_server_manager.read_resource_from_server(
             server=server,
+            user_api_key_auth=user_api_key_auth,
             url=url,
             mcp_auth_header=server_auth_header,
             extra_headers=extra_headers,
@@ -3723,6 +3728,7 @@ if MCP_AVAILABLE:
         user_api_key_auth: UserAPIKeyAuth | None,
         client_ip: str | None,
         allowed_server_ids: set[str] | None = None,
+        raw_headers: Mapping[str, str] | None = None,
     ) -> None:
         """Fail fast with HTTP 401 for MCP servers that need user auth but
         didn't receive it on this request. Covers both gateway-managed OAuth2
@@ -3867,6 +3873,7 @@ if MCP_AVAILABLE:
                     server=server,
                     oauth2_headers=oauth2_headers,
                     user_api_key_auth=user_api_key_auth,
+                    raw_headers=raw_headers,
                 )
 
             # Pass-through OAuth: when the admin has opted a server into
@@ -4195,6 +4202,7 @@ if MCP_AVAILABLE:
                 user_api_key_auth=user_api_key_auth,
                 client_ip=_client_ip,
                 allowed_server_ids=toolset_allowed_server_ids,
+                raw_headers=raw_headers,
             )
 
             # Pre-flight auth check for pass-through servers.  Must run after
@@ -4518,6 +4526,7 @@ if MCP_AVAILABLE:
                 user_api_key_auth=user_api_key_auth,
                 client_ip=_sse_client_ip,
                 allowed_server_ids=toolset_allowed_server_ids,
+                raw_headers=raw_headers,
             )
 
             # Pre-flight auth check for pass-through servers: surface upstream

litellm/proxy/_experimental/mcp_server/tool_search.py

@@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never
+from typing import TYPE_CHECKING, Any, Final, TypedDict
 
 from pydantic import ValidationError
-from typing_extensions import ReadOnly, Required
+from typing_extensions import ReadOnly, Required, assert_never
 
 import litellm
 from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K

litellm/proxy/_types.py

@@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
         None,
         description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.",
     )
+    missing_session_id: Literal["generate", "reject"] | None = Field(
+        None,
+        description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.",
+    )
     enable_public_model_hub: bool = Field(
         default=False,
         description="Public model hub for users to see what models they have access to, supported openai params, etc.",

litellm/proxy/agent_endpoints/endpoints.py

@@ -13,10 +13,10 @@ import os
 import uuid
 from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import Annotated, Final, TypedDict, assert_never
+from typing import Annotated, Final, TypedDict
 
 from fastapi import APIRouter, Depends, HTTPException, Query, Request
-from typing_extensions import ReadOnly, Required
+from typing_extensions import ReadOnly, Required, assert_never
 
 import litellm
 from litellm._logging import verbose_proxy_logger

litellm/proxy/common_utils/reset_budget_job.py

@@ -7,7 +7,9 @@ from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 from types import MappingProxyType
-from typing import Final, Literal, Protocol, TypeVar, assert_never
+from typing import Final, Literal, Protocol, TypeVar
+
+from typing_extensions import assert_never
 
 import litellm
 from litellm._logging import verbose_proxy_logger

litellm/proxy/guardrails/anthropic_sse.py

@@ -13,6 +13,19 @@ from typing import Final
 
 from litellm.types.utils import Choices, ModelResponse
 
+_ANTHROPIC_EVENT_TYPES: Final = frozenset(
+    {
+        "message_start",
+        "message_delta",
+        "message_stop",
+        "content_block_start",
+        "content_block_delta",
+        "content_block_stop",
+        "ping",
+        "error",
+    }
+)
+
 
 def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
     return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
@@ -30,23 +43,43 @@ def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
         return None
 
 
-def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
+def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]:
     from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
         AnthropicPassthroughLoggingHandler,
     )
 
+    return tuple(
+        event_data
+        for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream)  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses
+        if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses; a private import beats forking SSE parsing
+    )
+
+
+def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
     return next(
         (
             message
-            for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream)  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses
-            if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None  # pyright: ignore[reportPrivateUsage]  # same parser the assembler uses; a private import beats forking SSE parsing
-            and event_data.get("type") == "message_start"
-            and isinstance(message := event_data.get("message"), dict)
+            for event_data in _parsed_sse_events(sse_stream)
+            if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict)
         ),
         None,
     )
 
 
+def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool:
+    """Whether raw SSE frames are Anthropic Messages events.
+
+    ``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the
+    only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own
+    stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format
+    its client cannot parse, so the surface is decided on the event types actually present.
+    """
+    sse_stream: Final = _joined_sse_stream(all_chunks)
+    if sse_stream is None:
+        return False
+    return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream))
+
+
 def assemble_anthropic_sse_stream(
     all_chunks: Sequence[object], *, restore_identity: bool = False
 ) -> ModelResponse | None:
@@ -111,6 +144,27 @@ def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
     )
 
 
+def is_sse_error_stream(all_chunks: Sequence[object]) -> bool:
+    """Whether the buffered stream carries nothing but error frames.
+
+    post_call guardrails run in a chain, so a hook can be handed the terminal error frames an
+    earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing
+    them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the
+    Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload.
+    """
+    if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks):
+        # A stream mixing typed chunks with an error frame still carries content to scan, and the
+        # frames-only join below would drop exactly the part that has to be scanned
+        return False
+    sse_stream: Final = _joined_sse_stream(all_chunks)
+    if sse_stream is None:
+        return False
+    events: Final = _parsed_sse_events(sse_stream)
+    return len(events) > 0 and all(
+        event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events
+    )
+
+
 def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
     from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
         LiteLLMAnthropicMessagesAdapter,

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -8,7 +8,7 @@ import json
 import os
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timezone
-from types import UnionType
+from types import MappingProxyType, UnionType
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin
 from urllib.parse import urlparse
 
@@ -51,6 +51,9 @@ from litellm.types.guardrails import (
     SupportedGuardrailIntegrations,
     ToolPermissionGuardrailConfigModel,
 )
+from litellm.types.proxy.guardrails.guardrail_hooks.hide_secrets import (
+    HideSecretsGuardrailConfigModel,
+)
 
 if TYPE_CHECKING:
     from types import CodeType
@@ -1401,7 +1404,11 @@ async def get_guardrail_ui_settings():
         provider: [hook.value for hook in hooks]
         for provider, guardrail_class in guardrail_class_registry.items()
         if (hooks := guardrail_class.get_supported_event_hooks()) is not None
-    }
+    } | MappingProxyType(
+        # hide-secrets lives in the enterprise package, not in the registry
+        # above; it only runs on pre_call.
+        {SupportedGuardrailIntegrations.HIDE_SECRETS.value: [GuardrailEventHooks.pre_call.value]}
+    )
 
     return GuardrailUIAddGuardrailSettings(
         supported_entities=[entity.value for entity in PiiEntityType],
@@ -1953,12 +1960,18 @@ async def get_provider_specific_params():
 
     tool_permission_fields["ui_friendly_name"] = ToolPermissionGuardrailConfigModel.ui_friendly_name()
 
+    # hide-secrets lives in the enterprise package, not in the registry loop below.
+    hide_secrets_fields: Final = _get_fields_from_model(HideSecretsGuardrailConfigModel)
+
+    hide_secrets_fields["ui_friendly_name"] = HideSecretsGuardrailConfigModel.ui_friendly_name()
+
     # Return the provider-specific parameters
     provider_params: Final = {
         SupportedGuardrailIntegrations.BEDROCK.value: bedrock_fields,
         SupportedGuardrailIntegrations.PRESIDIO.value: presidio_fields,
         SupportedGuardrailIntegrations.LAKERA_V2.value: lakera_v2_fields,
         SupportedGuardrailIntegrations.TOOL_PERMISSION.value: tool_permission_fields,
+        SupportedGuardrailIntegrations.HIDE_SECRETS.value: hide_secrets_fields,
     }
 
     ### get the config model for the guardrail - go through the registry and get the config model for the guardrail

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

@@ -1,4 +1,5 @@
 from collections.abc import AsyncGenerator, Mapping, Sequence
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 import httpx
@@ -27,12 +28,25 @@ from litellm.llms.custom_httpx.http_handler import (
 )
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.guardrails.anthropic_sse import (
+    anthropic_sse_chunks_from_response,
+    anthropic_sse_error_frames,
+    assemble_anthropic_sse_stream,
+    is_anthropic_sse_stream,
+    is_raw_sse_stream,
+    is_sse_error_stream,
+)
 from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
     MODEL_ARMOR_MAX_FILE_SIZE_BYTES,
     plan_file_scans,
 )
 from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
-from litellm.types.llms.openai import AllMessageValues
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    ChatCompletionToolCallChunk,
+    ResponsesAPIResponse,
+    ResponsesAPIStreamEvents,
+)
 from litellm.types.utils import (
     CallTypes,
     CallTypesLiteral,
@@ -41,10 +55,33 @@ from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     StandardLoggingGuardrailInformation,
+    TextCompletionResponse,
 )
 
 GUARDRAIL_NAME: Final = "model_armor"
 
+# Only these carry the finished output; response.created carries an empty body
+_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
+
+# Every event whose ``delta`` is model output already on its way to the client. Read off the event
+# enum rather than listed, so an event added there cannot quietly fall out of the scan
+_RESPONSES_DELTA_EVENT_TYPES: Final = frozenset(
+    event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta")
+)
+
+# What makes two delta events part of the same field of the turn, rather than two fields that merely
+# streamed next to each other
+_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index")
+
+
+class _StreamSurface(Enum):
+    """Wire format of a buffered streaming response, which decides how it is read and how it is refused."""
+
+    CHAT_COMPLETIONS = auto()
+    ANTHROPIC_MESSAGES = auto()
+    RESPONSES = auto()
+    OPAQUE_SSE = auto()
+
 
 class ModelArmorAPIError(Exception):
     """Model Armor API failure (non-2xx), distinct from a content-block decision so
@@ -322,19 +359,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         else:
             return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
 
-    def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool:
+    def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool:
         """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
-        sanitization_result: Final = armor_response.get("sanitizationResult", {})
-        filter_results: Final = sanitization_result.get("filterResults", {})
-
-        # filterResults can be a dict (named keys) or a list (array of filter result dicts)
-        filter_result_items = []
-        if isinstance(filter_results, dict):
-            filter_result_items = list(filter_results.values())
-        elif isinstance(filter_results, list):
-            filter_result_items = filter_results
-
-        for filt in filter_result_items:
+        for filt in self._filter_result_items(armor_response):
             # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
             if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND":
                 return True
@@ -358,22 +385,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         # Fallback dict code removed; all cases handled above
         return False
 
-    def _get_sanitized_content(self, armor_response: dict) -> str | None:
+    def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None:
         """
         Get the sanitized content from a Model Armor response, if available.
         Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found.
         """
-        result: Final = armor_response.get("sanitizationResult", {})
-        filter_results: Final = result.get("filterResults", {})
-
-        # filterResults can be a dict (single filter) or a list (multiple filters)
-        filters: Final = (
-            list(filter_results.values())
-            if isinstance(filter_results, dict)
-            else filter_results
-            if isinstance(filter_results, list)
-            else []
-        )
+        filters: Final = self._filter_result_items(armor_response)
 
         # Prefer sanitized text from deidentifyResult if present
         for filter_entry in filters:
@@ -397,6 +414,61 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         # Fallback: if Model Armor put sanitized text at the root, use it
         return armor_response.get("sanitizedText") or armor_response.get("text")
 
+    @staticmethod
+    def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]:
+        """Every filter result in a scan response.
+
+        filterResults is a dict of named filters on most templates and a list on some, so both
+        shapes are flattened to the same list of filter entries.
+        """
+        filter_results: Final = armor_response.get("sanitizationResult", {}).get("filterResults", {})
+        if isinstance(filter_results, dict):
+            return list(filter_results.values())
+        if isinstance(filter_results, list):
+            return filter_results
+        return []
+
+    def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool:
+        """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
+        for filter_entry in self._filter_result_items(armor_response):
+            sdp = filter_entry.get("sdpFilterResult")
+            if sdp and sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND":
+                return True
+        return False
+
+    def _resolve_streaming_outcome(
+        self,
+        armor_response: Mapping[str, Any],
+        assembled_response: object,
+        content: str,
+    ) -> tuple[bool, str | None]:
+        """Whether to block the buffered stream, and the rewrite to emit when it is not blocked.
+
+        A de-identify match only reaches here unblocked because masking is on, so the redaction it
+        stands for has to be both resolvable and emittable. Where it is neither, the buffered
+        original still carries what Model Armor matched on, so this fails closed instead of
+        releasing it.
+        """
+        if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content):
+            return True, None
+        if not self.mask_response_content:
+            return False, None
+
+        sanitized_content: Final = self._get_sanitized_content(armor_response)
+        if not sanitized_content:
+            # No rewrite to apply. Harmless unless a match is outstanding, in which case applying
+            # nothing would hand back the very content that matched
+            return self._has_deidentify_match(armor_response), None
+        if sanitized_content == content:
+            return False, None
+        if not isinstance(assembled_response, ModelResponse):
+            verbose_proxy_logger.warning(
+                "Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, "
+                "blocking the response instead"
+            )
+            return True, None
+        return False, sanitized_content
+
     @staticmethod
     def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
         """Accumulate scan responses so a later text scan does not drop an earlier file scan.
@@ -831,6 +903,185 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
 
         return response
 
+    @staticmethod
+    def _is_terminal_error_stream(all_chunks: Sequence[object]) -> bool:
+        """Whether the buffered stream is only the refusal an earlier guardrail in the chain emitted.
+
+        post_call guardrails are composed, so this hook can be handed the terminal error items a
+        preceding one produced. They carry no message to scan, and replacing them would hide the
+        refusal the client is owed.
+        """
+        if all(getattr(chunk, "type", None) == "error" for chunk in all_chunks):
+            return True
+        return is_sse_error_stream(all_chunks)
+
+    @staticmethod
+    def _classify_stream(all_chunks: Sequence[object]) -> _StreamSurface:
+        """Wire format the buffered chunks belong to."""
+        if is_raw_sse_stream(all_chunks):
+            return (
+                _StreamSurface.ANTHROPIC_MESSAGES if is_anthropic_sse_stream(all_chunks) else _StreamSurface.OPAQUE_SSE
+            )
+        if any(
+            isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.")
+            for chunk in all_chunks
+        ):
+            return _StreamSurface.RESPONSES
+        return _StreamSurface.CHAT_COMPLETIONS
+
+    @staticmethod
+    def _final_responses_api_response(all_chunks: Sequence[object]) -> ResponsesAPIResponse | None:
+        """Response body carried by a terminal ``/v1/responses`` event.
+
+        A stream cut short before it completes has to read as unassembled rather than as a clean
+        empty response: ``response.created`` also carries a body, but an empty one, and scanning
+        that would release every buffered delta unscanned.
+        """
+        return next(
+            (
+                body
+                for chunk in reversed(all_chunks)
+                if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES
+                and isinstance(body := getattr(chunk, "response", None), ResponsesAPIResponse)
+            ),
+            None,
+        )
+
+    @staticmethod
+    def _responses_api_response_text(response: ResponsesAPIResponse) -> str:
+        """Text to scan in a Responses API response, tool-call arguments included.
+
+        Tool calls are folded in because ``get_content_from_model_response`` folds them into what
+        the chat surface scans, and a Responses turn can carry its whole payload in them.
+        """
+        from litellm.llms.openai.responses.guardrail_translation.handler import (
+            OpenAIResponsesHandler,
+        )
+
+        texts: Final[list[str]] = []  # mutable-ok: the shared extractor below appends into caller-owned lists
+        tool_calls: Final[list[ChatCompletionToolCallChunk]] = []  # mutable-ok: the same extractor's tool-call sink
+        handler: Final = OpenAIResponsesHandler()
+        for output_idx, output_item in enumerate(response.output or ()):
+            handler._extract_output_text_and_images(  # pyright: ignore[reportPrivateUsage]  # the shared Responses output extractor; forking it would duplicate per-item parsing
+                output_item=output_item,
+                output_idx=output_idx,
+                texts_to_check=texts,
+                images_to_check=[],  # mutable-ok: the extractor's images sink, unused here
+                task_mappings=[],  # mutable-ok: the extractor's task-mapping sink, unused here
+                tool_calls_to_check=tool_calls,
+            )
+        return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls)))
+
+    def _extract_streaming_content(self, assembled_response: object) -> str:
+        """Text to scan from an assembled stream, for every endpoint shape this hook serves."""
+        if isinstance(assembled_response, ResponsesAPIResponse):
+            return self._responses_api_response_text(assembled_response)
+        return self._extract_content_from_response(assembled_response)
+
+    @staticmethod
+    def _responses_delta_field(chunk: object) -> tuple[str, ...]:
+        """Which field of the turn a delta event belongs to."""
+        return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
+
+    @staticmethod
+    def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
+        """Text each field of a ``/v1/responses`` turn has already spelled out in its delta events.
+
+        One field's deltas are joined as they streamed, since a finding can be split across them,
+        and separate fields stay apart, so a reasoning summary running into the visible answer
+        cannot spell out a finding that neither of them carries.
+        """
+        deltas: Final = tuple(
+            (ModelArmorGuardrail._responses_delta_field(chunk), delta)
+            for chunk in all_chunks
+            if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
+            and isinstance(delta := getattr(chunk, "delta", None), str)
+        )
+        return tuple(
+            "".join(delta for field, delta in deltas if field == streamed_field)
+            for streamed_field in dict.fromkeys(field for field, _ in deltas)
+        )
+
+    def _streaming_content_to_scan(
+        self,
+        assembled_response: object,
+        all_chunks: Sequence[object],
+        surface: _StreamSurface,
+    ) -> str:
+        """Text to scan for a buffered stream, which is everything the client is about to receive.
+
+        A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in
+        delta events that its terminal body never repeats, so every delta field the body does not
+        already carry is scanned after it.
+        """
+        content: Final = self._extract_streaming_content(assembled_response)
+        if surface is not _StreamSurface.RESPONSES:
+            return content
+        unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
+        return "\n".join(part for part in (content, *unscanned) if part)
+
+    @staticmethod
+    def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:
+        """Replace every non-empty choice message with the Model Armor sanitized text."""
+        for choice in assembled_response.choices:
+            if isinstance(choice, Choices) and choice.message.content:
+                choice.message.content = sanitized_content
+
+    @staticmethod
+    def _assemble_chat_completion_stream(
+        all_chunks: list[object],  # mutable-ok: stream_chunk_builder only accepts a mutable list
+    ) -> ModelResponse | TextCompletionResponse | None:
+        """Assemble chat-completion chunks, returning ``None`` when they cannot be assembled."""
+        from litellm.main import stream_chunk_builder
+
+        try:
+            return stream_chunk_builder(chunks=all_chunks)
+        except Exception as exc:
+            verbose_proxy_logger.warning("Model Armor: chat-completion stream assembly failed (%s)", exc)
+            return None
+
+    def _assemble_stream(
+        self, all_chunks: Sequence[object], surface: _StreamSurface
+    ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None:
+        """Assemble the buffered stream into the scannable response its surface produces."""
+        if surface is _StreamSurface.ANTHROPIC_MESSAGES:
+            return assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
+        if surface is _StreamSurface.RESPONSES:
+            return self._final_responses_api_response(all_chunks)
+        if surface is _StreamSurface.OPAQUE_SSE:
+            return None
+        return self._assemble_chat_completion_stream(list(all_chunks))
+
+    @staticmethod
+    def _error_payload(exc: HTTPException) -> Mapping[str, object]:
+        """Error object for a terminal stream item, carrying the status the frame would otherwise lose."""
+        detail: Final = exc.detail if isinstance(exc.detail, Mapping) else {"message": str(exc.detail)}
+        error_value: Final = detail.get("error", detail)
+        return {
+            **(dict(error_value) if isinstance(error_value, Mapping) else {"message": str(error_value)}),
+            "code": str(exc.status_code),
+        }
+
+    @staticmethod
+    def _build_responses_error_items(exc: HTTPException) -> Sequence[object] | None:
+        """Responses API error events for a failure discovered after the stream started."""
+        from litellm.llms.openai.responses.guardrail_translation.handler import (
+            OpenAIResponsesHandler,
+        )
+
+        return OpenAIResponsesHandler().build_stream_error_items(exc, responses_so_far=None)
+
+    def _stream_error_items(self, exc: HTTPException, *, surface: _StreamSurface) -> Sequence[object]:
+        """Frame a guardrail failure as terminal stream items in this endpoint's wire format."""
+        payload: Final = self._error_payload(exc)
+        if surface is _StreamSurface.ANTHROPIC_MESSAGES:
+            return anthropic_sse_error_frames(str(payload.get("message", "")))
+        if surface is _StreamSurface.RESPONSES and (responses_items := self._build_responses_error_items(exc)):
+            return responses_items
+        # Also the fallback when a surface cannot frame its own error: create_response() reads the
+        # status back out of this form, so the refusal keeps its code instead of arriving as a 200
+        return (f"data: {json.dumps({'error': payload})}\n\n",)
+
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -840,97 +1091,125 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         """Process streaming response chunks."""
 
         from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
-        from litellm.main import stream_chunk_builder
+        from litellm.proxy.common_utils.callback_utils import (
+            add_guardrail_to_applied_guardrails_header,
+        )
 
         # Collect all chunks
-        all_chunks: Final[list[ModelResponseStream]] = []
+        all_chunks: Final[list[Any]] = []
         async for chunk in response:
             all_chunks.append(chunk)
 
+        if not all_chunks or self._is_terminal_error_stream(all_chunks):
+            for chunk in all_chunks:
+                yield chunk
+            return
+
+        surface: Final = self._classify_stream(all_chunks)
+
         # Build complete response
-        assembled_response: Final = stream_chunk_builder(chunks=all_chunks)
+        assembled_response: Final = self._assemble_stream(all_chunks, surface)
 
-        if isinstance(assembled_response, ModelResponse):
-            # Extract content
-            content: Final = self._extract_content_from_response(assembled_response)
+        if assembled_response is None:
+            if not self.optional_params.get("fail_on_error", True):
+                verbose_proxy_logger.warning(
+                    "Model Armor: streamed response could not be assembled for scanning, "
+                    "forwarding it unscanned because fail_on_error is disabled"
+                )
+                for chunk in all_chunks:
+                    yield chunk
+                return
 
-            if content:
-                try:
-                    # Check with Model Armor
-                    armor_response: Final = await self.make_model_armor_request(
-                        content=content,
-                        source="model_response",
-                        request_data=request_data,
-                    )
+            # Forwarding an unscannable stream would silently disable the guardrail, so fail closed
+            add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
+            for error_item in self._stream_error_items(
+                HTTPException(
+                    status_code=500,
+                    detail=f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it",
+                ),
+                surface=surface,
+            ):
+                yield error_item
+            return
 
-                    # Attach Model Armor response & status to this request's metadata to avoid race conditions
-                    if isinstance(request_data, dict):
-                        _, metadata = get_or_create_metadata_bucket(request_data)
-                        metadata["_model_armor_response"] = self._build_logging_response(armor_response)
-                        metadata["_model_armor_status"] = (
-                            "blocked" if self._should_block_content(armor_response) else "success"
-                        )
+        # Extract content
+        content: Final = self._streaming_content_to_scan(
+            assembled_response=assembled_response, all_chunks=all_chunks, surface=surface
+        )
 
-                    # Add guardrail to applied_guardrails BEFORE potential blocking
-                    # This ensures guardrail is recorded even when it blocks the request
-                    from litellm.proxy.common_utils.callback_utils import (
-                        add_guardrail_to_applied_guardrails_header,
-                    )
+        if not content:
+            verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail")
+            for chunk in all_chunks:
+                yield chunk
+            return
 
-                    add_guardrail_to_applied_guardrails_header(
-                        request_data=request_data, guardrail_name=self.guardrail_name
-                    )
+        try:
+            # Check with Model Armor
+            armor_response: Final = await self.make_model_armor_request(
+                content=content,
+                source="model_response",
+                request_data=request_data,
+            )
 
-                    # Check if blocked
-                    if self._should_block_content(armor_response):
-                        raise HTTPException(
-                            status_code=400,
-                            detail=self._build_block_error_detail(
-                                "Streaming response blocked by Model Armor",
-                                armor_response,
-                            ),
-                        )
+            # Decide the outcome before recording it. Mirrors the non-streaming sibling: with
+            # masking on, a de-identify match is a redaction to apply rather than a refusal, but
+            # that only holds while the redaction can actually be delivered
+            blocked, sanitized_content = self._resolve_streaming_outcome(
+                armor_response=armor_response,
+                assembled_response=assembled_response,
+                content=content,
+            )
 
-                    # Apply sanitization if enabled
-                    if self.mask_response_content:
-                        sanitized_content: Final = self._get_sanitized_content(armor_response)
-                        if sanitized_content and sanitized_content != content:
-                            # Update assembled response
-                            for choice in assembled_response.choices:
-                                if isinstance(choice, Choices):
-                                    if choice.message.content:
-                                        choice.message.content = sanitized_content
+            # Attach Model Armor response & status to this request's metadata to avoid race conditions
+            if isinstance(request_data, dict):
+                _, metadata = get_or_create_metadata_bucket(request_data)
+                metadata["_model_armor_response"] = self._build_logging_response(armor_response)
+                metadata["_model_armor_status"] = "blocked" if blocked else "success"
 
-                            # Return sanitized stream
-                            mock_response: Final = MockResponseIterator(model_response=assembled_response)
-                            async for chunk in mock_response:
-                                yield chunk
-                            return
+            # Add guardrail to applied_guardrails BEFORE potential blocking
+            # This ensures guardrail is recorded even when it blocks the request
+            add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
 
-                except ModelArmorAPIError as e:
-                    if self.optional_params.get("fail_on_error", True):
-                        error_obj = {"message": e.detail, "code": "500"}
-                        yield f"data: {json.dumps({'error': error_obj})}\n\n"
-                        return
-                except HTTPException as e:
-                    # Yield error as SSE event so create_response() detects it and
-                    # returns a proper JSON error response with the correct status code.
-                    # (Raising from a generator hits create_response's generic except → 500.)
-                    detail: Final = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)}
-                    error_value: Final = detail.get("error", detail)
-                    if isinstance(error_value, dict):
-                        error_obj = dict(error_value)
-                    else:
-                        error_obj = {"message": str(error_value)}
-                    error_obj["code"] = str(e.status_code)
-                    yield f"data: {json.dumps({'error': error_obj})}\n\n"
+            if blocked:
+                raise HTTPException(
+                    status_code=400,
+                    detail=self._build_block_error_detail(
+                        "Streaming response blocked by Model Armor",
+                        armor_response,
+                    ),
+                )
+
+            if sanitized_content is not None and isinstance(assembled_response, ModelResponse):
+                self._apply_sanitized_content(assembled_response, sanitized_content)
+
+                # Return sanitized stream
+                if surface is _StreamSurface.ANTHROPIC_MESSAGES:
+                    for sse_chunk in anthropic_sse_chunks_from_response(assembled_response):
+                        yield sse_chunk
                     return
-                except Exception as e:
-                    verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True)
-                    if self.optional_params.get("fail_on_error", True):
-                        raise
-            else:
-                verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail")
+                mock_response: Final = MockResponseIterator(model_response=assembled_response)
+                async for chunk in mock_response:
+                    yield chunk
+                return
+
+        except ModelArmorAPIError as e:
+            if self.optional_params.get("fail_on_error", True):
+                for error_item in self._stream_error_items(
+                    HTTPException(status_code=500, detail=e.detail), surface=surface
+                ):
+                    yield error_item
+                return
+        except HTTPException as e:
+            # Yield the error as a terminal stream item so create_response() detects it and returns
+            # a proper JSON error response with the correct status code. Raising from a generator
+            # instead hits create_response's generic except and becomes a 500.
+            for error_item in self._stream_error_items(e, surface=surface):
+                yield error_item
+            return
+        except Exception as e:
+            verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True)
+            if self.optional_params.get("fail_on_error", True):
+                raise
 
         # Return original chunks if no sanitization needed
         for chunk in all_chunks:

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -36,6 +36,7 @@ if TYPE_CHECKING:
     from litellm.integrations.custom_guardrail import ModifyResponseException
     from litellm.llms.base_llm.guardrail_translation.base_translation import (
         BaseTranslation,
+        StreamingScanKey,
     )
 
 # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error
@@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol):
     @property
     def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ...
 
+    @property
+    def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ...
+
     @property
     def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
 
@@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]:
     return choices
 
 
+def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
+    if scan_key is None:
+        return False
+    return scan_key == last_scan_key or scan_key.has_nothing_to_scan
+
+
 class _StreamTerminated(Exception):
     """Internal signal that the incremental transform stream has already emitted
     its terminal chunks (block message or in-stream error) and must stop."""
@@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Drives how a block terminates the stream: continue the in-progress
         # message (True) vs emit a standalone block message (False, buffered).
         chunks_yielded = False
+        last_scan_key: StreamingScanKey | None = None  # rebind-ok: replaced after every scan round
 
         async for item in response:
             chunk_counter += 1
@@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger):
 
             # Process chunk based on sampling rate
             if chunk_counter % sampling_rate == 0:
+                endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+                scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
+                if _is_redundant_scan(scan_key, last_scan_key):
+                    verbose_proxy_logger.debug(
+                        "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
+                        chunk_counter,
+                        guardrail_to_apply.guardrail_name,
+                    )
+                    chunks_yielded = True
+                    responses_yielded.append(item)
+                    yield item
+                    continue
+
                 verbose_proxy_logger.debug(
                     "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
                     chunk_counter,
@@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger):
                 # string, permanently losing this chunk's content.
                 original_item = copy.deepcopy(item)
 
-                endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
-
                 try:
                     await endpoint_translation.process_output_streaming_response(
                         responses_so_far=responses_so_far,
@@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger):
                     ):
                         yield error_item
                     return
+                if scan_key is not None:
+                    last_scan_key = scan_key
                 chunks_yielded = True
                 responses_yielded.append(original_item)
                 yield original_item
@@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger):
             # preserve the list, not clone every chunk (deepcopy would double
             # peak memory for large responses).
             buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None
+            end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
+            if _is_redundant_scan(end_scan_key, last_scan_key):
+                verbose_proxy_logger.debug(
+                    "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
+                    guardrail_to_apply.guardrail_name,
+                )
+                for buffered_item in buffered_items or ():
+                    yield buffered_item
+                for pending_item in pending_end_of_stream_items:
+                    responses_yielded.append(pending_item)
+                    yield pending_item
+                return
 
             try:
                 await endpoint_translation.process_output_streaming_response(

litellm/proxy/litellm_pre_call_utils.py

@@ -16,6 +16,7 @@ from starlette.datastructures import Headers
 import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm._service_logger import ServiceLogging
+from litellm._uuid import uuid
 from litellm.constants import (
     CONSUMED_REQUEST_TAGS_METADATA_KEY,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
@@ -23,6 +24,7 @@ from litellm.constants import (
     OTEL_SERVICE_NAME_METADATA_KEYS,
     PRE_CALL_EXECUTED_GUARDRAILS_KEY,
     SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
+    SESSION_ID_GENERATED_METADATA_KEY,
 )
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@@ -40,6 +42,7 @@ from litellm.proxy._types import (
     AddTeamCallback,
     CommonProxyErrors,
     LitellmDataForBackendLLMCall,
+    LiteLLMRoutes,
     LitellmUserRoles,
     ProxyErrorTypes,
     ProxyException,
@@ -47,6 +50,8 @@ from litellm.proxy._types import (
     TeamCallbackMetadata,
     UserAPIKeyAuth,
 )
+from litellm.proxy.auth.auth_utils import get_request_route
+from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.common_utils.callback_utils import (
     decrypt_callback_vars,
     get_metadata_variable_name_from_kwargs,
@@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
     return session_id
 
 
+def _is_llm_inference_route(request: Request) -> bool:
+    route: Final = get_request_route(request)
+    return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access(
+        route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
+    )
+
+
+def apply_missing_session_id_policy(
+    data: dict[str, object],  # mutable-ok: stamps session ids in place on the request body the pipeline threads through
+    _metadata_variable_name: str,
+    general_settings: Mapping[str, object] | None,
+    request: Request,
+) -> None:
+    policy: Final = general_settings.get("missing_session_id") if general_settings else None
+    if policy is None or not _is_llm_inference_route(request):
+        return
+    metadata: Final = data.get(_metadata_variable_name)
+    if not isinstance(metadata, dict):
+        return
+    if data.get("litellm_session_id") or metadata.get("session_id"):
+        return
+    match policy:
+        case "generate":
+            session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4())
+            data["litellm_session_id"] = session_id  # rebind-ok: data is an out-param
+            data.setdefault("litellm_trace_id", session_id)
+            metadata["session_id"] = session_id
+            metadata[SESSION_ID_GENERATED_METADATA_KEY] = True
+        case "reject":
+            raise ProxyException(
+                message=(
+                    "Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. "
+                    "Required by `general_settings.missing_session_id: reject`."
+                ),
+                type=ProxyErrorTypes.bad_request_error,
+                param="session_id",
+                code=400,
+            )
+        case _:
+            verbose_proxy_logger.warning(
+                "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy
+            )
+
+
 def is_claude_code_user_agent(user_agent: str) -> bool:
     """Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
     extensions and the Agent SDK run through the same CLI and share that prefix."""
@@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request(
         data=data,
         _metadata_variable_name=_metadata_variable_name,
     )
+    apply_missing_session_id_policy(
+        data=data,
+        _metadata_variable_name=_metadata_variable_name,
+        general_settings=general_settings,
+        request=request,
+    )
 
     # Expose request headers under the metadata field for guardrails (fixes #17477)
     if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict):

litellm/proxy/management_endpoints/sso/saml_sso.py

@@ -443,7 +443,9 @@ class SAMLAuthHandler:
         last_name: Final = SAMLAuthHandler._attribute_value(
             attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES
         )
-        role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES)
+        role_values: Final = SAMLAuthHandler._attribute_values(
+            attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES
+        )
         team_ids: Final = SAMLAuthHandler._attribute_values(
             attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES
         )
@@ -464,7 +466,7 @@ class SAMLAuthHandler:
                 picture=None,
                 provider="saml",
                 team_ids=team_ids,
-                user_role=get_litellm_user_role(role_value) if role_value else None,
+                user_role=get_litellm_user_role(role_values),
             )
         except ValidationError as e:
             raise HTTPException(

litellm/proxy/management_endpoints/types.py

@@ -4,12 +4,44 @@ Types for the management endpoints
 Might include fastapi/proxy requirements.txt related imports
 """
 
+from collections.abc import Iterable, Sequence
 from typing import Any, Final, cast
 
 from fastapi_sso.sso.base import OpenID
 
 from litellm.proxy._types import LitellmUserRoles
 
+# Ordered highest to lowest privilege
+LITELLM_USER_ROLE_HIERARCHY: Final = (
+    LitellmUserRoles.PROXY_ADMIN,
+    LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
+    LitellmUserRoles.INTERNAL_USER,
+    LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
+)
+
+
+def highest_privilege_role(roles: Iterable[LitellmUserRoles]) -> LitellmUserRoles | None:
+    """
+    Pick the highest privilege role out of the roles an IdP asserted for one user.
+
+    IdPs do not guarantee ordering within a multi-valued role claim, so a user holding
+    several roles resolves to the most privileged one rather than whichever came first.
+    Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay
+    deterministic.
+
+    Args:
+        roles: The roles resolved from the claim
+
+    Returns:
+        The highest privilege role, or None if `roles` is empty
+    """
+    resolved: Final = frozenset(roles)
+    if not resolved:
+        return None
+
+    ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None)
+    return ranked if ranked is not None else min(resolved, key=lambda role: role.value)
+
 
 def is_valid_litellm_user_role(role_str: str) -> bool:
     """
@@ -28,12 +60,22 @@ def is_valid_litellm_user_role(role_str: str) -> bool:
         return False
 
 
-def get_litellm_user_role(role_str) -> LitellmUserRoles | None:
+def _role_from_claim_value(role_str: object) -> LitellmUserRoles | None:
+    if not isinstance(role_str, str):
+        return None
+    # Use _value2member_map_ for O(1) lookup, case-insensitive
+    result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower())
+    return cast(LitellmUserRoles | None, result)
+
+
+def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None:
     """
     Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive).
 
     Handles list inputs since some SSO providers (e.g., Keycloak) return roles
-    as arrays like ["proxy_admin"] instead of plain strings.
+    as arrays like ["proxy_admin"] instead of plain strings. A claim carrying several
+    roles resolves to the highest privilege one, so a user does not lose access just
+    because the IdP listed a weaker role first.
 
     Args:
         role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"])
@@ -41,16 +83,12 @@ def get_litellm_user_role(role_str) -> LitellmUserRoles | None:
     Returns:
         LitellmUserRoles enum if valid, None otherwise
     """
-    try:
-        if isinstance(role_str, list):
-            if len(role_str) == 0:
-                return None
-            role_str = role_str[0]
-        # Use _value2member_map_ for O(1) lookup, case-insensitive
-        result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower())
-        return cast(LitellmUserRoles | None, result)
-    except Exception:
-        return None
+    if isinstance(role_str, (list, tuple)):
+        entries: Final = cast(Sequence[object], role_str)  # cast-ok: isinstance narrows the claim, not its elements
+        return highest_privilege_role(
+            role for role in (_role_from_claim_value(entry) for entry in entries) if role is not None
+        )
+    return _role_from_claim_value(role_str)
 
 
 class CustomOpenID(OpenID):

litellm/proxy/management_endpoints/ui_sso.py

@@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import (
 )
 from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add
 from litellm.proxy.management_endpoints.types import (
+    LITELLM_USER_ROLE_HIERARCHY,
     CustomOpenID,
     get_litellm_user_role,
     is_valid_litellm_user_role,
@@ -809,15 +810,6 @@ def normalize_email(email: str | None) -> str | None:
     return email.lower() if isinstance(email, str) else email
 
 
-# Ordered highest to lowest privilege
-LITELLM_USER_ROLE_HIERARCHY: Final = (
-    LitellmUserRoles.PROXY_ADMIN,
-    LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
-    LitellmUserRoles.INTERNAL_USER,
-    LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
-)
-
-
 def determine_role_from_groups(
     user_groups: list[str],
     role_mappings: "RoleMappings",
@@ -4312,14 +4304,7 @@ class MicrosoftSSOHandler:
         listed first. Roles the hierarchy does not rank (org_admin, team, customer)
         resolve by name to stay deterministic
         """
-        resolved: Final = frozenset(
-            role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None
-        )
-        if not resolved:
-            return None
-
-        ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None)
-        return ranked if ranked is not None else min(resolved, key=lambda role: role.value)
+        return get_litellm_user_role(tuple(app_roles or ()))
 
     @staticmethod
     def get_app_roles_from_id_token(id_token: str | None) -> list[str]:

litellm/proxy/openai_files_endpoints/batch_file_validation.py

@@ -2,7 +2,9 @@ import json
 from collections.abc import Iterator
 from dataclasses import dataclass
 from itertools import chain
-from typing import BinaryIO, Final, NoReturn, assert_never
+from typing import BinaryIO, Final, NoReturn
+
+from typing_extensions import assert_never
 
 from litellm.proxy._types import ProxyException
 

litellm/proxy/openai_files_endpoints/general_upload_validation.py

@@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose.
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import BinaryIO, Final, NoReturn, assert_never
+from typing import BinaryIO, Final, NoReturn
+
+from typing_extensions import assert_never
 
 from litellm.proxy._types import ProxyException
 from litellm.proxy.common_utils.path_utils import safe_filename

litellm/proxy/public_endpoints/autorouter_presets.json

@@ -1,4 +1,26 @@
 {
+  "1m_context": {
+    "label": "1M Context",
+    "description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Opus 5 for complex, Opus 5 at high thinking for reasoning.",
+    "complexity_router_config": {
+      "tiers": {
+        "SIMPLE": ["gpt-5.6-luna"],
+        "MEDIUM": ["gpt-5.6-terra"],
+        "COMPLEX": ["claude-opus-5"],
+        "REASONING": ["claude-opus-5"]
+      },
+      "tier_model_configs": {
+        "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }]
+      },
+      "classifier_type": "heuristic_v2",
+      "escalation_keywords": ["LITELLM ESCALATE"],
+      "classification_mode": "every_request",
+      "session_affinity": false,
+      "modality_routing": false,
+      "modality_pin_override": false,
+      "deployment_affinity": true
+    }
+  },
   "anthropic_family": {
     "label": "Anthropic Family",
     "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.",
@@ -17,6 +39,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -35,6 +58,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -63,6 +87,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -84,6 +109,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   }

litellm/router.py

@@ -150,8 +150,10 @@ from litellm.router_utils.fallback_event_handlers import (
     _check_non_standard_fallback_format,
     clear_pre_routing_selection,
     fallback_lookup_groups,
+    fallbacks_disabled_for_request,
     get_fallback_model_group_for_lookup_groups,
     get_pre_routing_selection,
+    record_disable_fallbacks,
     record_pre_routing_selection,
     run_async_fallback,
 )
@@ -5193,7 +5195,7 @@ class Router:
                         if not has_generated_content and error_event is None
                         else None
                     )
-                    if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs):
+                    if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs):
                         refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
                         raise MidStreamFallbackError(
                             message=refusal_error.message,
@@ -7266,6 +7268,7 @@ class Router:
                     _fallback_metadata["original_model_group"] = model_group
         include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True
         disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False)
+        record_disable_fallbacks(kwargs, disable_fallbacks is True)
         fallbacks: Final[list | None] = kwargs.get("fallbacks", self.fallbacks)
         context_window_fallbacks: list | None = kwargs.get("context_window_fallbacks", self.context_window_fallbacks)
         content_policy_fallbacks: list | None = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
@@ -8131,6 +8134,29 @@ class Router:
         )
         return False
 
+    def _refusal_fallback_available(self, model_group: str, kwargs: Mapping[str, Any]) -> bool:
+        """
+        Whether a safeguard refusal can actually be recovered by the dispatcher. A configured
+        content-policy list is authoritative; with none configured at all, the dispatcher falls
+        through to the generic fallbacks lookup, so the gate mirrors that reachability and arms
+        on a resolving generic chain (tier first, then the requested group, then "*").
+        """
+        if fallbacks_disabled_for_request(kwargs):
+            return False
+        content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks)
+        if content_policy_fallbacks is not None:
+            return self._has_content_policy_fallback(model_group, kwargs)
+        if self._has_default_fallbacks():
+            return True
+        fallbacks: Final = kwargs.get("fallbacks", self.fallbacks)
+        if fallbacks is None:
+            return False
+        resolved, _ = get_fallback_model_group_for_lookup_groups(
+            fallbacks=fallbacks,
+            lookup_groups=fallback_lookup_groups(kwargs, model_group),
+        )
+        return resolved is not None
+
     def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool:
         """
         Determines if a content policy error should be raised.
@@ -8162,7 +8188,7 @@ class Router:
             return False
         if get_safeguard_refusal_stop_details(response) is None:
             return False
-        return self._has_content_policy_fallback(model, kwargs)
+        return self._refusal_fallback_available(model, kwargs)
 
     def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None):
         _all_deployments: list = []

litellm/router_strategy/complexity_router/README.md

@@ -187,6 +187,9 @@ model_list:
 
         # Replace a routed model that cannot take image input (default: false)
         modality_routing: true
+
+        # Let that replacement also override a kept session pin, for image turns only (default: false)
+        modality_pin_override: true
 ```
 
 ## Usage
@@ -227,9 +230,16 @@ vision model sits below the decided tier gets the 400 and an actionable message
 A same-tier re-pick keeps the decision's cause and adds `modality:image` to `signals`; a tier
 change or default takeover records `cause: modality_escalation` with the displaced placement
 (`modality_escalated_from:<TIER>` or `modality_displaced_default_model`). Escalations are never
-pinned by session affinity, and a KEPT session pin bypasses the gate entirely: a session pinned
+pinned by session affinity, and by default a KEPT session pin bypasses the gate: a session pinned
 to a text-only model keeps it even when an image arrives.
 
+Add `modality_pin_override: true` to lift that last exemption. The image turn is then re-placed
+the same way every other decision is, and records `cause: modality_pin_override` whether or not
+the tier moved, since the model left the pin either way. The pin itself is untouched: the session
+affinity write happens upstream of the gate and stores the session's own model, so the next text
+turn replays the original pin and the override is never pinned in its place. It does nothing
+unless `modality_routing` is also on.
+
 ### Heuristic-first chaining
 
 `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM

litellm/router_strategy/complexity_router/complexity_router.py

@@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
 from pydantic import BaseModel, create_model
 
 from litellm._logging import verbose_router_logger
-from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY
+from litellm.constants import (
+    EMPTY_MAPPING,
+    RETURN_RAW_MODEL_NAME_METADATA_KEY,
+    SESSION_ID_GENERATED_METADATA_KEY,
+)
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
 from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
@@ -747,7 +751,8 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
 
     A modality escalation is transient the same way: it describes what this one call carries (an
     image), not what the session's traffic looks like, and pinning it would hold every following
-    text turn on the vision-capable model the image forced.
+    text turn on the vision-capable model the image forced. A modality pin override is the same
+    fact on a session that already holds a pin, so it must not overwrite the pin it displaced.
     """
     return decision is None or (
         decision.get("cause")
@@ -756,6 +761,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
             "plan_mode",
             "housekeeping",
             "modality_escalation",
+            "modality_pin_override",
         )
         and not decision.get("context_escalated")
     )
@@ -2389,8 +2395,11 @@ class ComplexityRouter(CustomLogger):
         """Replace a routed model that cannot accept this request's image input.
 
         The single modality owner, applied to the decided response at the hook's exits so every
-        routing path is covered uniformly. A KEPT session pin is exempt by design (its cause);
-        replacement picks and every other path are just responses. The re-placement walks
+        routing path is covered uniformly. A KEPT session pin is exempt by design (its cause)
+        unless modality_pin_override is set, in which case the image turn is re-placed and reported
+        as modality_pin_override while the stored pin, written upstream from the session's own
+        model, is left for the next text turn; replacement picks and every other path are just
+        responses. The re-placement walks
         UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks
         through `_pick_model_for_tier` so routing plugins still apply, then falls to
         default_model (never on plugin routers, and never on a plan-floored decision, since
@@ -2403,7 +2412,11 @@ class ComplexityRouter(CustomLogger):
             not self.config.modality_routing
             or not resolved_messages
             or response.model is None
-            or (decision is not None and decision.get("cause") == "session_affinity_pin")
+            or (
+                decision is not None
+                and decision.get("cause") == "session_affinity_pin"
+                and not self.config.modality_pin_override
+            )
             or not request_contains_image_content(resolved_messages)
             or self._model_accepts_image_input(response.model)
         ):
@@ -2445,6 +2458,10 @@ class ComplexityRouter(CustomLogger):
         self._restamp_adaptive_choice(request_kwargs, response.model, new_model)
         same_tier: Final = capable is not None and decided == capable
         base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback"
+        # Reaching here on a kept pin means modality_pin_override is on, since the guard above
+        # returns otherwise. The model moved off the pin even on a same-tier repick, so reporting
+        # the pin's own cause would claim the session's model served a request it did not.
+        displaced_pin: Final = base_cause == "session_affinity_pin"
         displaced_default: Final = decided is None and response.model == self.config.default_model
         markers: Final = (
             "modality:image",
@@ -2454,7 +2471,7 @@ class ComplexityRouter(CustomLogger):
         old_signals: Final = tuple(decision.get("signals") or ()) if decision is not None else ()
         new_decision: Final = self._build_routing_decision(
             routed_model=new_model,
-            cause=base_cause if same_tier else "modality_escalation",
+            cause="modality_pin_override" if displaced_pin else (base_cause if same_tier else "modality_escalation"),
             tier=new_tier,
             score=decision.get("score") if decision is not None else None,
             signals=(*old_signals, *markers),
@@ -2712,7 +2729,7 @@ class ComplexityRouter(CustomLogger):
         """Resolve a client-supplied session_id."""
         for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
             session_id = metadata.get("session_id")
-            if session_id is not None:
+            if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
                 return str(session_id)
         return None
 

litellm/router_strategy/complexity_router/config.py

@@ -883,7 +883,20 @@ class ComplexityRouterConfig(BaseModel):
             "a routed model explicitly declared supports_vision false (deployment model_info "
             "or the model cost map; unmapped names stay routable) is replaced by the nearest "
             "HIGHER tier holding a capable model, then default_model, else a clear 400. A kept "
-            "session-affinity pin still wins even when an image arrives."
+            "session-affinity pin still wins even when an image arrives, unless "
+            "modality_pin_override is also enabled."
+        ),
+    )
+    modality_pin_override: bool = Field(
+        default=False,
+        description=(
+            "Let modality_routing replace a kept session-affinity pin on the turns that carry an "
+            "image. Without this, a session pinned to a text-only model fails every image turn with "
+            "a provider 400, since the pin is exempt from the modality gate. When enabled, such a "
+            "turn routes to a capable model for that request only and the stored pin is left "
+            "untouched, so the next text turn replays the session's own model; the override is "
+            "reported as cause modality_pin_override and is never itself pinned. Inert unless "
+            "modality_routing is also enabled."
         ),
     )
 

litellm/router_utils/fallback_event_handlers.py

@@ -263,6 +263,38 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None:
     return next((selected for selected in selections if isinstance(selected, str) and selected), None)
 
 
+DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks"
+
+
+def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None:
+    """
+    Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata
+    bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal
+    gate (which decides whether to convert a refusal into a recoverable error) needs this
+    carrier to know recovery is impossible.
+    """
+    from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
+
+    if request_kwargs is None:
+        return
+    bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))
+    if not isinstance(bucket, dict):
+        return
+    if disabled:
+        bucket[DISABLE_FALLBACKS_METADATA_KEY] = True
+    else:
+        bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None)
+
+
+def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool:
+    """True when this request opted out of fallbacks, read from the raw kwarg (pre-pop
+    snapshots keep it) or the router-internal bucket the wrapper stamps after popping it."""
+    if kwargs.get("disable_fallbacks") is True:
+        return True
+    buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS)
+    return any(isinstance(bucket, dict) and bucket.get(DISABLE_FALLBACKS_METADATA_KEY) is True for bucket in buckets)
+
+
 def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]:
     """
     Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins,

litellm/router_utils/pre_call_checks/deployment_affinity_check.py

@@ -21,7 +21,7 @@ from typing_extensions import TypedDict
 
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
-from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
+from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.integrations.custom_logger import CustomLogger, Span
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import AllMessageValues
@@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger):
     @staticmethod
     def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
         session_id: Final = metadata.get("session_id")
-        if session_id is None:
+        if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
             return None
         return str(session_id)
 

litellm/types/integrations/datadog_llm_obs.py

@@ -86,6 +86,7 @@ class LLMMetrics(TypedDict, total=False):
     cache_read_input_tokens: ReadOnly[float]
     cache_write_input_tokens: ReadOnly[float]
     non_cached_input_tokens: ReadOnly[float]
+    reasoning_output_tokens: ReadOnly[float]
 
 
 class LLMObsPayload(TypedDict, total=False):

litellm/types/llms/gemini_audio_transcription.py

@@ -1,7 +1,7 @@
-from typing import Literal, Required
+from typing import Literal
 
 from pydantic import BaseModel, ConfigDict
-from typing_extensions import ReadOnly, TypedDict
+from typing_extensions import ReadOnly, Required, TypedDict
 
 
 class GeminiTranscriptionAudioInput(TypedDict):

litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py

@@ -0,0 +1,20 @@
+"""Types for the Hide Secrets guardrail."""
+
+from pydantic import Field
+
+from .base import GuardrailConfigModel
+
+
+class HideSecretsGuardrailConfigModel(GuardrailConfigModel):
+    """Configuration for the Hide Secrets guardrail. Detection runs in-process
+    on the detect-secrets library; ``detect_secrets_config`` overrides the
+    bundled plugin set."""
+
+    detect_secrets_config: dict | None = Field(  # mutable-ok: UI type derivation maps dict to "object"
+        default=None,
+        description="Optional detect-secrets configuration (plugins_used, filters_used) overriding the bundled plugin set",
+    )
+
+    @staticmethod
+    def ui_friendly_name() -> str:
+        return "Hide Secrets"

litellm/types/utils.py

@@ -2879,6 +2879,10 @@ RoutingDecisionCause = Literal[
     # routed model does not accept image input, so the nearest higher capable tier or
     # default_model served instead. The displaced placement rides in signals.
     "modality_escalation",
+    # modality_pin_override replaced a KEPT session-affinity pin for this request only: the turn
+    # carries an image the pinned model cannot accept. The stored pin is untouched, so the next
+    # text turn replays it. Distinct from "modality_escalation", which never displaces a pin.
+    "modality_pin_override",
     "session_affinity_pin",
     "session_affinity_escalation",
     # classification_mode 'user_turn': the request is an agent loop's continuation turn (no new

ruff-strict-budget.json

@@ -1,6 +1,6 @@
 {
   "ANN001": {
-    "limit": 2985
+    "limit": 2984
   },
   "ANN002": {
     "limit": 71
@@ -57,7 +57,7 @@
     "limit": 3
   },
   "BLE001": {
-    "limit": 2917
+    "limit": 2916
   },
   "C401": {
     "limit": 8

tests/code_coverage_tests/check_py310_typing_imports.py

@@ -0,0 +1,150 @@
+import ast
+import os
+import sys
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset(
+    {
+        "NotRequired",
+        "Required",
+        "Self",
+        "LiteralString",
+        "Never",
+        "assert_never",
+        "assert_type",
+        "reveal_type",
+        "TypeVarTuple",
+        "Unpack",
+        "dataclass_transform",
+        "override",
+        "TypeAliasType",
+        "get_original_bases",
+        "ReadOnly",
+        "TypeIs",
+        "NoDefault",
+        "get_protocol_members",
+        "is_protocol",
+        "evaluate_forward_ref",
+        "TypeForm",
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class TypingImportViolation:
+    file: str
+    line: int
+    name: str
+
+
+def _walk_with_ancestors(
+    node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = ()
+) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]:
+    yield node, ancestors
+    for field_name, field_value in ast.iter_fields(node):
+        if isinstance(field_value, ast.AST):
+            yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name)))
+        elif isinstance(field_value, list):
+            for child in field_value:
+                if isinstance(child, ast.AST):
+                    yield from _walk_with_ancestors(child, (*ancestors, (node, field_name)))
+
+
+def _is_sys_version_info(node: ast.AST) -> bool:
+    return (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id == "sys"
+        and node.attr == "version_info"
+    )
+
+
+def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool:
+    nearest_if: Final[tuple[ast.If, str] | None] = next(
+        (
+            (ancestor, field_name)
+            for ancestor, field_name in reversed(ancestors)
+            if isinstance(ancestor, ast.If)
+        ),
+        None,
+    )
+    if nearest_if is None:
+        return False
+    enclosing_if, branch = nearest_if
+    test: Final[ast.expr] = enclosing_if.test
+    if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left):
+        return False
+    operator: Final[ast.cmpop] = test.ops[0]
+    return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or (
+        isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse"
+    )
+
+
+def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]:
+    path: Final[Path] = Path(file_path)
+    tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    return tuple(
+        violation
+        for node, ancestors in _walk_with_ancestors(tree)
+        if not _is_version_guarded(ancestors)
+        for violation in _violations_for_node(node, path)
+    )
+
+
+def _violations_for_node(
+    node: ast.AST, path: Path
+) -> tuple[TypingImportViolation, ...]:
+    if isinstance(node, ast.ImportFrom) and node.module == "typing":
+        return tuple(
+            TypingImportViolation(file=str(path), line=node.lineno, name=alias.name)
+            for alias in node.names
+            if alias.name in PY311_PLUS_TYPING_NAMES
+        )
+    if (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id == "typing"
+        and node.attr in PY311_PLUS_TYPING_NAMES
+    ):
+        return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),)
+    return ()
+
+
+def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]:
+    base_path: Final[Path] = Path(base_dir)
+    return tuple(
+        violation
+        for directory in (
+            base_path / "litellm",
+            base_path / "enterprise",
+            base_path / "litellm-proxy-extras" / "litellm_proxy_extras",
+        )
+        if directory.exists()
+        for path in directory.rglob("*.py")
+        for violation in scan_file(path)
+    )
+
+
+def main() -> None:
+    violations: Final[tuple[TypingImportViolation, ...]] = scan_directory()
+    if violations:
+        message: Final[str] = "\n".join(
+            (
+                "Python 3.10-incompatible typing imports found:",
+                *(
+                    f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; "
+                    "import it from typing_extensions instead because litellm supports Python 3.10"
+                    for violation in violations
+                ),
+            )
+        )
+        sys.stdout.write(f"{message}\n")
+        raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead")
+    sys.stdout.write("No Python 3.10-incompatible typing imports found.\n")
+
+
+if __name__ == "__main__":
+    main()

tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py

@@ -370,12 +370,12 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
         custom_llm_provider,
         litellm_params,
         logging_obj,
+        embedding_executor=None,
         extra_headers=None,
         extra_body=None,
         timeout=None,
         client=None,
         _is_async=False,
-        embedding_executor=None,
     ):
         litellm_params_dict = (
             litellm_params.model_dump(exclude_none=False)

tests/openai_endpoints_tests/test_e2e_openai_responses_api.py

@@ -1,5 +1,5 @@
 import httpx
-from openai import OpenAI, BadRequestError, APIStatusError
+from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError
 import pytest
 
 
@@ -105,10 +105,9 @@ def test_streaming_response():
     assert len(collected_chunks) > 0
 
 
-def test_bad_request_error():
+def test_model_not_found_error():
     client = get_test_client()
-    with pytest.raises(BadRequestError):
-        # Trigger error with invalid model name
+    with pytest.raises(NotFoundError):
         client.responses.create(model="non-existent-model", input="This should fail")
 
 

tests/router_unit_tests/test_router_anthropic_messages_fallback.py

@@ -338,6 +338,112 @@ def test_record_pre_routing_selection_writes_only_the_internal_bucket():
     assert kwargs["metadata"] == {"user_id": "u1"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
+async def test_generic_only_row_recovers_safeguard_refusal(stream):
+    """With no content-policy list configured, a generic fallback row covers safeguard refusals,
+    so the dashboard's generic fallbacks work without config-only content_policy rows."""
+    fake = FakeAnthropicUpstream()
+    router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
+
+    with fake.install():
+        response = await router.aanthropic_messages(
+            model="fable-tier", max_tokens=16, stream=stream, messages=[{"role": "user", "content": "hi"}]
+        )
+        body = await _collect(response) if stream else response
+
+    if stream:
+        assert b'"refusal"' not in body
+        assert b"text_delta" in body
+    else:
+        assert body["stop_reason"] == "end_turn"
+    assert len(fake.calls) == 2
+    assert "claude-opus-5" in fake.calls[1]
+
+
+@pytest.mark.asyncio
+async def test_configured_content_policy_list_stays_authoritative_over_generic_rows():
+    fake = FakeAnthropicUpstream()
+    router = Router(
+        model_list=[FABLE_TIER, OPUS_TARGET],
+        fallbacks=[{"fable-tier": ["opus-target"]}],
+        content_policy_fallbacks=[{"unrelated-group": ["opus-target"]}],
+    )
+
+    with fake.install():
+        response = await router.aanthropic_messages(
+            model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}]
+        )
+
+    assert response["stop_reason"] == "refusal"
+    assert len(fake.calls) == 1
+
+
+def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy():
+    router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}])
+    stamped = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
+
+    assert router._refusal_fallback_available("router-group", stamped) is True
+    assert router._refusal_fallback_available("router-group", {}) is False
+    assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
+
+
+def test_chat_content_filter_gate_unchanged_by_generic_rows():
+    """The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's
+    content_filter gate keeps its long-standing content-policy-only semantics."""
+    from litellm.types.utils import Choices, ModelResponse
+
+    router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
+    response = ModelResponse(choices=[Choices(finish_reason="content_filter")])
+
+    assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
+async def test_disable_fallbacks_returns_the_refusal_instead_of_raising(stream):
+    """A request that opted out of fallbacks must receive the provider's refusal response,
+    never a ContentPolicyViolationError the dispatcher refuses to recover."""
+    fake = FakeAnthropicUpstream()
+    router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
+
+    with fake.install():
+        response = await router.aanthropic_messages(
+            model="fable-tier",
+            max_tokens=16,
+            stream=stream,
+            disable_fallbacks=True,
+            messages=[{"role": "user", "content": "hi"}],
+        )
+        body = await _collect(response) if stream else response
+
+    if stream:
+        assert b'"stop_reason": "refusal"' in body
+    else:
+        assert body["stop_reason"] == "refusal"
+    assert len(fake.calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_disable_fallbacks_beats_a_content_policy_row_too():
+    fake = FakeAnthropicUpstream()
+    router = Router(
+        model_list=[FABLE_TIER, OPUS_TARGET],
+        content_policy_fallbacks=[{"fable-tier": ["opus-target"]}],
+    )
+
+    with fake.install():
+        response = await router.aanthropic_messages(
+            model="fable-tier",
+            max_tokens=16,
+            disable_fallbacks=True,
+            messages=[{"role": "user", "content": "hi"}],
+        )
+
+    assert response["stop_reason"] == "refusal"
+    assert len(fake.calls) == 1
+
+
 def test_refusal_gate_keys_on_pre_routing_tier_stamp():
     router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])
 

tests/rust-python-harness/README.md

@@ -8,14 +8,17 @@ The matrix always has these SDK columns:
 - `messages / amessages`
 - `responses / aresponses`
 - `count_tokens`
+- `chat_completions / acompletion`
+- `transcription / atranscription`
 
-The harness has three deliberately broad test-strategy folders:
+The harness has four deliberately broad test-strategy folders:
 
 | Strategy | Folder |
 | --- | --- |
 | Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) |
 | Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) |
 | Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) |
+| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) |
 
 ## Run it
 
@@ -112,7 +115,7 @@ The initial end-to-end entries deliberately show `◐`: the repository has Rust
 
 ## Attach parity tests
 
-Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
+Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
 
 ```json
 {
@@ -123,7 +126,7 @@ Each of the three folders contains a concise `README.md` and a `strategy.json`.
 }
 ```
 
-Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
+Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
 
 Use these coverage values:
 

tests/rust-python-harness/cli.py

@@ -6,9 +6,10 @@ from collections.abc import Sequence
 from pathlib import Path
 
 from .catalog import load_catalog
-from .models import HarnessCase, Strategy
+from .models import SDK_FUNCTIONS, HarnessCase, Strategy
 from .runner import run_pytest
 from .ui import make_dashboard
+from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness"
@@ -40,9 +41,17 @@ def _parser() -> argparse.ArgumentParser:
         action="append",
         default=[],
         dest="sdk_functions",
-        choices=("ocr", "messages", "responses", "count_tokens"),
+        choices=SDK_FUNCTIONS,
         help="run only this SDK function",
     )
+    parser.add_argument(
+        "--validate-ledger",
+        action="store_true",
+        help=(
+            "report Python<->Rust test-parity ledger gaps and drift instead of "
+            "running the dashboard; narrow with --function"
+        ),
+    )
     parser.add_argument(
         "--plain",
         action="store_true",
@@ -100,7 +109,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[
     )
     sdk_functions = _pick_values(
         "SDK functions",
-        [(name, name) for name in ("ocr", "messages", "responses", "count_tokens")],
+        [(name, name) for name in SDK_FUNCTIONS],
     )
     return strategy_ids, sdk_functions
 
@@ -131,6 +140,38 @@ def _print_catalog(strategies: Sequence[Strategy]) -> None:
             print(f"  {case.sdk_function:12} {case.coverage.value:14} {selectors}")
 
 
+def _print_function_report(report: FunctionReport) -> None:
+    print(f"\n{report.sdk_function}")
+    if report.ledger is None or report.audit is None:
+        print("  no ledger yet")
+        return
+    ledger, audit = report.ledger, report.audit
+    print(
+        f"  {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust "
+        f"({ledger.percentage}%)"
+    )
+    print(f"  {len(ledger.rust_only_tests)} rust-only tests with no python counterpart")
+    if audit.is_clean:
+        print("  ledger is in sync with the live test files")
+        return
+    for label, items in (
+        ("ledger references a python test that no longer exists", audit.missing_python_tests),
+        ("python test exists but is not tracked in the ledger", audit.stale_python_tests),
+        ("ledger references a rust test that no longer exists", audit.missing_rust_tests),
+        ("rust test exists but is not tracked in the ledger", audit.stale_rust_tests),
+    ):
+        for item in items:
+            print(f"  {label}: {item}")
+
+
+def _validate_ledger(sdk_functions: set[str]) -> int:
+    functions = sdk_functions or set(SDK_FUNCTIONS)
+    reports = tuple(build_function_report(function) for function in sorted(functions))
+    for report in reports:
+        _print_function_report(report)
+    return 0 if all(report.is_clean for report in reports) else 1
+
+
 def main(argv: Sequence[str] | None = None) -> int:
     args = _parser().parse_args(argv)
     if args.coverage and importlib.util.find_spec("pytest_cov") is None:
@@ -138,6 +179,8 @@ def main(argv: Sequence[str] | None = None) -> int:
             "--coverage requires the project's pytest-cov dependency; run with "
             "`poetry run python -m tests.rust-python-harness --coverage`"
         )
+    if args.validate_ledger:
+        return _validate_ledger(set(args.sdk_functions))
     strategies = load_catalog()
     if args.list:
         _print_catalog(strategies)

tests/rust-python-harness/e2e_fuzz_tests/strategy.json

@@ -7,6 +7,8 @@
     "ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
     "messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
     "responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."},
-    "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."}
+    "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."},
+    "chat_completions": {"coverage": "partial", "selectors": ["tests/test_litellm/rust_bridge/test_chat_completions.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
+    "transcription": {"coverage": "partial", "selectors": ["tests/test_litellm/test_audio_transcription_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}
   }
 }

tests/rust-python-harness/existing_e2e_test_sdk/README.md

@@ -0,0 +1,3 @@
+# Existing e2e SDK tests
+
+Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically.

tests/rust-python-harness/existing_e2e_test_sdk/strategy.json

@@ -0,0 +1,14 @@
+{
+  "order": 40,
+  "id": "existing_e2e_test_sdk",
+  "label": "Existing e2e SDK tests",
+  "description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.",
+  "functions": {
+    "ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."},
+    "messages": {"coverage": "planned", "selectors": []},
+    "responses": {"coverage": "planned", "selectors": []},
+    "count_tokens": {"coverage": "planned", "selectors": []},
+    "chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."},
+    "transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."}
+  }
+}

tests/rust-python-harness/models.py

@@ -33,7 +33,7 @@ class ConfidenceLevel(str, Enum):
     LOW = "LOW"
 
 
-SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens")
+SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription")
 
 
 @dataclass(frozen=True)

tests/rust-python-harness/runner.py

@@ -15,6 +15,8 @@ UpdateCallback = Callable[[HarnessRun], None]
 def selector_matches_node(selector: str, nodeid: str) -> bool:
     normalized_selector = selector.replace("\\", "/")
     normalized_nodeid = nodeid.replace("\\", "/")
+    if normalized_selector.endswith("/"):
+        return normalized_nodeid.startswith(normalized_selector)
     if "::" in normalized_selector:
         return normalized_nodeid == normalized_selector or normalized_nodeid.startswith(
             f"{normalized_selector}["

tests/rust-python-harness/shared/parity/ledger.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+@dataclass(frozen=True, slots=True)
+class LedgerEntry:
+    python_file: str
+    python_test: str
+    status: str
+    rust_file: str
+    rust_test: str
+    justification: str
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class RustOnlyEntry:
+    rust_file: str
+    rust_test: str
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class TestLedger:
+    sdk_function: str
+    python_scope: tuple[str, ...]
+    rust_scope: tuple[str, ...]
+    entries: tuple[LedgerEntry, ...]
+    rust_only_tests: tuple[RustOnlyEntry, ...]
+
+    @property
+    def mapped_count(self) -> int:
+        return sum(1 for entry in self.entries if entry.status == "mapped")
+
+    @property
+    def total_count(self) -> int:
+        return len(self.entries)
+
+    @property
+    def percentage(self) -> float:
+        if self.total_count == 0:
+            return 0.0
+        return round(100.0 * self.mapped_count / self.total_count, 1)
+
+
+def _require_string(value: Any, field: str, source: Path) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f"{source}: {field} must be a non-empty string")
+    return value
+
+
+def _require_string_list(value: Any, field: str, source: Path) -> tuple[str, ...]:
+    if not isinstance(value, list) or not all(isinstance(item, str) and item for item in value):
+        raise ValueError(f"{source}: {field} must be a list of non-empty strings")
+    return tuple(value)
+
+
+def _load_entry(data: Any, index: int, source: Path) -> LedgerEntry:
+    if not isinstance(data, dict):
+        raise ValueError(f"{source}: entries[{index}] must be an object")
+    python_file = _require_string(data.get("python_file"), f"entries[{index}].python_file", source)
+    python_test = _require_string(data.get("python_test"), f"entries[{index}].python_test", source)
+    status = data.get("status")
+    if status not in ("mapped", "unmapped"):
+        raise ValueError(f"{source}: entries[{index}].status must be 'mapped' or 'unmapped'")
+
+    if status == "mapped":
+        rust_file = _require_string(data.get("rust_file"), f"entries[{index}].rust_file", source)
+        rust_test = _require_string(data.get("rust_test"), f"entries[{index}].rust_test", source)
+        justification = _require_string(
+            data.get("justification"), f"entries[{index}].justification", source
+        )
+        return LedgerEntry(
+            python_file=python_file,
+            python_test=python_test,
+            status=status,
+            rust_file=rust_file,
+            rust_test=rust_test,
+            justification=justification,
+            reason="",
+        )
+
+    reason = _require_string(data.get("reason"), f"entries[{index}].reason", source)
+    return LedgerEntry(
+        python_file=python_file,
+        python_test=python_test,
+        status=status,
+        rust_file="",
+        rust_test="",
+        justification="",
+        reason=reason,
+    )
+
+
+def _load_rust_only_entry(data: Any, index: int, source: Path) -> RustOnlyEntry:
+    if not isinstance(data, dict):
+        raise ValueError(f"{source}: rust_only_tests[{index}] must be an object")
+    return RustOnlyEntry(
+        rust_file=_require_string(data.get("rust_file"), f"rust_only_tests[{index}].rust_file", source),
+        rust_test=_require_string(data.get("rust_test"), f"rust_only_tests[{index}].rust_test", source),
+        reason=_require_string(data.get("reason"), f"rust_only_tests[{index}].reason", source),
+    )
+
+
+def load_ledger(path: Path) -> TestLedger:
+    with path.open(encoding="utf-8") as stream:
+        data = json.load(stream)
+
+    sdk_function = _require_string(data.get("sdk_function"), "sdk_function", path)
+    python_scope = _require_string_list(data.get("python_scope"), "python_scope", path)
+    rust_scope = _require_string_list(data.get("rust_scope"), "rust_scope", path)
+
+    entries_data = data.get("entries")
+    if not isinstance(entries_data, list):
+        raise ValueError(f"{path}: entries must be a list")
+    entries = tuple(
+        _load_entry(entry, index, path) for index, entry in enumerate(entries_data)
+    )
+
+    rust_only_data = data.get("rust_only_tests")
+    if not isinstance(rust_only_data, list):
+        raise ValueError(f"{path}: rust_only_tests must be a list")
+    rust_only_tests = tuple(
+        _load_rust_only_entry(entry, index, path) for index, entry in enumerate(rust_only_data)
+    )
+
+    return TestLedger(
+        sdk_function=sdk_function,
+        python_scope=python_scope,
+        rust_scope=rust_scope,
+        entries=entries,
+        rust_only_tests=rust_only_tests,
+    )

tests/rust-python-harness/strategies/unit_tests/mapping_validator.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from ...shared.parity.ledger import TestLedger, load_ledger
+from .python_runner import enumerate_python_tests
+from .rust_runner import enumerate_rust_tests
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+LEDGER_ROOT = Path(__file__).parent / "ledgers"
+
+
+def ledger_path_for(sdk_function: str) -> Path:
+    return LEDGER_ROOT / sdk_function / f"{sdk_function}_test_ledger.json"
+
+
+@dataclass(frozen=True, slots=True)
+class AuditReport:
+    missing_python_tests: tuple[str, ...]
+    stale_python_tests: tuple[str, ...]
+    missing_rust_tests: tuple[str, ...]
+    stale_rust_tests: tuple[str, ...]
+
+    @property
+    def is_clean(self) -> bool:
+        return not (
+            self.missing_python_tests
+            or self.stale_python_tests
+            or self.missing_rust_tests
+            or self.stale_rust_tests
+        )
+
+
+def _ledger_python_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]:
+    grouping: dict[str, set[str]] = {path: set() for path in ledger.python_scope}
+    for entry in ledger.entries:
+        grouping.setdefault(entry.python_file, set()).add(entry.python_test)
+    return grouping
+
+
+def _ledger_rust_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]:
+    grouping: dict[str, set[str]] = {path: set() for path in ledger.rust_scope}
+    for entry in ledger.entries:
+        if entry.status == "mapped":
+            grouping.setdefault(entry.rust_file, set()).add(entry.rust_test)
+    for rust_only in ledger.rust_only_tests:
+        grouping.setdefault(rust_only.rust_file, set()).add(rust_only.rust_test)
+    return grouping
+
+
+def audit_ledger(ledger: TestLedger, repo_root: Path = REPO_ROOT) -> AuditReport:
+    missing_python: list[str] = []
+    stale_python: list[str] = []
+    for python_file, ledger_tests in _ledger_python_tests_by_file(ledger).items():
+        actual_tests = enumerate_python_tests(repo_root, python_file)
+        for missing in sorted(ledger_tests - actual_tests):
+            missing_python.append(f"{python_file}:{missing}")
+        for stale in sorted(actual_tests - ledger_tests):
+            stale_python.append(f"{python_file}:{stale}")
+
+    missing_rust: list[str] = []
+    stale_rust: list[str] = []
+    for rust_file, ledger_tests in _ledger_rust_tests_by_file(ledger).items():
+        actual_tests = enumerate_rust_tests(repo_root, rust_file)
+        for missing in sorted(ledger_tests - actual_tests):
+            missing_rust.append(f"{rust_file}:{missing}")
+        for stale in sorted(actual_tests - ledger_tests):
+            stale_rust.append(f"{rust_file}:{stale}")
+
+    return AuditReport(
+        missing_python_tests=tuple(missing_python),
+        stale_python_tests=tuple(stale_python),
+        missing_rust_tests=tuple(missing_rust),
+        stale_rust_tests=tuple(stale_rust),
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class FunctionReport:
+    sdk_function: str
+    ledger: TestLedger | None
+    audit: AuditReport | None
+
+    @property
+    def has_ledger(self) -> bool:
+        return self.ledger is not None
+
+    @property
+    def is_clean(self) -> bool:
+        return self.audit is None or self.audit.is_clean
+
+
+def build_function_report(sdk_function: str, repo_root: Path = REPO_ROOT) -> FunctionReport:
+    path = ledger_path_for(sdk_function)
+    if not path.exists():
+        return FunctionReport(sdk_function=sdk_function, ledger=None, audit=None)
+    ledger = load_ledger(path)
+    return FunctionReport(
+        sdk_function=sdk_function, ledger=ledger, audit=audit_ledger(ledger, repo_root)
+    )

tests/rust-python-harness/strategies/unit_tests/python_runner.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+
+def enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
+    source = (repo_root / relative_path).read_text(encoding="utf-8")
+    tree = ast.parse(source, filename=relative_path)
+
+    module_level: list[str] = []
+    for node in ast.iter_child_nodes(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"):
+            module_level.append(node.name)
+        elif isinstance(node, ast.ClassDef):
+            for child in ast.iter_child_nodes(node):
+                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith(
+                    "test_"
+                ):
+                    module_level.append(f"{node.name}::{child.name}")
+
+    return frozenset(module_level)

tests/rust-python-harness/strategies/unit_tests/rust_runner.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+_RUST_TEST_PATTERN = re.compile(
+    r"#\[(?:test|tokio::test)\][^\n]*\n(?:[^\n]*\n)*?\s*(?:async\s+)?fn\s+(\w+)\s*\("
+)
+
+
+def enumerate_rust_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
+    source = (repo_root / relative_path).read_text(encoding="utf-8")
+    return frozenset(match.group(1) for match in _RUST_TEST_PATTERN.finditer(source))

tests/rust-python-harness/ui.py

@@ -119,7 +119,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
 
         table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function")
         table.add_column("Strategy", ratio=3)
-        for label in ("ocr/aocr", "messages", "responses", "count_tokens"):
+        for label in SDK_FUNCTIONS:
             table.add_column(label, justify="center", ratio=1)
         for strategy in self.strategies:
             cells = []

tests/rust-python-harness/unit_tests_rust/strategy.json

@@ -7,6 +7,8 @@
     "ocr": {"coverage": "planned", "selectors": []},
     "messages": {"coverage": "planned", "selectors": []},
     "responses": {"coverage": "planned", "selectors": []},
-    "count_tokens": {"coverage": "planned", "selectors": []}
+    "count_tokens": {"coverage": "planned", "selectors": []},
+    "chat_completions": {"coverage": "planned", "selectors": []},
+    "transcription": {"coverage": "planned", "selectors": []}
   }
 }

tests/rust-python-harness/validate_sub_methods/strategy.json

@@ -7,6 +7,8 @@
     "ocr": {"coverage": "planned", "selectors": []},
     "messages": {"coverage": "planned", "selectors": []},
     "responses": {"coverage": "planned", "selectors": []},
-    "count_tokens": {"coverage": "planned", "selectors": []}
+    "count_tokens": {"coverage": "planned", "selectors": []},
+    "chat_completions": {"coverage": "planned", "selectors": []},
+    "transcription": {"coverage": "planned", "selectors": []}
   }
 }

tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py

@@ -0,0 +1,274 @@
+"""Tests for the hide-secrets guardrail (LIT-3548).
+
+Covers the three defects from the ticket:
+- ``apply_guardrail`` (the UI test playground path) must redact, not echo.
+- Guardrail runs must record ``standard_logging_guardrail_information`` so
+  Spend Logs / the guardrails monitor show activity, with hits ("mask" +
+  masked_entity_count) distinguishable from clean requests ("allow").
+- Defining ``apply_guardrail`` must NOT reroute proxied traffic off the
+  native ``async_pre_call_hook`` (per-key opt-out and ``data["prompt"]``
+  handling live only on the native path).
+"""
+
+import pytest
+
+from litellm_enterprise.enterprise_callbacks.secret_detection import (
+    _ENTERPRISE_SecretDetection,
+)
+from litellm.caching.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+
+AWS_KEY = "AKIAIOSFODNN7EXAMPLE"
+
+
+def _guardrail() -> _ENTERPRISE_SecretDetection:
+    return _ENTERPRISE_SecretDetection(
+        guardrail_name="hide-secrets", event_hook="pre_call", default_on=True
+    )
+
+
+def _recorded(request_data: dict) -> dict:
+    entries = request_data["metadata"]["standard_logging_guardrail_information"]
+    assert len(entries) == 1
+    return entries[0]
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_redacts_secrets():
+    """Playground path: the returned texts must carry [REDACTED], not the secret."""
+    guardrail = _guardrail()
+    request_data: dict = {"metadata": {}}
+
+    result = await guardrail.apply_guardrail(
+        inputs={"texts": [f"my key is {AWS_KEY}, keep it safe"]},
+        request_data=request_data,
+        input_type="request",
+    )
+
+    assert result["texts"] == ["my key is [REDACTED], keep it safe"]
+
+    recorded = _recorded(request_data)
+    assert recorded["guardrail_status"] == "success"
+    assert recorded["guardrail_response"] == "mask"
+    assert recorded["guardrail_provider"] == "hide-secrets"
+    assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_clean_text_records_allow():
+    guardrail = _guardrail()
+    request_data: dict = {"metadata": {}}
+
+    result = await guardrail.apply_guardrail(
+        inputs={"texts": ["nothing sensitive here"]},
+        request_data=request_data,
+        input_type="request",
+    )
+
+    assert result["texts"] == ["nothing sensitive here"]
+
+    recorded = _recorded(request_data)
+    assert recorded["guardrail_status"] == "success"
+    assert recorded["guardrail_response"] == "allow"
+    assert recorded["masked_entity_count"] == {}
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_records_mask_with_entity_count():
+    """Live-traffic path: a redaction must be visible in spend-log telemetry."""
+    guardrail = _guardrail()
+    data = {
+        "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}],
+        "metadata": {},
+    }
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    assert data["messages"][0]["content"] == "use [REDACTED] for auth"
+
+    recorded = _recorded(data)
+    assert recorded["guardrail_status"] == "success"
+    assert recorded["guardrail_response"] == "mask"
+    assert recorded["guardrail_provider"] == "hide-secrets"
+    assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_clean_request_records_allow():
+    """A request with no secrets must be distinguishable from a redacted one."""
+    guardrail = _guardrail()
+    data = {
+        "messages": [{"role": "user", "content": "what's the weather"}],
+        "metadata": {},
+    }
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    recorded = _recorded(data)
+    assert recorded["guardrail_status"] == "success"
+    assert recorded["guardrail_response"] == "allow"
+    assert recorded["masked_entity_count"] == {}
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_opt_out_records_nothing():
+    """A key with permissions={"hide_secrets": False} skips redaction, so no
+    telemetry is recorded: every reader of a recorded entry (guardrail usage
+    tracking, compliance checks, the spend-log viewer) counts it as a run."""
+    guardrail = _guardrail()
+    content = f"my key is {AWS_KEY}"
+    data = {"messages": [{"role": "user", "content": content}], "metadata": {}}
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(permissions={"hide_secrets": False}),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    assert data["messages"][0]["content"] == content  # untouched
+    assert "standard_logging_guardrail_information" not in data["metadata"]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_still_redacts_text_completion_prompt():
+    """data["prompt"] (str and list) is a native-hook-only surface; it must
+    keep redacting now that the class also implements apply_guardrail."""
+    guardrail = _guardrail()
+    data = {"prompt": f"key {AWS_KEY} end", "metadata": {}}
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+    assert data["prompt"] == "key [REDACTED] end"
+
+    guardrail = _guardrail()
+    data = {"prompt": [f"key {AWS_KEY}", "clean"], "metadata": {}}
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+    assert data["prompt"] == ["key [REDACTED]", "clean"]
+
+
+def test_proxied_traffic_stays_on_native_hooks():
+    """Implementing apply_guardrail must not reroute proxied requests onto the
+    unified path: that path skips ``should_run_check`` (per-key opt-out) and
+    never sees ``data["prompt"]``."""
+    guardrail = _guardrail()
+    assert guardrail.uses_apply_guardrail_interface() is True
+    assert guardrail._deployment_pre_call_target() is guardrail
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_without_texts_records_nothing():
+    """No inputs means nothing was inspected, so no "allow" row is recorded.
+    Empty strings count as no input: there is no content to inspect."""
+    guardrail = _guardrail()
+
+    empty_variants: list[list[str]] = [[], ["", ""]]
+    for texts in empty_variants:
+        request_data: dict = {"metadata": {}}
+        result = await guardrail.apply_guardrail(
+            inputs={"texts": texts}, request_data=request_data, input_type="request"
+        )
+        assert result == {"texts": texts}
+        assert "standard_logging_guardrail_information" not in request_data["metadata"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "data",
+    [
+        pytest.param(
+            {
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image_url", "image_url": {"url": "https://x/y.png"}}
+                        ],
+                    }
+                ],
+                "metadata": {},
+            },
+            id="image_only",
+        ),
+        pytest.param(
+            {"messages": [{"role": "user", "content": ""}], "metadata": {}},
+            id="empty_message",
+        ),
+        pytest.param({"prompt": "", "metadata": {}}, id="empty_prompt"),
+        pytest.param({"prompt": ["", ""], "metadata": {}}, id="empty_prompt_list"),
+    ],
+)
+async def test_pre_call_hook_without_inspectable_text_records_nothing(data: dict):
+    """A payload the guardrail could not inspect (image-only content, empty
+    strings) must not record an "allow" run: monitoring would count a check
+    that never looked at any text."""
+    guardrail = _guardrail()
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    assert "standard_logging_guardrail_information" not in data["metadata"]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_mixed_prompt_list_still_redacts_and_records():
+    """A prompt list mixing empty and real strings is inspected, so the run is
+    recorded and the non-empty entry is still redacted."""
+    guardrail = _guardrail()
+    data = {"prompt": ["", f"key {AWS_KEY}"], "metadata": {}}
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    assert data["prompt"] == ["", "key [REDACTED]"]
+    recorded = _recorded(data)
+    assert recorded["guardrail_response"] == "mask"
+    assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
+
+
+@pytest.mark.asyncio
+async def test_legacy_nameless_instance_records_nothing():
+    """``litellm_settings.callbacks: ["hide_secrets"]`` builds an arg-less
+    instance with no guardrail_name. It still redacts, but recording a nameless
+    entry would flip every spend row's guardrail status with nothing to join on."""
+    guardrail = _ENTERPRISE_SecretDetection()
+    data = {
+        "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}],
+        "metadata": {},
+    }
+
+    await guardrail.async_pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        cache=DualCache(),
+        data=data,
+        call_type="completion",
+    )
+
+    assert data["messages"][0]["content"] == "use [REDACTED] for auth"
+    assert "standard_logging_guardrail_information" not in data["metadata"]

tests/test_litellm/integrations/azure_storage/test_azure_storage.py

@@ -1,10 +1,15 @@
+import asyncio
 import sys
+import threading
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-
-from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
+from litellm.integrations.azure_storage.azure_storage import (
+    AzureBlobStorageLogger,
+    _cached_credential_chain_token_provider,
+)
+from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType
 from litellm.types.utils import StandardLoggingPayload
 
 
@@ -25,6 +30,26 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch):
     monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", "core.usgovcloudapi.net")
 
 
+@pytest.fixture
+def workload_identity_env_vars(monkeypatch):
+    monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account")
+    monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container")
+    for unset in (
+        "AZURE_STORAGE_TENANT_ID",
+        "AZURE_STORAGE_CLIENT_ID",
+        "AZURE_STORAGE_CLIENT_SECRET",
+        "AZURE_STORAGE_ACCOUNT_KEY",
+        "AZURE_STORAGE_ENDPOINT_SUFFIX",
+        "AZURE_CLIENT_SECRET",
+        "AZURE_CREDENTIAL",
+        "AZURE_SCOPE",
+    ):
+        monkeypatch.delenv(unset, raising=False)
+    monkeypatch.setenv("AZURE_CLIENT_ID", "workload-identity-client-id")
+    monkeypatch.setenv("AZURE_TENANT_ID", "workload-identity-tenant-id")
+    monkeypatch.setenv("AZURE_FEDERATED_TOKEN_FILE", "/var/run/secrets/azure/tokens/azure-identity-token")
+
+
 @pytest.mark.asyncio
 async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
     """
@@ -32,17 +57,12 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
     a payload to Azure Blob Storage using the 3-step process (create, append, flush).
     """
     with (
-        patch(
-            "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
-        ) as mock_get_client,
-        patch(
-            "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id"
-        ) as mock_get_token,
+        patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client,
+        patch("litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id") as mock_get_token,
     ):
         # Create mock HTTP client
         mock_http_client = AsyncMock()
-        mock_response = AsyncMock()
-        mock_response.raise_for_status = AsyncMock()
+        mock_response = MagicMock()
         mock_http_client.put.return_value = mock_response
         mock_http_client.patch.return_value = mock_response
         mock_get_client.return_value = mock_http_client
@@ -79,9 +99,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
         put_call_args = mock_http_client.put.call_args
         assert put_call_args[0][0] == f"{expected_base_url}?resource=file"
         assert put_call_args[1]["headers"]["x-ms-version"] is not None
-        assert (
-            put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
-        )
+        assert put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
 
         # Step 2: Append data
         assert mock_http_client.patch.call_count == 2  # Called for append and flush
@@ -89,9 +107,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars):
         assert append_call[0][0] == f"{expected_base_url}?action=append&position=0"
         assert append_call[1]["headers"]["x-ms-version"] is not None
         assert append_call[1]["headers"]["Content-Type"] == "application/json"
-        assert (
-            append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
-        )
+        assert append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token"
         assert "test-log-id-123" in append_call[1]["data"]
 
         # Step 3: Flush data
@@ -110,9 +126,7 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env
     AZURE_STORAGE_ENDPOINT_SUFFIX must reach the Entra-ID REST upload path so a
     sovereign-cloud account is addressed instead of the commercial dfs host.
     """
-    with patch(
-        "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
-    ) as mock_get_client:
+    with patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client:
         mock_http_client = AsyncMock()
         mock_response = MagicMock()
         mock_http_client.put.return_value = mock_response
@@ -127,17 +141,10 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env
 
         await logger.async_upload_payload_to_azure_blob_storage(test_payload)
 
-        expected_base_url = (
-            "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json"
-        )
+        expected_base_url = "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json"
         assert mock_http_client.put.call_args[0][0] == f"{expected_base_url}?resource=file"
-        assert (
-            mock_http_client.patch.call_args_list[0][0][0]
-            == f"{expected_base_url}?action=append&position=0"
-        )
-        assert mock_http_client.patch.call_args_list[1][0][0].startswith(
-            f"{expected_base_url}?action=flush"
-        )
+        assert mock_http_client.patch.call_args_list[0][0][0] == f"{expected_base_url}?action=append&position=0"
+        assert mock_http_client.patch.call_args_list[1][0][0].startswith(f"{expected_base_url}?action=flush")
 
 
 @pytest.mark.asyncio
@@ -148,9 +155,7 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars)
     """
     fake_aio_module = MagicMock()
 
-    with patch.dict(
-        sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}
-    ):
+    with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
         logger = AzureBlobStorageLogger()
         await logger.get_service_client()
 
@@ -160,14 +165,180 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars)
     )
 
 
+@pytest.mark.asyncio
+async def test_upload_authenticates_through_the_credential_chain_under_workload_identity(
+    workload_identity_env_vars,
+):
+    build_provider = MagicMock(return_value=lambda: "workload-identity-token")
+    with patch(  # test-quality-ok: REST client is created inside the method; assert emitted request headers
+        "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
+    ) as mock_get_client:
+        mock_http_client = AsyncMock()
+        mock_http_client.put.return_value = MagicMock()
+        mock_http_client.patch.return_value = MagicMock()
+        mock_get_client.return_value = mock_http_client
+
+        logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
+        await logger.async_upload_payload_to_azure_blob_storage({"id": "wif-log-id"})
+
+    build_provider.assert_called_once_with()
+    assert logger.azure_auth_token == "workload-identity-token"
+    sent_headers = [mock_http_client.put.call_args[1]["headers"]] + [
+        call[1]["headers"] for call in mock_http_client.patch.call_args_list
+    ]
+    assert len(sent_headers) == 3
+    assert all(headers["Authorization"] == "Bearer workload-identity-token" for headers in sent_headers)
+
+
+def test_default_chain_provider_is_storage_scoped_and_built_once_per_process():
+    _cached_credential_chain_token_provider.cache_clear()
+    with (
+        patch(  # test-quality-ok: assert the default factory's fixed scope and credential type without constructing Azure SDK credentials
+            "litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_provider",
+            return_value=lambda: "chain-token",
+        ) as mock_builder
+    ):
+        first = _cached_credential_chain_token_provider()
+        second = _cached_credential_chain_token_provider()
+    _cached_credential_chain_token_provider.cache_clear()
+
+    assert first is second
+    assert first() == "chain-token"
+    mock_builder.assert_called_once_with(
+        azure_scope="https://storage.azure.com/.default",
+        azure_credential=AzureCredentialType.DefaultAzureCredential,
+    )
+
+
+@pytest.mark.asyncio
+async def test_chain_tokens_are_read_from_the_provider_on_every_refresh(
+    workload_identity_env_vars,
+):
+    provider = MagicMock(side_effect=["chain-token-1", "chain-token-2"])
+    logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider))
+    await logger.set_valid_azure_ad_token()
+    first_token = logger.azure_auth_token
+    await logger.set_valid_azure_ad_token()
+
+    assert first_token == "chain-token-1"
+    assert logger.azure_auth_token == "chain-token-2"
+    assert provider.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_chain_token_read_yields_to_the_event_loop(workload_identity_env_vars):
+    """
+    The chain walk is blocking I/O (IMDS probe, CLI subprocess), so reading the provider
+    inline would stall every request on the worker. Prove other coroutines run during the read.
+    """
+    loop_was_free = threading.Event()
+
+    def provider() -> str:
+        if not loop_was_free.wait(timeout=5):
+            raise TimeoutError("the event loop never ran the observer while the token was being read")
+        return "chain-token"
+
+    async def observer():
+        loop_was_free.set()
+
+    logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider))
+    observer_task = asyncio.create_task(observer())
+    await logger.set_valid_azure_ad_token()
+    await observer_task
+
+    assert logger.azure_auth_token == "chain-token"
+
+
+@pytest.mark.asyncio
+async def test_empty_string_service_principal_vars_still_use_the_credential_chain(
+    workload_identity_env_vars, monkeypatch
+):
+    for name in ("AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"):
+        monkeypatch.setenv(name, "")
+
+    logger = AzureBlobStorageLogger(
+        build_credential_chain_token_provider=MagicMock(return_value=lambda: "workload-identity-token")
+    )
+    await logger.set_valid_azure_ad_token()
+
+    assert logger.azure_auth_token == "workload-identity-token"
+
+
+@pytest.mark.asyncio
+async def test_client_secret_auth_still_uses_the_storage_scoped_service_principal(mock_env_vars):
+    build_provider = MagicMock()
+    with (
+        patch(  # test-quality-ok: assert the storage scope passed to the shared token factory without making an external auth call
+            "litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id",
+            return_value=lambda: "client-secret-token",
+        ) as mock_entra_id
+    ):
+        logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
+        await logger.set_valid_azure_ad_token()
+
+    assert logger.azure_auth_token == "client-secret-token"
+    build_provider.assert_not_called()
+    assert mock_entra_id.call_args.kwargs == {
+        "tenant_id": "test-tenant-id",
+        "client_id": "test-client-id",
+        "client_secret": "test-client-secret",
+        "scope": "https://storage.azure.com/.default",
+    }
+
+
+@pytest.mark.parametrize(
+    "missing_var",
+    ["AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"],
+)
+@pytest.mark.asyncio
+async def test_partially_configured_service_principal_still_names_the_missing_variable(
+    mock_env_vars, monkeypatch, missing_var
+):
+    monkeypatch.delenv(missing_var)
+
+    build_provider = MagicMock()
+    logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
+    with pytest.raises(ValueError, match=f"Missing required environment variable: {missing_var}"):
+        await logger.set_valid_azure_ad_token()
+
+    build_provider.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_account_key_auth_never_requests_a_token(workload_identity_env_vars, monkeypatch):
+    monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=")
+
+    file_client = MagicMock()
+    file_client.create_file = AsyncMock()
+    file_client.append_data = AsyncMock()
+    file_client.flush_data = AsyncMock()
+    directory_client = MagicMock()
+    directory_client.exists = AsyncMock(return_value=True)
+    directory_client.get_file_client = MagicMock(return_value=file_client)
+    file_system_client = MagicMock()
+    file_system_client.get_directory_client = MagicMock(return_value=directory_client)
+    service_client = MagicMock()
+    service_client.get_file_system_client = MagicMock(return_value=file_system_client)
+    fake_aio_module = MagicMock()
+    fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client)
+
+    build_provider = MagicMock()
+    with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
+        logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider)
+        await logger.async_upload_payload_to_azure_blob_storage({"id": "account-key-log-id"})
+
+    build_provider.assert_not_called()
+    assert logger.azure_auth_token is None
+    file_client.flush_data.assert_awaited_once()
+    assert fake_aio_module.DataLakeServiceClient.call_args.kwargs["credential"] == "dGVzdC1rZXk="
+
+
 @pytest.mark.asyncio
 async def test_service_client_defaults_to_commercial_endpoint(mock_env_vars):
     """Unset AZURE_STORAGE_ENDPOINT_SUFFIX keeps the pre-existing commercial host"""
     fake_aio_module = MagicMock()
 
-    with patch.dict(
-        sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}
-    ):
+    with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
         logger = AzureBlobStorageLogger()
         await logger.get_service_client()
 

tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py

@@ -17,6 +17,7 @@ from unittest.mock import patch
 
 import pytest
 
+import litellm
 from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
@@ -55,15 +56,22 @@ def build_payload(
     response_message: dict[str, Any] | None = None,
     usage_object: dict[str, Any] | None = None,
     model_parameters: dict[str, Any] | None = None,
+    metadata: dict[str, Any] | None = None,
+    model_group: str | None = None,
     prompt_tokens: int = 4447,
 ) -> dict[str, Any]:
+    standard_logging_metadata: dict[str, Any] = {
+        **(metadata or {}),
+        **({"usage_object": usage_object} if usage_object is not None else {}),
+    }
     return {
         "standard_logging_object": {
             "call_type": "acompletion",
             "messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages,
             "response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]},
             "model_parameters": model_parameters or {},
-            "metadata": {"usage_object": usage_object} if usage_object is not None else {},
+            "metadata": standard_logging_metadata,
+            "model_group": model_group,
             "prompt_tokens": prompt_tokens,
             "completion_tokens": 507,
             "total_tokens": prompt_tokens + 507,
@@ -244,6 +252,43 @@ def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMOb
     assert "non_cached_input_tokens" not in payload["metrics"]
 
 
+def test_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 128}})
+
+    assert payload["metrics"]["reasoning_output_tokens"] == 128.0
+
+
+def test_responses_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, usage_object={"output_tokens_details": {"reasoning_tokens": 64}})
+
+    assert payload["metrics"]["reasoning_output_tokens"] == 64.0
+
+
+def test_zero_reasoning_tokens_are_not_reported(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 0}})
+
+    assert "reasoning_output_tokens" not in payload["metrics"]
+
+
+def test_reasoning_tokens_come_from_the_spelling_that_reports_them(logger: DataDogLLMObsLogger) -> None:
+    """A chat-details mapping without the count must not shadow the responses spelling that has it."""
+    payload = build(
+        logger,
+        usage_object={
+            "completion_tokens_details": {"accepted_prediction_tokens": 5},
+            "output_tokens_details": {"reasoning_tokens": 64},
+        },
+    )
+
+    assert payload["metrics"]["reasoning_output_tokens"] == 64.0
+
+
+def test_boolean_reasoning_tokens_are_not_a_count(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": True}})
+
+    assert "reasoning_output_tokens" not in payload["metrics"]
+
+
 def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None:
     payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]})
 
@@ -256,6 +301,340 @@ def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None:
     ]
 
 
+def test_cost_tags_include_present_categories_and_dimensions(logger: DataDogLLMObsLogger) -> None:
+    payload = build(
+        logger,
+        metadata={
+            "user_api_key_user_id": "User 42",
+            "user_api_key_alias": "Primary Key",
+            "team_alias": "Platform",
+            "routing_decision": {
+                "tier": "premium",
+                "cause": "high_complexity",
+                "score": 0.91,
+                "escalated": True,
+                "signals": ["long prompt"],
+                "routed_model": "openai/gpt-5",
+            },
+        },
+        model_group="premium-models",
+    )
+
+    assert payload["tags"][-8:] == [
+        "team:platform",
+        "user:user_42",
+        "key_alias:primary_key",
+        "model_group:premium-models",
+        "router_tier:premium",
+        "router_cause:high_complexity",
+        "router_escalated:true",
+        "routed_model:openai/gpt-5",
+    ]
+    assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == [
+        "team",
+        "user",
+        "key_alias",
+        "model_group",
+        "router_tier",
+        "router_cause",
+        "router_escalated",
+        "routed_model",
+    ]
+
+
+def test_missing_cost_tag_values_are_not_declared(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, metadata={"team_alias": "Platform"})
+
+    assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["team"]
+    assert not any(tag.startswith(("user:", "key_alias:", "model_group:")) for tag in payload["tags"])
+
+
+def test_values_that_normalize_to_empty_are_not_tagged_or_declared(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, metadata={"user_api_key_user_id": "___", "user_api_key_alias": "!!!"}, model_group="tier-1")
+
+    assert not any(tag in ("user:", "key_alias:") for tag in payload["tags"])
+    assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"]
+
+
+def test_a_valueless_tag_from_the_shared_builder_is_not_declared(logger: DataDogLLMObsLogger) -> None:
+    """The team tag comes from the shared builder, which emits it bare when the alias normalizes away."""
+    payload = build(logger, metadata={"team_alias": "!!!"}, model_group="tier-1")
+
+    assert "team:" in payload["tags"]
+    assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"]
+
+
+def test_router_fields_are_flattened(logger: DataDogLLMObsLogger) -> None:
+    payload = build(
+        logger,
+        metadata={
+            "routing_decision": {
+                "tier": "premium",
+                "cause": "high_complexity",
+                "score": 0.91,
+                "escalated": True,
+                "signals": ["secret prompt text"],
+                "routed_model": "openai/gpt-5",
+            }
+        },
+        model_group="premium-models",
+    )
+
+    assert payload["meta"]["metadata"]["router_tier"] == "premium"
+    assert payload["meta"]["metadata"]["router_cause"] == "high_complexity"
+    assert payload["meta"]["metadata"]["router_score"] == 0.91
+    assert payload["meta"]["metadata"]["router_escalated"] is True
+    assert payload["meta"]["metadata"]["router_signals"] == ["secret prompt text"]
+    assert payload["meta"]["metadata"]["routed_model"] == "openai/gpt-5"
+
+
+def test_a_context_escalated_route_reports_as_escalated(logger: DataDogLLMObsLogger) -> None:
+    """The router records a size-driven escalation under its own key, and it is still an escalation."""
+    payload = build(logger, metadata={"routing_decision": {"tier": "premium", "context_escalated": True}})
+
+    assert payload["meta"]["metadata"]["router_escalated"] is True
+    assert "router_escalated:true" in payload["tags"]
+
+
+def test_a_routed_request_that_did_not_escalate_reports_false(logger: DataDogLLMObsLogger) -> None:
+    """Without this the escalation dimension is absent on ordinary traffic, so nothing can group by it."""
+    payload = build(logger, metadata={"routing_decision": {"tier": "simple", "cause": "heuristic_scorer"}})
+
+    assert payload["meta"]["metadata"]["router_escalated"] is False
+    assert "router_escalated:false" in payload["tags"]
+    assert "router_escalated" in payload["meta"]["metadata"]["_dd"]["cost_tags"]
+
+
+def test_a_request_that_never_reached_a_router_has_no_router_fields(logger: DataDogLLMObsLogger) -> None:
+    payload = build(logger, model_group="premium-models")
+
+    assert "router_escalated" not in payload["meta"]["metadata"]
+    assert not any(tag.startswith("router_") for tag in payload["tags"])
+
+
+def test_redacted_payload_keeps_metrics_and_removes_sensitive_fields(logger: DataDogLLMObsLogger) -> None:
+    payload = build_payload(
+        messages=[{"role": "user", "content": "secret prompt"}],
+        response_message={"role": "assistant", "content": "secret response"},
+        usage_object={"prompt_tokens_details": {"cached_tokens": 128}},
+        metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}},
+        model_parameters={"tools": [TOOL_DEFINITION]},
+    )
+    with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
+        with patch("asyncio.create_task"):
+            redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
+    redacted_payload = redacted_logger.redact_standard_logging_payload_from_model_call_details(payload)
+    result = json.loads(
+        safe_dumps(
+            redacted_logger.create_llm_obs_payload(
+                redacted_payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
+            )
+        )
+    )
+
+    assert result["meta"]["input"]["messages"][0]["content"] == "redacted-by-litellm"
+    assert result["meta"]["output"]["messages"][0]["content"] == "redacted-by-litellm"
+    assert result["meta"]["metadata"]["router_tier"] == "premium"
+    assert "router_signals" not in result["meta"]["metadata"]
+    assert "routing_decision" not in result["meta"]["metadata"]
+    assert "tool_definitions" not in result["meta"]
+    assert result["metrics"]["cache_read_input_tokens"] == 128.0
+    assert result["metrics"]["total_cost"] == 0.02
+
+
+def test_redaction_drops_the_routing_record_carried_in_metadata(logger: DataDogLLMObsLogger) -> None:
+    """The whole routing record rides along in metadata, so dropping the flat copy alone leaks the prompt."""
+    with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
+        with patch("asyncio.create_task"):
+            redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
+    result = json.loads(
+        safe_dumps(
+            redacted_logger.create_llm_obs_payload(
+                build_payload(
+                    metadata={
+                        "routing_decision": {
+                            "tier": "premium",
+                            "cause": "keyword_rule",
+                            "signals": ["secret prompt text"],
+                            "matched_keyword": "secret keyword",
+                            "escalation_keyword": "secret escalation",
+                        }
+                    }
+                ),
+                datetime(2026, 9, 1, 12, 0, 0),
+                datetime(2026, 9, 1, 12, 0, 2),
+            )
+        )
+    )
+
+    assert "routing_decision" not in result["meta"]["metadata"]
+    assert result["meta"]["metadata"]["router_tier"] == "premium"
+    assert result["meta"]["metadata"]["router_cause"] == "keyword_rule"
+    assert "secret" not in safe_dumps(result["meta"]["metadata"])
+
+
+def test_a_failure_span_redacts_its_messages(logger: DataDogLLMObsLogger) -> None:
+    """The redaction hook only runs on success, so the failure span has to redact for itself."""
+    failed = build_payload(messages=[{"role": "user", "content": "secret prompt"}])
+    failed["standard_logging_object"]["status"] = "failure"
+    failed["standard_logging_object"]["response"] = None
+    failed["standard_logging_object"]["error_information"] = {"error_message": "boom", "error_class": "BadRequestError"}
+    with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
+        with patch("asyncio.create_task"):
+            redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True)
+    result = json.loads(
+        safe_dumps(
+            redacted_logger.create_llm_obs_payload(
+                failed, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
+            )
+        )
+    )
+
+    assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
+    assert result["meta"]["output"]["messages"] == []
+    assert result["status"] == "error"
+
+
+def test_excluding_messages_from_the_logging_payload_still_ships_the_span(logger: DataDogLLMObsLogger) -> None:
+    """`standard_logging_payload_excluded_fields` deletes the key, and a span with no prompt is still a span."""
+    payload = build_payload()
+    del payload["standard_logging_object"]["messages"]
+
+    span = json.loads(
+        safe_dumps(
+            logger.create_llm_obs_payload(payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2))
+        )
+    )
+
+    assert span["meta"]["input"]["messages"] == []
+    assert span["metrics"]["total_cost"] == 0.02
+
+
+def test_an_explicit_redaction_setting_survives_the_global_params(logger: DataDogLLMObsLogger) -> None:
+    """Global params carry defaults for keys the operator never set, and those must not win."""
+    with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
+        with patch("asyncio.create_task"):
+            with patch.object(  # test-quality-ok: the ctor reads this module global with no injection seam
+                litellm, "datadog_llm_observability_params", {}
+            ):
+                configured_logger = DataDogLLMObsLogger(
+                    turn_off_message_logging=True
+                )  # test-quality-ok: verifies ctor setting
+
+    assert configured_logger.turn_off_message_logging is True
+
+
+def _redacting_logger(
+    **kwargs: Any,
+) -> DataDogLLMObsLogger:  # test-quality-ok: shared test factory accepts init variants
+    with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True):
+        with patch("asyncio.create_task"):
+            return DataDogLLMObsLogger(**kwargs)
+
+
+def _span_json(logger_under_test: DataDogLLMObsLogger, payload: dict[str, Any]) -> dict[str, Any]:
+    span = logger_under_test.create_llm_obs_payload(
+        payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)
+    )
+    return json.loads(safe_dumps(span))
+
+
+def test_redaction_keeps_the_conversation_shape_without_its_content() -> None:
+    """Roles and message count survive so the trace stays legible; contents and tool payloads do not."""
+    result = _span_json(
+        _redacting_logger(turn_off_message_logging=True),
+        build_payload(
+            messages=[
+                {"role": "user", "content": "secret prompt"},
+                {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]},
+            ],
+            response_message={"role": "assistant", "content": "secret response"},
+        ),
+    )
+
+    assert result["meta"]["input"]["messages"] == [
+        {"role": "user", "content": "redacted-by-litellm"},
+        {"role": "assistant", "content": "redacted-by-litellm"},
+    ]
+    assert result["meta"]["output"]["messages"] == [{"role": "assistant", "content": "redacted-by-litellm"}]
+
+
+def test_redaction_drops_unrecognized_and_malformed_message_roles() -> None:
+    """Caller-controlled role values must not bypass redaction or crash span creation."""
+    result = _span_json(
+        _redacting_logger(turn_off_message_logging=True),
+        build_payload(
+            messages=[
+                {"role": "SECRET-39402", "content": "hello"},
+                {"role": ["SECRET-39402"], "content": "hello"},
+                {"role": {"secret": "SECRET-39402"}, "content": "hello"},
+                {"role": "agent", "content": "hello"},
+            ]
+        ),
+    )
+
+    assert result["meta"]["input"]["messages"] == [
+        {"role": "", "content": "redacted-by-litellm"},
+        {"role": "", "content": "redacted-by-litellm"},
+        {"role": "", "content": "redacted-by-litellm"},
+        {"role": "agent", "content": "redacted-by-litellm"},
+    ]
+    assert "SECRET-39402" not in safe_dumps(result)
+
+
+def test_the_deprecated_message_logging_flag_engages_the_same_redaction() -> None:
+    """The platform redacts for `message_logging is not True`, so this callback's own gate must agree."""
+    result = _span_json(
+        _redacting_logger(message_logging=False),
+        build_payload(
+            messages=[{"role": "user", "content": "secret prompt"}],
+            model_parameters={"tools": [TOOL_DEFINITION]},
+            metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}},
+        ),
+    )
+
+    assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
+    assert "tool_definitions" not in result["meta"]
+    assert "routing_decision" not in result["meta"]["metadata"]
+
+
+def test_a_truthy_redaction_setting_redacts_like_the_shared_hook() -> None:
+    """The shared hook redacts on truthiness, so a config-provided string must not half-redact the span."""
+    result = _span_json(
+        _redacting_logger(turn_off_message_logging="yes"),
+        build_payload(messages=[{"role": "user", "content": "secret prompt"}]),
+    )
+
+    assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}]
+
+
+def test_redaction_drops_every_prompt_carrying_metadata_record(logger: DataDogLLMObsLogger) -> None:
+    """Tool arguments, retrieved text, and the guardrail's copy of the request ride in metadata records too."""
+    sensitive_metadata: dict[str, Any] = {
+        "requester_metadata": {"note": "secret prompt text"},
+        "prompt_management_metadata": {"prompt_id": "p1", "prompt_variables": {"topic": "secret"}},
+        "mcp_tool_call_metadata": {"name": "search", "arguments": {"query": "secret"}},
+        "vector_store_request_metadata": [{"query": "secret"}],
+    }
+
+    def sensitive_payload() -> dict[str, Any]:
+        payload = build_payload(metadata=sensitive_metadata)
+        payload["standard_logging_object"]["guardrail_information"] = [
+            {"guardrail_name": "g", "guardrail_request": {"messages": [{"content": "secret prompt"}]}}
+        ]
+        return payload
+
+    redacted = _span_json(_redacting_logger(turn_off_message_logging=True), sensitive_payload())
+    unredacted = _span_json(logger, sensitive_payload())
+
+    assert "secret" not in safe_dumps(redacted["meta"]["metadata"])
+    for record in sensitive_metadata:
+        assert record not in redacted["meta"]["metadata"]
+        assert record in unredacted["meta"]["metadata"]
+    assert redacted["meta"]["metadata"]["guardrail_information"] is None
+    assert unredacted["meta"]["metadata"]["guardrail_information"] is not None
+
+
 def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None:
     """The Anthropic surface declares tools unwrapped, with input_schema instead of parameters."""
     payload = build(
@@ -272,6 +651,15 @@ def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogL
     assert "tool_definitions" not in build(logger)["meta"]
 
 
+def test_a_ddtrace_integer_parent_id_is_forwarded_as_its_string(logger: DataDogLLMObsLogger) -> None:
+    """ddtrace hands span ids as ints; dropping them detaches the span from its APM trace."""
+    kwargs = build_payload()
+    kwargs["litellm_params"]["metadata"]["parent_id"] = 8675309
+    start = datetime(2026, 9, 1, 12, 0, 0)
+    span = json.loads(safe_dumps(logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=2))))
+    assert span["parent_id"] == "8675309"
+
+
 def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None:
     """A truncated argument string is still the only record of what the model tried to call."""
     payload = build(

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -21,6 +21,7 @@ from litellm.litellm_core_utils.streaming_handler import (
 from litellm.types.utils import (
     CompletionTokensDetailsWrapper,
     Delta,
+    ModelResponse,
     ModelResponseStream,
     PromptTokensDetailsWrapper,
     StandardLoggingPayload,
@@ -1750,7 +1751,7 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params():
     assert complete_response.usage.cost == 0.00025
 
     # Use the real propagation method from CustomStreamWrapper
-    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
+    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openrouter")
 
     assert "additional_headers" in complete_response._hidden_params
     assert (
@@ -1769,14 +1770,12 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params():
     assert provider_cost == 0.00025
 
 
-def test_perplexity_streaming_dict_cost_propagates_to_hidden_params():
-    """
-    Regression: Perplexity reports usage.cost as a breakdown object, which used to
-    blow up the end of the stream with
-    `float() argument must be a string or a real number, not 'dict'`.
-    """
+def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator():
     import litellm
-    from litellm.cost_calculator import get_response_cost_from_hidden_params
+    from litellm.cost_calculator import (
+        get_response_cost_from_hidden_params,
+        response_cost_calculator,
+    )
 
     chunks = [
         ModelResponseStream(
@@ -1828,13 +1827,81 @@ def test_perplexity_streaming_dict_cost_propagates_to_hidden_params():
 
     assert complete_response is not None
 
-    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
+    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "perplexity")
 
-    assert (
-        get_response_cost_from_hidden_params(complete_response._hidden_params)
-        == 0.00503
+    assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None
+    assert response_cost_calculator(
+        response_object=complete_response,
+        model="perplexity/sonar",
+        custom_llm_provider="perplexity",
+        call_type="completion",
+        optional_params={},
+    ) == pytest.approx(0.00503)
+
+
+def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
+    import litellm
+    from litellm.cost_calculator import (
+        get_response_cost_from_hidden_params,
+        response_cost_calculator,
     )
 
+    model = "openai/streams-cost-in-nanodollars"
+    litellm.register_model(
+        {
+            model: {
+                "input_cost_per_token": 1e-6,
+                "output_cost_per_token": 2e-6,
+                "litellm_provider": "openai",
+                "mode": "chat",
+            }
+        }
+    )
+    complete_response = ModelResponse(
+        id="chatcmpl-openai-compatible",
+        model=model,
+        choices=[],
+        usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000),
+    )
+
+    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai")
+
+    assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None
+    assert response_cost_calculator(
+        response_object=complete_response,
+        model=model,
+        custom_llm_provider="openai",
+        call_type="completion",
+        optional_params={},
+    ) == pytest.approx(2e-5)
+
+
+def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch):
+    import litellm
+    from litellm.cost_calculator import (
+        get_response_cost_from_hidden_params,
+        response_cost_calculator,
+    )
+
+    complete_response = ModelResponse(
+        id="chatcmpl-xai",
+        model="grok-4-latest",
+        choices=[],
+        usage=Usage(completion_tokens=353, prompt_tokens=198, total_tokens=551, cost=0.0009956),
+    )
+
+    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "xai")
+
+    assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None
+    monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5})
+    assert response_cost_calculator(
+        response_object=complete_response,
+        model="xai/grok-4-latest",
+        custom_llm_provider="xai",
+        call_type="completion",
+        optional_params={},
+    ) == pytest.approx(0.0009956 * 1.5)
+
 
 def test_provider_reported_cost_ignores_unusable_shapes():
     assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None

tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py

@@ -0,0 +1,35 @@
+"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py."""
+
+import json
+
+from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
+from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
+
+
+def _text_event(text: str) -> str:
+    return json.dumps(
+        {
+            "jsonrpc": "2.0",
+            "id": "req-1",
+            "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]},
+        }
+    )
+
+
+def _status_event() -> str:
+    return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}})
+
+
+class TestA2AGuardrailHandlerStreamingScanKey:
+    def test_key_joins_the_text_of_every_message_event(self):
+        key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")])
+        assert key == StreamingScanKey(texts=("hello world",))
+
+    def test_events_without_text_leave_the_key_unchanged(self):
+        handler = A2AGuardrailHandler()
+        events = [_text_event("hello")]
+        assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events)
+
+    def test_unparseable_items_are_ignored(self):
+        key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"])
+        assert key.texts == ("hi",)

tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -13,6 +13,7 @@ import pytest
 
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
 from litellm.llms.anthropic.chat.guardrail_translation.handler import (
     AnthropicMessagesHandler,
 )
@@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults:
         }
         later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)]
         assert {"type": "text", "text": "Now fetch the page."} in later_blocks
+
+
+class TestAnthropicMessagesHandlerStreamingScanKey:
+    """get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
+
+    @staticmethod
+    def _sse(event_type, data):
+        return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+    def _text_delta(self, text):
+        return self._sse(
+            "content_block_delta",
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
+        )
+
+    def test_key_is_empty_before_any_text_arrives(self):
+        head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}})
+        key = AnthropicMessagesHandler().get_streaming_scan_key([head])
+        assert key == StreamingScanKey(texts=("",))
+
+    def test_key_accumulates_text_deltas(self):
+        key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")])
+        assert key.texts == ("hello world",)
+        assert key.stream_ended is False
+
+    def _stop(self, stop_reason):
+        return self._sse(
+            "message_delta",
+            {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}},
+        )
+
+    def test_stop_without_tool_use_scans_the_same_payload(self):
+        handler = AnthropicMessagesHandler()
+        open_key = handler.get_streaming_scan_key([self._text_delta("hi")])
+        ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")])
+        assert ended_key.stream_ended is True
+        assert ended_key == open_key
+
+    def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self):
+        handler = AnthropicMessagesHandler()
+        tool_use = self._sse(
+            "content_block_start",
+            {
+                "type": "content_block_start",
+                "index": 1,
+                "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}},
+            },
+        )
+        open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
+        ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
+        assert open_key == StreamingScanKey(texts=("hi",))
+        assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key != open_key

tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py

@@ -27,6 +27,20 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch):
     monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", GOV_SUFFIX)
 
 
+@pytest.fixture
+def credential_chain_env_vars(monkeypatch):
+    monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account")
+    monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container")
+    for name in (
+        "AZURE_STORAGE_TENANT_ID",
+        "AZURE_STORAGE_CLIENT_ID",
+        "AZURE_STORAGE_CLIENT_SECRET",
+        "AZURE_STORAGE_ACCOUNT_KEY",
+        "AZURE_STORAGE_ENDPOINT_SUFFIX",
+    ):
+        monkeypatch.delenv(name, raising=False)
+
+
 def _make_backend() -> AzureBlobStorageBackend:
     backend = AzureBlobStorageBackend()
     backend.azure_auth_token = "mock-azure-ad-token"
@@ -42,6 +56,29 @@ def _mock_upload_client() -> AsyncMock:
     return client
 
 
+@pytest.mark.asyncio
+async def test_upload_file_with_credential_chain(credential_chain_env_vars):
+    client = _mock_upload_client()
+    build_provider = MagicMock(return_value=lambda: "workload-identity-token")
+
+    with patch(  # test-quality-ok: the backend creates its REST client internally; assert the emitted authorization header
+        "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client
+    ):
+        backend = AzureBlobStorageBackend(build_credential_chain_token_provider=build_provider)
+        storage_url = await backend.upload_file(
+            file_content=b"hello",
+            filename="report.json",
+            content_type="application/json",
+            path_prefix="logs",
+            file_naming_strategy="original_filename",
+        )
+
+    build_provider.assert_called_once_with()
+    assert storage_url == "https://test-account.blob.core.windows.net/test-container/logs/report.json"
+    assert client.put.call_args[1]["headers"]["Authorization"] == "Bearer workload-identity-token"
+    assert client.patch.call_count == 2
+
+
 @pytest.mark.parametrize(
     "env_fixture, expected_suffix",
     [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)],
@@ -125,10 +162,7 @@ async def test_download_file_accepts_url_persisted_before_the_suffix_was_set(moc
         )
 
     assert content == b"file-bytes"
-    assert (
-        client.get.call_args[0][0]
-        == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json"
-    )
+    assert client.get.call_args[0][0] == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json"
 
 
 @pytest.mark.parametrize(
@@ -178,10 +212,7 @@ async def test_download_file_drops_query_string_from_the_stored_url(mock_env_var
             "https://test-account.blob.core.windows.net/test-container/logs/report.json?sig=redacted&se=2026"
         )
 
-    assert (
-        client.get.call_args[0][0]
-        == "https://test-account.blob.core.windows.net/test-container/logs/report.json"
-    )
+    assert client.get.call_args[0][0] == "https://test-account.blob.core.windows.net/test-container/logs/report.json"
 
 
 @pytest.mark.parametrize(

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -8,6 +8,7 @@ import litellm
 
 from litellm import get_model_info, supports_reasoning, supports_vision
 from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
+from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
 from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
 from litellm.types.utils import (
     ChatCompletionMessageToolCall,
@@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id():
     )
 
 
+def test_get_fireworks_session_id_ignores_proxy_generated_session_id():
+    """general_settings.missing_session_id: generate stamps a fresh id per request; sending it
+    as x-session-affinity would pin every request to a different node."""
+    assert (
+        get_fireworks_session_id(
+            {
+                "litellm_session_id": "generated-1",
+                "litellm_trace_id": "generated-1",
+                "metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True},
+            }
+        )
+        is None
+    )
+
+
 def test_handle_message_content_with_tool_calls():
     config = FireworksAIConfig()
     message = Message(

tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -12,6 +12,7 @@ import pytest
 
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
 from litellm.llms.openai.chat.guardrail_translation.handler import (
     OpenAIChatCompletionsHandler,
 )
@@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded:
             )
         ]
         assert handler._check_streaming_has_ended(chunks) is True
+
+
+class TestStreamingScanKey:
+    """get_streaming_scan_key identifies what a sampled round would scan so the
+    unified hook can skip rounds that would re-scan already-cleared text"""
+
+    @staticmethod
+    def _chunk(content, finish_reason=None, index=0):
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+        return ModelResponseStream(
+            choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)]
+        )
+
+    def test_key_carries_accumulated_text_and_open_stream(self):
+        handler = OpenAIChatCompletionsHandler()
+        key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
+        assert key == StreamingScanKey(texts=("hello",))
+
+    def test_chunks_without_text_leave_the_key_unchanged(self):
+        handler = OpenAIChatCompletionsHandler()
+        before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
+        after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)])
+        assert after == before
+
+    def test_finish_chunk_without_tool_calls_scans_the_same_payload(self):
+        handler = OpenAIChatCompletionsHandler()
+        open_key = handler.get_streaming_scan_key([self._chunk("hi")])
+        ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")])
+        assert open_key.stream_ended is False
+        assert ended_key.stream_ended is True
+        assert ended_key == open_key
+
+    def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self):
+        from litellm.types.utils import (
+            ChatCompletionDeltaToolCall,
+            Delta,
+            Function,
+            ModelResponseStream,
+            StreamingChoices,
+        )
+
+        handler = OpenAIChatCompletionsHandler()
+        tool_call = ChatCompletionDeltaToolCall(
+            id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
+        )
+        tool_chunk = ModelResponseStream(
+            choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)]
+        )
+        open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk])
+        ended_key = handler.get_streaming_scan_key(
+            [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
+        )
+        assert open_key == StreamingScanKey(texts=("hi",))
+        assert ended_key.texts == ("hi",)
+        assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key != open_key
+
+    def test_text_after_the_first_choice_finishes_still_changes_the_key(self):
+        handler = OpenAIChatCompletionsHandler()
+        first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)]
+        key_at_first_finish = handler.get_streaming_scan_key(first_done)
+        key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)])
+        assert key_at_first_finish.stream_ended is True
+        assert key_after_more_text.stream_ended is True
+        assert key_after_more_text != key_at_first_finish
+
+    def test_non_stream_items_are_ignored(self):
+        handler = OpenAIChatCompletionsHandler()
+        key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"])
+        assert key.texts == ("hi",)

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks:
         dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"]
         assert len(dones) == 1
         assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy."
+
+
+class TestOpenAIResponsesHandlerStreamingScanKey:
+    """get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
+
+    @staticmethod
+    def _delta(sequence_number, text):
+        return {
+            "type": "response.output_text.delta",
+            "sequence_number": sequence_number,
+            "item_id": "msg_1",
+            "output_index": 0,
+            "content_index": 0,
+            "delta": text,
+        }
+
+    def test_no_events_yields_no_key(self):
+        assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None
+
+    def test_key_accumulates_deltas_while_the_stream_is_open(self):
+        from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
+
+        key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")])
+        assert key == StreamingScanKey(texts=("hello",))
+
+    def test_typed_delta_events_accumulate_like_dicts(self):
+        from litellm.types.llms.openai import OutputTextDeltaEvent
+
+        events = [
+            OutputTextDeltaEvent(
+                type="response.output_text.delta",
+                item_id="msg_1",
+                output_index=0,
+                content_index=0,
+                delta=text,
+                sequence_number=i,
+            )
+            for i, text in enumerate(("hel", "lo"))
+        ]
+        key = OpenAIResponsesHandler().get_streaming_scan_key(events)
+        assert key.texts == ("hello",)
+        assert key.stream_ended is False
+
+    def test_events_without_text_leave_the_key_unchanged(self):
+        handler = OpenAIResponsesHandler()
+        events = [self._delta(0, "hi")]
+        quiet = events + [{"type": "response.in_progress", "sequence_number": 1}]
+        assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events)
+
+    @staticmethod
+    def _completed(sequence_number, output):
+        return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}}
+
+    def test_completed_event_keys_on_the_final_output_text(self):
+        handler = OpenAIResponsesHandler()
+        message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
+        open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
+        ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])])
+        assert ended_key.stream_ended is True
+        assert ended_key == open_key
+
+    def test_completed_event_with_a_function_call_changes_the_key(self):
+        handler = OpenAIResponsesHandler()
+        message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
+        function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
+        open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
+        ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])])
+        assert ended_key.texts == ("hi",)
+        assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
+        assert ended_key != open_key
+
+    def test_completed_event_reads_every_output_text_part(self):
+        from litellm.types.responses.main import GenericResponseOutputItem, OutputText
+
+        item = GenericResponseOutputItem(
+            type="message",
+            id="msg_1",
+            status="completed",
+            role="assistant",
+            content=[
+                OutputText(type="output_text", text="one", annotations=[]),
+                OutputText(type="output_text", text="two", annotations=[]),
+            ],
+        )
+        key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])])
+        assert key.texts == ("one", "two")
+
+    def test_output_item_done_round_is_never_deduped(self):
+        done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
+        assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None

tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py

@@ -7,12 +7,13 @@ transformations for the Responses API.
 Source: litellm/llms/xai/responses/transformation.py
 """
 
-from unittest.mock import MagicMock
-
+from unittest.mock import MagicMock, Mock
 
+import httpx
 import pytest
 
 import litellm
+from litellm.llms.xai.cost_calculator import cost_per_token
 from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig
 from litellm.responses.utils import ResponseAPILoggingUtils
 from litellm.types.llms.openai import (
@@ -400,3 +401,94 @@ class TestXAIResponsesWebSearchBilling:
 
         bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage)
         assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS
+
+
+class TestXAIResponsesReportedCost:
+    """xAI reports what it charged; the transformation moves it to where litellm bills from.
+
+    ``ResponseAPILoggingUtils`` copies ``usage.cost`` onto the chat Usage that cost
+    tracking prices, so restating ``cost_in_usd_ticks`` there is what makes /v1/responses
+    bill the reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756.
+    """
+
+    @staticmethod
+    def _response_body(usage: dict) -> dict:
+        return {
+            "id": "resp_xai",
+            "object": "response",
+            "created_at": 0,
+            "model": "grok-4-latest",
+            "status": "completed",
+            "output": [],
+            "parallel_tool_calls": False,
+            "tool_choice": "auto",
+            "tools": [],
+            "usage": usage,
+        }
+
+    def _transformed_usage(self, usage: dict) -> ResponseAPIUsage | None:
+        raw_response = httpx.Response(status_code=200, json=self._response_body(usage))
+
+        response = XAIResponsesAPIConfig().transform_response_api_response(
+            model="grok-4-latest",
+            raw_response=raw_response,
+            logging_obj=Mock(),
+        )
+        return response.usage
+
+    def test_reported_cost_reaches_the_cost_calculator(self):
+        usage = self._transformed_usage(
+            {
+                "input_tokens": 100,
+                "output_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": 37756000,
+            }
+        )
+
+        assert usage.cost == 0.0037756
+
+        chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage)
+        assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
+
+    def test_streamed_reported_cost_reaches_the_cost_calculator(self):
+        event = XAIResponsesAPIConfig().transform_streaming_response(
+            model="grok-4-latest",
+            parsed_chunk={
+                "type": "response.completed",
+                "sequence_number": 7,
+                "response": self._response_body(
+                    {
+                        "input_tokens": 100,
+                        "output_tokens": 200,
+                        "total_tokens": 300,
+                        "cost_in_usd_ticks": 37756000,
+                    }
+                ),
+            },
+            logging_obj=Mock(),
+        )
+
+        assert isinstance(event, ResponseCompletedEvent)
+        chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage)
+        assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
+
+    def test_usage_without_a_reported_cost_is_left_alone(self):
+        usage = self._transformed_usage(
+            {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300}
+        )
+
+        assert usage.cost is None
+
+    def test_negative_reported_cost_is_not_carried(self):
+        """A caller who can set api_base must not be able to report negative spend."""
+        usage = self._transformed_usage(
+            {
+                "input_tokens": 100,
+                "output_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": -37756000,
+            }
+        )
+
+        assert usage.cost is None

tests/test_litellm/llms/xai/test_xai_chat_transformation.py

@@ -1,9 +1,14 @@
+from unittest.mock import Mock
 
-
+import httpx
 import pytest
 
 import litellm
-from litellm.llms.xai.chat.transformation import XAIChatConfig
+from litellm.llms.xai.chat.transformation import (
+    XAIChatCompletionStreamingHandler,
+    XAIChatConfig,
+)
+from litellm.llms.xai.cost_calculator import cost_per_token
 from litellm.types.utils import (
     CompletionTokensDetailsWrapper,
     ModelResponse,
@@ -195,3 +200,113 @@ class TestXAIChatWebSearchBilling:
         )
 
         assert with_search - without_search == pytest.approx(3 * 5.0 / 1000.0)
+
+
+class TestXAIReportedCost:
+    """xAI reports what it charged; the transformation moves it to where litellm bills from.
+
+    ``cost`` is the field litellm already carries a provider stated cost in, so restating
+    ``cost_in_usd_ticks`` there is what lets ``llms/xai/cost_calculator.py`` bill the
+    reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756.
+    """
+
+    @staticmethod
+    def _transformed_usage(usage: dict) -> Usage:
+        raw_response = httpx.Response(
+            status_code=200,
+            json={
+                "id": "chatcmpl-xai",
+                "object": "chat.completion",
+                "created": 0,
+                "model": "grok-4-latest",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "hi"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": usage,
+            },
+        )
+
+        response = XAIChatConfig().transform_response(
+            model="grok-4-latest",
+            raw_response=raw_response,
+            model_response=ModelResponse(),
+            logging_obj=Mock(),
+            request_data={},
+            messages=[{"role": "user", "content": "hi"}],
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+        return response.usage
+
+    def test_reported_cost_reaches_the_cost_calculator(self):
+        usage = self._transformed_usage(
+            {
+                "prompt_tokens": 100,
+                "completion_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": 37756000,
+            }
+        )
+
+        assert usage.cost == 0.0037756
+        assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756)
+
+    def test_usage_without_a_reported_cost_is_left_alone(self):
+        usage = self._transformed_usage(
+            {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
+        )
+
+        assert getattr(usage, "cost", None) is None
+
+    def test_negative_reported_cost_is_not_carried(self):
+        """A caller who can set api_base must not be able to report negative spend."""
+        usage = self._transformed_usage(
+            {
+                "prompt_tokens": 100,
+                "completion_tokens": 200,
+                "total_tokens": 300,
+                "cost_in_usd_ticks": -37756000,
+            }
+        )
+
+        assert getattr(usage, "cost", None) is None
+
+    def test_streamed_reported_cost_survives_chunk_aggregation(self):
+        """Streamed spend only matches if the conversion happens on the chunk.
+
+        Chunk aggregation rebuilds usage from the fields it models plus ``cost``, so a
+        chunk still carrying only ``cost_in_usd_ticks`` loses the reported amount.
+        """
+        handler = XAIChatCompletionStreamingHandler(
+            streaming_response=iter([]), sync_stream=True
+        )
+
+        parsed = handler.chunk_parser(
+            {
+                "id": "chatcmpl-xai",
+                "object": "chat.completion.chunk",
+                "created": 0,
+                "model": "grok-4-latest",
+                "choices": [],
+                "usage": {
+                    "prompt_tokens": 100,
+                    "completion_tokens": 200,
+                    "total_tokens": 300,
+                    "cost_in_usd_ticks": 37756000,
+                },
+            }
+        )
+
+        assert parsed.usage.cost == 0.0037756
+
+        assembled = litellm.stream_chunk_builder(chunks=[parsed])
+        assert assembled.usage.cost == 0.0037756
+        assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == (
+            0.0,
+            0.0037756,
+        )

tests/test_litellm/llms/xai/test_xai_cost_calculator.py

@@ -7,7 +7,10 @@ import os
 
 import litellm
 from litellm.types.utils import (
+    Choices,
     CompletionTokensDetailsWrapper,
+    Message,
+    ModelResponse,
     PromptTokensDetailsWrapper,
     Usage,
 )
@@ -361,6 +364,145 @@ class TestXAICostCalculator:
             response_object=object(), usage=usage
         )
 
+    def test_reported_cost_is_preferred_over_token_math(self):
+        """The amount xAI reported, carried on usage.cost by the transformation, is billed.
+
+        It lands entirely on completion cost because xAI does not split its total by
+        direction, the same shape the perplexity calculator returns.
+        """
+        usage = Usage(
+            prompt_tokens=100,
+            completion_tokens=200,
+            total_tokens=300,
+            cost=0.0037756,
+        )
+
+        prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+        assert prompt_cost == 0.0
+        assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10)
+
+    def test_reported_cost_suppresses_web_search_surcharge(self):
+        """The reported total already covers server-side tool calls.
+
+        Without the suppression these 3 searches would be billed a second time on
+        top of the total xAI already charged.
+        """
+        usage = Usage(
+            prompt_tokens=100,
+            completion_tokens=50,
+            total_tokens=150,
+            prompt_tokens_details=PromptTokensDetailsWrapper(
+                text_tokens=100,
+                web_search_requests=3,
+            ),
+            cost=0.0037756,
+        )
+
+        assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0
+
+    def test_web_search_surcharge_suppressed_through_the_dispatcher(self):
+        """The suppression has to hold on the path cost tracking actually uses.
+
+        Legacy behaviour stays intact when xAI reports no cost.
+        """
+        from litellm.llms import get_cost_for_web_search_request
+
+        usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
+        setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3})
+
+        assert get_cost_for_web_search_request("xai", usage, {}) > 0.0
+
+        reported = Usage(
+            prompt_tokens=100, completion_tokens=50, total_tokens=150, cost=0.0037756
+        )
+        setattr(reported, "server_side_tool_usage_details", {"web_search_calls": 3})
+        assert get_cost_for_web_search_request("xai", reported, {}) == 0.0
+
+    def test_no_reported_cost_falls_back_to_token_math(self):
+        """Absent the provider figure, nothing changes for existing callers."""
+        usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
+
+        prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+        assert prompt_cost > 0.0
+        assert completion_cost > 0.0
+
+    def test_malformed_reported_cost_falls_back_to_token_math(self):
+        """A junk value must not fail the request, fall back to calculating."""
+        usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
+        setattr(usage, "cost", "not-a-number")
+
+        prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+        assert prompt_cost > 0.0
+        assert completion_cost > 0.0
+
+    def test_boolean_reported_cost_falls_back_to_token_math(self):
+        """True is an int in python and would otherwise be billed as $1."""
+        usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
+        setattr(usage, "cost", True)
+
+        prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+        assert prompt_cost > 0.0
+        assert completion_cost > 0.0
+        assert completion_cost != 1.0
+
+    def test_negative_reported_cost_is_rejected(self):
+        """A negative amount must never reach spend tracking.
+
+        A caller who can set api_base controls the response body, so trusting a
+        negative figure would let them subtract from their own recorded spend and
+        slip past a budget. Fall back to token pricing instead, and keep charging
+        the web search surcharge, since no trustworthy total was reported.
+        """
+        usage = Usage(
+            prompt_tokens=100,
+            completion_tokens=200,
+            total_tokens=300,
+            cost=-0.0037756,
+        )
+        setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3})
+
+        prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+        assert prompt_cost > 0.0
+        assert completion_cost > 0.0
+        assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0
+
+    def test_non_finite_reported_cost_is_rejected(self):
+        """NaN compares false against every budget threshold.
+
+        Usage stores a provider supplied cost without validating it, so a caller who
+        controls the response body could report NaN and leave spend >= max_budget
+        false for the life of the key rather than mispricing one request. The
+        infinities are refused alongside it. Fall back to token pricing and keep
+        charging the web search surcharge, since no trustworthy total was reported.
+        """
+        for reported_cost in (float("nan"), float("inf"), float("-inf")):
+            usage = Usage(
+                prompt_tokens=100,
+                completion_tokens=200,
+                total_tokens=300,
+                cost=reported_cost,
+            )
+            setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3})
+
+            prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+
+            assert math.isfinite(prompt_cost), reported_cost
+            assert math.isfinite(completion_cost), reported_cost
+            assert prompt_cost > 0.0, reported_cost
+            assert completion_cost > 0.0, reported_cost
+            assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0, reported_cost
+
+    def test_zero_reported_cost_is_honoured(self):
+        """A reported zero is a real answer, not a missing value."""
+        usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300, cost=0.0)
+
+        assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0)
+
     def test_grok_4_20_beta_reasoning_cost_calculation(self):
         """Test cost calculation for grok-4.20-beta-0309-reasoning model."""
         usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
@@ -437,6 +579,48 @@ class TestXAICostCalculator:
         assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
         assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
 
+    def test_custom_pricing_beats_the_reported_cost(self):
+        response = ModelResponse(
+            id="chatcmpl-xai",
+            model="grok-4-latest",
+            choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")],
+            usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956),
+        )
+
+        billed = litellm.completion_cost(
+            completion_response=response,
+            model="xai/grok-4-latest",
+            custom_llm_provider="xai",
+            custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.001},
+            custom_pricing=True,
+        )
+
+        assert math.isclose(billed, 0.551, rel_tol=1e-10)
+
+    def test_deployment_custom_pricing_beats_the_reported_cost(self, monkeypatch):
+        deployment_id = "xai-deployment-priced-by-the-operator"
+        monkeypatch.setitem(
+            litellm.model_cost,
+            deployment_id,
+            {"input_cost_per_token": 0.001, "output_cost_per_token": 0.001, "litellm_provider": "xai", "mode": "chat"},
+        )
+        response = ModelResponse(
+            id="chatcmpl-xai",
+            model="grok-4-latest",
+            choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")],
+            usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956),
+        )
+
+        billed = litellm.completion_cost(
+            completion_response=response,
+            model="xai/grok-4-latest",
+            custom_llm_provider="xai",
+            custom_pricing=True,
+            router_model_id=deployment_id,
+        )
+
+        assert math.isclose(billed, 0.551, rel_tol=1e-10)
+
 
 class TestXAIWebSearchCostHelpers:
     """Focused coverage for web_search / tool-usage helpers in cost_calculator.py."""

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py

@@ -914,6 +914,7 @@ async def test_mcp_get_prompt_success():
     )
     mock_manager.get_prompt_from_server.assert_awaited_once_with(
         server=server,
+        user_api_key_auth=user_api_key_auth,
         prompt_name="hello",
         arguments={"foo": "bar"},
         mcp_auth_header={"Authorization": "token"},
@@ -976,6 +977,7 @@ async def test_mcp_read_resource_success():
     )
     mock_manager.read_resource_from_server.assert_awaited_once_with(
         server=server,
+        user_api_key_auth=user_api_key_auth,
         url="https://example.com/resource",
         mcp_auth_header={"Authorization": "token"},
         extra_headers={"X-Test": "1"},
@@ -8268,7 +8270,9 @@ class TestOboPreflightScopedToAllowedServers:
 
         _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
 
-        preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key)
+        preflight.assert_awaited_once_with(
+            server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None
+        )
 
 
 @pytest.mark.asyncio

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -30,6 +30,7 @@ from mcp.types import (
     TextResourceContents,
 )
 from mcp.types import Tool as MCPTool
+from pydantic import AnyUrl
 
 from litellm.constants import MCP_METADATA_TIMEOUT
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@@ -2270,7 +2271,9 @@ class TestMCPServerManager:
         """prompts/list on an OBO server must exchange the caller's bearer, not connect with none."""
         server = self._token_exchange_server("te-prompts")
         st = await self._capture_subject_token(
-            lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
+            lambda m: m.get_prompts_from_server(
+                server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
+            )
         )
         assert st == "subj-jwt"
 
@@ -2279,7 +2282,9 @@ class TestMCPServerManager:
         """resources/list on an OBO server must exchange the caller's bearer."""
         server = self._token_exchange_server("te-resources")
         st = await self._capture_subject_token(
-            lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
+            lambda m: m.get_resources_from_server(
+                server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
+            )
         )
         assert st == "subj-jwt"
 
@@ -2290,6 +2295,7 @@ class TestMCPServerManager:
         st = await self._capture_subject_token(
             lambda m: m.read_resource_from_server(
                 server=server,
+                user_api_key_auth=None,
                 url="https://up.example.com/r",
                 raw_headers={"authorization": "Bearer subj-jwt"},
             )
@@ -2307,7 +2313,9 @@ class TestMCPServerManager:
             auth_type=MCPAuth.none,
         )
         st = await self._capture_subject_token(
-            lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
+            lambda m: m.get_prompts_from_server(
+                server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
+            )
         )
         assert st is None
 
@@ -3254,7 +3262,7 @@ class TestMCPServerManager:
             new_callable=AsyncMock,
             return_value=mock_client,
         ):
-            prompts = await manager.get_prompts_from_server(server, add_prefix=True)
+            prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True)
 
         mock_client.list_prompts.assert_awaited_once()
         assert len(prompts) == 1
@@ -3289,6 +3297,7 @@ class TestMCPServerManager:
         ):
             result = await manager.get_prompt_from_server(
                 server=server,
+                user_api_key_auth=None,
                 prompt_name="hello",
                 arguments={"tone": "casual"},
             )
@@ -3334,6 +3343,7 @@ class TestMCPServerManager:
         ):
             result = await manager.get_resources_from_server(
                 server=server,
+                user_api_key_auth=None,
                 mcp_auth_header="auth",
                 extra_headers={"X-Test": "1"},
                 add_prefix=True,
@@ -3391,6 +3401,7 @@ class TestMCPServerManager:
         ):
             result = await manager.get_resource_templates_from_server(
                 server=server,
+                user_api_key_auth=None,
                 mcp_auth_header="auth",
                 extra_headers=None,
                 add_prefix=False,
@@ -3441,6 +3452,7 @@ class TestMCPServerManager:
         ) as mock_create_client:
             result = await manager.read_resource_from_server(
                 server=server,
+                user_api_key_auth=None,
                 url="https://example.com/resource",
                 mcp_auth_header="auth",
                 extra_headers={"X-Test": "1"},
@@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth:
 
         assert result.isError is True
         assert "upstream returned HTTP 503" in result.content[0].text
+
+
+class TestLitellmAdmissionKeyIsNeverTheSubjectToken:
+    """The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the
+    RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it
+    present, ``Authorization`` is the caller's own identity token and is exchanged as before."""
+
+    _ADMISSION_KEY: Final = "sk-litellm-virtual-key"
+    _USER_TOKEN: Final = "user-idp-jwt"
+
+    @staticmethod
+    def _token_exchange_server(server_id: str) -> MCPServer:
+        return MCPServer(
+            server_id=server_id,
+            name=f"{server_id}-server",
+            url="https://up.example.com/mcp",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_token_exchange,
+            token_exchange_endpoint="https://idp.example.com/token",
+            client_id="cid",
+            client_secret="csec",
+        )
+
+    @staticmethod
+    def _id_jag_server(server_id: str) -> MCPServer:
+        return MCPServer(
+            server_id=server_id,
+            name=f"{server_id}-server",
+            url="https://up.example.com/mcp",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_id_jag,
+            client_id="cid",
+            client_secret="csec",
+            token_exchange_endpoint="https://idp.example.com/token",
+            id_jag_resource_token_endpoint="https://resource-as.example.com/token",
+        )
+
+    @staticmethod
+    def _recording_provider() -> MagicMock:
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
+
+        provider: Final = MagicMock()
+        provider.resolve_credentials = AsyncMock(
+            return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
+        )
+        return provider
+
+    @staticmethod
+    def _subjects_seen_by(provider: MagicMock) -> list[str | None]:
+        return [
+            call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None
+            for call in provider.resolve_credentials.call_args_list
+        ]
+
+    @staticmethod
+    def _manager_with_recording_client() -> MCPServerManager:
+        manager: Final = MCPServerManager()
+        client: Final = AsyncMock()
+        client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
+        client.list_prompts = AsyncMock(return_value=[])
+        client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[]))
+        manager._create_mcp_client = AsyncMock(return_value=client)
+        return manager
+
+    @staticmethod
+    def _subject_token_given_to_client(manager: MCPServerManager) -> str | None:
+        return manager._create_mcp_client.call_args.kwargs["subject_token"]
+
+    async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth):
+        manager: Final = self._manager_with_recording_client()
+        await manager._call_regular_mcp_tool(
+            mcp_server=server,
+            original_tool_name="tool",
+            arguments={},
+            tasks=[],
+            mcp_auth_header=None,
+            mcp_server_auth_headers=None,
+            oauth2_headers=oauth2_headers,
+            raw_headers=raw_headers,
+            proxy_logging_obj=None,
+            user_api_key_auth=user_api_key_auth,
+        )
+        return self._subject_token_given_to_client(manager)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag])
+    async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type):
+        server = (
+            self._token_exchange_server("te-call")
+            if auth_type == MCPAuth.oauth2_token_exchange
+            else self._id_jag_server("jag-call")
+        )
+        subject_token = await self._call_tool_subject(
+            server,
+            oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token is None
+
+    @pytest.mark.asyncio
+    async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self):
+        """The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only."""
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-rest"),
+            oauth2_headers=None,
+            raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token is None
+
+    @pytest.mark.asyncio
+    async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-split"),
+            oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
+            raw_headers={
+                "X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}",
+                "authorization": f"Bearer {self._USER_TOKEN}",
+            },
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token == self._USER_TOKEN
+
+    @pytest.mark.asyncio
+    async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self):
+        """Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead."""
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-empty-header"),
+            oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token is None
+
+    @pytest.mark.asyncio
+    async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self):
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-same-key"),
+            oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            raw_headers={
+                "x-litellm-api-key": self._ADMISSION_KEY,
+                "authorization": f"Bearer {self._ADMISSION_KEY}",
+            },
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token is None
+
+    @pytest.mark.asyncio
+    async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self):
+        """A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential."""
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-second-key"),
+            oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"},
+            raw_headers={
+                "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
+                "authorization": "Bearer sk-another-virtual-key",
+            },
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert subject_token is None
+
+    @pytest.mark.asyncio
+    async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self):
+        subject_token = await self._call_tool_subject(
+            self._token_exchange_server("te-jwt"),
+            oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
+            raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"),
+        )
+        assert subject_token == self._USER_TOKEN
+
+    @pytest.mark.asyncio
+    async def test_tools_list_with_only_the_litellm_key_has_no_subject(self):
+        manager: Final = self._manager_with_recording_client()
+        manager._fetch_tools_with_timeout = AsyncMock(return_value=[])
+        await manager._get_tools_from_server(
+            server=self._token_exchange_server("te-list-key"),
+            oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+        )
+        assert self._subject_token_given_to_client(manager) is None
+
+    @pytest.mark.asyncio
+    async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self):
+        manager: Final = self._manager_with_recording_client()
+        await manager.get_prompts_from_server(
+            server=self._token_exchange_server("te-prompts-key"),
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+            raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+        )
+        assert self._subject_token_given_to_client(manager) is None
+
+    @pytest.mark.asyncio
+    async def test_resource_read_with_only_the_litellm_key_has_no_subject(self):
+        manager: Final = self._manager_with_recording_client()
+        await manager.read_resource_from_server(
+            server=self._token_exchange_server("te-read-key"),
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+            url=AnyUrl("file:///notes.txt"),
+            raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+        )
+        assert self._subject_token_given_to_client(manager) is None
+
+    @pytest.mark.asyncio
+    async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
+        manager: Final = self._manager_with_recording_client()
+        await manager.read_resource_from_server(
+            server=self._token_exchange_server("te-read-split"),
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+            url=AnyUrl("file:///notes.txt"),
+            raw_headers={
+                "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
+                "authorization": f"Bearer {self._USER_TOKEN}",
+            },
+        )
+        assert self._subject_token_given_to_client(manager) == self._USER_TOKEN
+
+    @pytest.mark.asyncio
+    async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self):
+        provider: Final = self._recording_provider()
+        manager = MCPServerManager(cred_provider=provider)
+        server = MCPServer(
+            server_id="te-openapi",
+            name="te_openapi",
+            server_name="te_openapi",
+            url=None,
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_token_exchange,
+            token_exchange_endpoint="https://idp.example.com/token",
+            client_id="cid",
+            client_secret="csec",
+            spec_path="https://api.example.com/openapi.json",
+        )
+        user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice")
+
+        await manager.resolve_openapi_upstream_auth(
+            mcp_server=server,
+            oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+            raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+            mcp_auth_header=None,
+            user_api_key_auth=user_auth,
+            forwarded_headers=None,
+        )
+        await manager.resolve_openapi_upstream_auth(
+            mcp_server=server,
+            oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
+            raw_headers={
+                "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
+                "authorization": f"Bearer {self._USER_TOKEN}",
+            },
+            mcp_auth_header=None,
+            user_api_key_auth=user_auth,
+            forwarded_headers=None,
+        )
+        assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN]
+
+    @pytest.mark.asyncio
+    async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self):
+        provider: Final = self._recording_provider()
+        manager = MCPServerManager(cred_provider=provider)
+
+        with pytest.raises(HTTPException) as exc_info:
+            await manager.preflight_token_exchange(
+                server=self._token_exchange_server("te-preflight-key"),
+                oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
+                user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+                raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
+            )
+        assert exc_info.value.status_code == 401
+        headers = exc_info.value.headers or {}
+        assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "")
+        assert self._subjects_seen_by(provider) == []
+
+    @pytest.mark.asyncio
+    async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
+        provider: Final = self._recording_provider()
+        manager = MCPServerManager(cred_provider=provider)
+
+        await manager.preflight_token_exchange(
+            server=self._token_exchange_server("te-preflight-split"),
+            oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
+            raw_headers={
+                "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
+                "authorization": f"Bearer {self._USER_TOKEN}",
+            },
+        )
+        assert self._subjects_seen_by(provider) == [self._USER_TOKEN]

tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py

@@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None):
 @pytest.mark.asyncio
 async def test_openai_moderation_streaming_default_uses_sampled_cadence():
     """Default config samples every 5th streamed chunk and runs a final aggregate
-    pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream
-    calls, plus 1 final = 3 total.
+    pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream
+    round is skipped because chunk 10 already scanned the full text, for 2 total calls
     """
     import litellm
 
@@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence():
             ):
                 pass
 
-        assert patched_make_request.await_count == 3, (
-            f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), "
+        assert patched_make_request.await_count == 2, (
+            f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; "
+            f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
             f"got {patched_make_request.await_count}"
         )
 
@@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder
 @pytest.mark.asyncio
 async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled():
     """With streaming_end_of_stream_only=False and streaming_sampling_rate=2,
-    moderation runs every 2nd chunk during the stream, plus once more at end.
+    moderation runs every 2nd chunk during the stream. The terminal chunk scan covers
+    the final aggregate, for 3 total calls
     """
     import litellm
 
@@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab
             ):
                 pass
 
-        # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls),
-        # plus the final aggregate pass after the stream ends (1 call) = 4 total.
-        assert patched_make_request.await_count == 4, (
-            f"Expected 4 moderation calls (3 sampled + 1 final aggregate), "
+        assert patched_make_request.await_count == 3, (
+            f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped "
+            f"because chunk 6 already scanned the full text), "
             f"got {patched_make_request.await_count}"
         )

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

@@ -1703,8 +1703,8 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts:
 @pytest.mark.parametrize(
     ("configured", "expected_calls"),
     [
-        ({}, 3),
-        ({"streaming_sampling_rate": 2}, 6),
+        ({}, 2),
+        ({"streaming_sampling_rate": 2}, 5),
         ({"streaming_end_of_stream_only": True}, 1),
         ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1),
     ],
@@ -1712,7 +1712,10 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts:
 async def test_streaming_params_from_config_control_output_scan_cadence(
     configured: dict[str, object], expected_calls: int
 ) -> None:
-    """10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once."""
+    """10 chunks: default samples at 5 and 10, rate 2 samples 5 times, end-of-stream scans once.
+
+    The final pass is skipped because chunk 10 already scanned the complete output.
+    """
     handler = _initialize_from_config(mode="post_call", **configured)
 
     assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls

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

@@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
 
     @pytest.mark.asyncio
     async def test_streaming_default_uses_sampled_cadence(self):
-        """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3."""
+        """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10
+        cover the full text, so the end-of-stream round is skipped and there are 2 calls
+        """
         from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
             UnifiedLLMGuardrails,
         )
@@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
             ):
                 pass
 
-        assert mock_post.await_count == 3, (
-            f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), "
+        assert mock_post.await_count == 2, (
+            f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; "
+            f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
             f"got {mock_post.await_count}"
         )
         for call in mock_post.await_args_list:
@@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
 
     @pytest.mark.asyncio
     async def test_streaming_sampling_rate_override(self):
-        """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls."""
+        """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so
+        the end-of-stream round is skipped and there are 3 calls
+        """
         from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
             UnifiedLLMGuardrails,
         )
@@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
             ):
                 pass
 
-        assert mock_post.await_count == 4, (
-            f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), "
+        assert mock_post.await_count == 3, (
+            f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped "
+            f"because chunk 6 already scanned the full text), "
             f"got {mock_post.await_count}"
         )
 

tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py

@@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket:
         assert recorded[0]["guardrail_name"] == "audit-recorder"
         assert recorded[0]["guardrail_status"] == "success"
         assert request_data["metadata"]["user_api_key_user_id"] == "user-1"
+
+
+class _ScanCountingGuardrail(CustomGuardrail):
+    """Pass-through guardrail that records every response-side scan payload."""
+
+    def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False):
+        super().__init__(guardrail_name="scan-counter")
+        self.streaming_sampling_rate = sampling_rate
+        self.streaming_end_of_stream_only = end_of_stream_only
+        self.streaming_buffer_until_moderated = buffer_until_moderated
+        self.guardrail_config = {}
+        self.scans: tuple[dict[str, object], ...] = ()
+
+    def should_run_guardrail(self, data, event_type):  # type: ignore[override]
+        return True
+
+    async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
+        self.scans = (
+            *self.scans,
+            {
+                "texts": list(inputs.get("texts") or []),
+                "tool_calls": list(inputs.get("tool_calls") or []),
+                "model": inputs.get("model"),
+            },
+        )
+        return inputs
+
+
+def _responses_delta(sequence_number, text):
+    return {
+        "type": "response.output_text.delta",
+        "sequence_number": sequence_number,
+        "item_id": "msg_1",
+        "output_index": 0,
+        "content_index": 0,
+        "delta": text,
+    }
+
+
+def _responses_tail(sequence_number, text):
+    return [
+        {
+            "type": "response.output_text.done",
+            "sequence_number": sequence_number,
+            "item_id": "msg_1",
+            "output_index": 0,
+            "content_index": 0,
+            "text": text,
+        },
+        {
+            "type": "response.completed",
+            "sequence_number": sequence_number + 1,
+            "response": {
+                "model": "gpt-5.6",
+                "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}],
+            },
+        },
+    ]
+
+
+class TestStreamingScanDedup:
+    """A sampled round whose scan payload matches the previous round (or carries
+    no text yet) is skipped, so a stream is never re-scanned for output the
+    guardrail already cleared. Regression for LIT-6692."""
+
+    @pytest.fixture(autouse=True)
+    def _use_real_mappings(self, monkeypatch):
+        monkeypatch.setattr(
+            unified_module,
+            "endpoint_guardrail_translation_mappings",
+            load_guardrail_translation_mappings(),
+        )
+
+    @pytest.mark.asyncio
+    async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=3)
+        chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert len(out) == 3
+        assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
+
+    @pytest.mark.asyncio
+    async def test_chat_round_with_unchanged_text_is_skipped(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=3)
+        chunks = [
+            _stream_chunk("a"),
+            _stream_chunk("b"),
+            _stream_chunk("c"),
+            _stream_chunk(None),
+            _stream_chunk(None),
+            _stream_chunk(None),
+            _stream_chunk("d", finish_reason="stop"),
+        ]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert len(out) == 7
+        assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]]
+
+    @pytest.mark.asyncio
+    async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=3)
+        chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert len(out) == 4
+        assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
+
+    @pytest.mark.asyncio
+    async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self):
+        from litellm.types.utils import ChatCompletionDeltaToolCall, Function
+
+        guardrail = _ScanCountingGuardrail(sampling_rate=3)
+        tool_call = ChatCompletionDeltaToolCall(
+            id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
+        )
+        finish = ModelResponseStream(
+            choices=[
+                StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls")
+            ]
+        )
+        chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert len(out) == 4
+        assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]]
+        assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
+
+    @pytest.mark.asyncio
+    async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=3)
+        chunks = [
+            _stream_chunk("a", index=0),
+            _stream_chunk("x", index=1),
+            _stream_chunk("b", finish_reason="stop", index=0),
+            _stream_chunk("y", index=1),
+            _stream_chunk("z", finish_reason="stop", index=1),
+        ]
+
+        await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert len(guardrail.scans) == 2
+        assert any("yz" in text for text in guardrail.scans[-1]["texts"])
+
+    @pytest.mark.asyncio
+    async def test_responses_completed_event_on_sampled_index_is_scanned_once(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=5)
+        deltas = [_responses_delta(i, f"t{i}") for i in range(8)]
+        full_text = "".join(f"t{i}" for i in range(8))
+        chunks = deltas + _responses_tail(8, full_text)
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
+
+        assert len(out) == 10
+        assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]]
+        assert guardrail.scans[-1]["model"] == "gpt-5.6"
+
+    @pytest.mark.asyncio
+    async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=5)
+        deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
+        chunks = deltas + _responses_tail(5, "t0t1t2t3t4")
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
+
+        assert len(out) == 7
+        assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]]
+
+    @pytest.mark.asyncio
+    async def test_responses_completed_carrying_a_function_call_is_still_scanned(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=5)
+        deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
+        completed = {
+            "type": "response.completed",
+            "sequence_number": 5,
+            "response": {
+                "model": "gpt-5.6",
+                "output": [
+                    {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]},
+                    {
+                        "type": "function_call",
+                        "id": "fc_1",
+                        "call_id": "call_1",
+                        "name": "get_weather",
+                        "arguments": '{"city": "Paris"}',
+                        "status": "completed",
+                    },
+                ],
+            },
+        }
+        chunks = deltas + [completed]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
+
+        assert len(out) == 6
+        assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]]
+        assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
+
+    @pytest.mark.asyncio
+    async def test_responses_round_with_unchanged_text_is_skipped(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=5)
+        deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
+        quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)]
+        chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4")
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
+
+        assert len(out) == 12
+        assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},)
+
+    @pytest.mark.asyncio
+    async def test_responses_tool_call_done_event_is_still_scanned(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=2)
+        tool_call_done = {
+            "type": "response.output_item.done",
+            "sequence_number": 1,
+            "output_index": 1,
+            "item": {
+                "type": "function_call",
+                "id": "fc_1",
+                "call_id": "call_1",
+                "name": "get_weather",
+                "arguments": '{"city": "Paris"}',
+                "status": "completed",
+            },
+        }
+        chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi")
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
+
+        assert len(out) == 4
+        assert len(guardrail.scans) == 2
+        assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"]
+        assert guardrail.scans[1]["texts"] == ["hi"]
+
+    @pytest.mark.asyncio
+    async def test_anthropic_skips_empty_round_and_terminal_duplicate(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=2)
+        chunks = _anthropic_message_chunks(["hello ", "world"])
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
+
+        assert out == chunks
+        assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
+
+    @pytest.mark.asyncio
+    async def test_end_of_stream_only_still_scans_exactly_once(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True)
+        chunks = _anthropic_message_chunks(["hello ", "world"])
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
+
+        assert out == chunks
+        assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
+
+    @pytest.mark.asyncio
+    async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self):
+        guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True)
+        chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
+
+        out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
+
+        assert out == chunks
+        assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -670,6 +670,37 @@ def test_get_provider_specific_params():
     )  # Literal type should be select
 
 
+@pytest.mark.asyncio
+async def test_provider_specific_params_includes_hide_secrets():
+    """hide-secrets lives in the enterprise package so it is not in
+    guardrail_class_registry; the endpoint must still advertise it or the
+    Add Guardrail UI dropdown never offers it (LIT-3548)."""
+    from litellm.proxy.guardrails.guardrail_endpoints import (
+        get_provider_specific_params,
+    )
+
+    provider_params = await get_provider_specific_params()
+
+    assert "hide-secrets" in provider_params
+    # populateGuardrailProviders() in the dashboard only lists providers whose
+    # entry carries a ui_friendly_name.
+    assert provider_params["hide-secrets"]["ui_friendly_name"] == "Hide Secrets"
+    assert provider_params["hide-secrets"]["detect_secrets_config"]["required"] is False
+
+
+@pytest.mark.asyncio
+async def test_add_guardrail_settings_restricts_hide_secrets_to_pre_call():
+    """hide-secrets only implements async_pre_call_hook, so offering the other
+    modes in the UI would create configs that boot clean and never run."""
+    from litellm.proxy.guardrails.guardrail_endpoints import (
+        get_guardrail_ui_settings,
+    )
+
+    settings = await get_guardrail_ui_settings()
+
+    assert settings.supported_modes_by_provider["hide-secrets"] == ["pre_call"]
+
+
 def test_optional_params_not_returned_when_not_overridden():
     """Test that optional_params is not returned when the config model doesn't override it"""
     from typing import Optional

tests/test_litellm/proxy/management_endpoints/test_saml_sso.py

@@ -516,6 +516,39 @@ async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated):
     assert result.team_ids == ["team-a", "team-b"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "roles",
+    [
+        ["internal_user", "proxy_admin_viewer"],
+        ["proxy_admin_viewer", "internal_user"],
+    ],
+)
+async def test_multi_valued_role_attribute_resolves_to_highest_privilege(saml_env_idp_initiated, roles):
+    """An assertion carrying several roles must not depend on the order the IdP emitted them in."""
+    key_pem, cert_pem = saml_env_idp_initiated
+    resp = _build_signed_response(
+        key_pem,
+        cert_pem,
+        attributes={
+            "email": ["dave@example.com"],
+            "role": roles,
+        },
+    )
+
+    result = await _acs(_b64(resp), _shared_cache())
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
+
+
+@pytest.mark.asyncio
+async def test_assertion_without_role_attribute_has_no_user_role(saml_env_idp_initiated):
+    key_pem, cert_pem = saml_env_idp_initiated
+    resp = _build_signed_response(key_pem, cert_pem, attributes={"email": ["erin@example.com"]})
+
+    result = await _acs(_b64(resp), _shared_cache())
+    assert result.user_role is None
+
+
 @pytest.mark.asyncio
 async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env):
     cache = DualCache()

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -6598,13 +6598,94 @@ def test_get_litellm_user_role_with_invalid_role():
     assert result is None
 
 
-def test_get_litellm_user_role_with_list_multiple_roles():
-    """Test that get_litellm_user_role takes the first element from a multi-element list."""
+@pytest.mark.parametrize(
+    "role_claim",
+    [
+        ["proxy_admin", "internal_user"],
+        ["internal_user", "proxy_admin"],
+    ],
+)
+def test_get_litellm_user_role_picks_highest_privilege_regardless_of_order(role_claim):
+    """A multi-valued role claim resolves to the most privileged role, not the first one listed."""
     from litellm.proxy._types import LitellmUserRoles
     from litellm.proxy.management_endpoints.types import get_litellm_user_role
 
-    result = get_litellm_user_role(["proxy_admin", "internal_user"])
-    assert result == LitellmUserRoles.PROXY_ADMIN
+    assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN
+
+
+@pytest.mark.parametrize(
+    "role_claim",
+    [
+        ["proxy_admin_viewer", "internal_user"],
+        ["internal_user", "proxy_admin_viewer"],
+    ],
+)
+def test_get_litellm_user_role_keeps_org_spend_visibility_for_mixed_roles(role_claim):
+    """
+    Regression for LIT-6077: a user holding both proxy_admin_viewer and internal_user kept
+    losing org-level spend visibility whenever the IdP happened to list internal_user first.
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
+
+
+def test_get_litellm_user_role_ignores_unrecognised_entries():
+    """Roles LiteLLM does not know about are skipped rather than swallowing the whole claim."""
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(["some_idp_group", "internal_user"]) == LitellmUserRoles.INTERNAL_USER
+    assert get_litellm_user_role(["some_idp_group", "another_group"]) is None
+
+
+def test_get_litellm_user_role_list_lookup_is_case_insensitive():
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(["INTERNAL_USER", "Proxy_Admin"]) == LitellmUserRoles.PROXY_ADMIN
+
+
+@pytest.mark.parametrize(
+    "role_claim",
+    [
+        ["org_admin", "team"],
+        ["team", "org_admin"],
+    ],
+)
+def test_get_litellm_user_role_is_deterministic_for_unranked_roles(role_claim):
+    """Roles outside the privilege hierarchy still resolve the same way in either claim order."""
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(role_claim) == LitellmUserRoles.ORG_ADMIN
+
+
+@pytest.mark.parametrize(
+    "role_claim",
+    [
+        ["org_admin", "internal_user"],
+        ["internal_user", "org_admin"],
+    ],
+)
+def test_get_litellm_user_role_prefers_a_ranked_role_over_an_unranked_one(role_claim):
+    """
+    org_admin, team and customer sit outside the privilege ladder, so a claim mixing one of
+    them with a ranked role settles on the ranked role in either order. Same rule the Entra
+    app_roles and role_mappings paths already follow.
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(role_claim) == LitellmUserRoles.INTERNAL_USER
+
+
+def test_get_litellm_user_role_returns_none_for_non_string_claims():
+    from litellm.proxy.management_endpoints.types import get_litellm_user_role
+
+    assert get_litellm_user_role(None) is None
+    assert get_litellm_user_role({"role": "proxy_admin"}) is None
 
 
 # ============================================================================
@@ -6654,6 +6735,46 @@ def test_process_sso_jwt_access_token_extracts_role_from_access_token():
     assert result.user_role == LitellmUserRoles.PROXY_ADMIN
 
 
+@pytest.mark.parametrize(
+    "role_claim",
+    [
+        ["internal_user", "proxy_admin_viewer"],
+        ["proxy_admin_viewer", "internal_user"],
+    ],
+)
+def test_process_sso_jwt_access_token_resolves_highest_privilege_role(role_claim):
+    """
+    The generic SSO access-token path must land on the same role for a user whose role
+    claim holds several roles, whichever order the IdP emitted them in.
+    """
+    import jwt as pyjwt
+
+    from litellm.proxy._types import LitellmUserRoles
+
+    access_token_str = pyjwt.encode(
+        {"sub": "user-123", "email": "mixed@test.com", "litellm_role": role_claim},
+        "secret",
+        algorithm="HS256",
+    )
+    result = CustomOpenID(
+        id="user-123",
+        email="mixed@test.com",
+        display_name="Mixed Role User",
+        team_ids=[],
+        user_role=None,
+    )
+
+    with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}):
+        process_sso_jwt_access_token(
+            access_token_str=access_token_str,
+            sso_jwt_handler=None,
+            result=result,
+            role_mappings=None,
+        )
+
+    assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
+
+
 def test_process_sso_jwt_access_token_does_not_override_existing_role():
     """
     Test that process_sso_jwt_access_token does NOT override a role that was

tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py

@@ -1104,6 +1104,16 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog(
     assert response.status_code == 200
     payload = response.json()
     assert "anthropic_family" in payload
+    assert payload["1m_context"]["complexity_router_config"]["classifier_type"] == "heuristic_v2"
+    assert payload["1m_context"]["complexity_router_config"]["tiers"] == {
+        "SIMPLE": ["gpt-5.6-luna"],
+        "MEDIUM": ["gpt-5.6-terra"],
+        "COMPLEX": ["claude-opus-5"],
+        "REASONING": ["claude-opus-5"],
+    }
+    assert payload["1m_context"]["complexity_router_config"]["tier_model_configs"] == {
+        "REASONING": [{"model_name": "claude-opus-5", "litellm_params": {"reasoning_effort": "high"}}]
+    }
     for preset in payload.values():
         assert isinstance(preset["label"], str)
         assert isinstance(preset["description"], str)

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

@@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import (
 from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
     TRUSTED_CALLBACK_VARS_FIELD,
 )
+from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
 from litellm.types.utils import CredentialItem
 
 
@@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge():
     }
 
     assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
+
+
+def _request_for(path: str) -> MagicMock:
+    request = MagicMock(spec=Request)
+    request.scope = {"path": path}
+    request.url = MagicMock()
+    request.url.path = path
+    request.url.__str__.return_value = f"http://localhost{path}"
+    request.method = "POST"
+    request.query_params = {}
+    request.headers = {"Content-Type": "application/json"}
+    request.client = MagicMock()
+    request.client.host = "127.0.0.1"
+    return request
+
+
+def _spend_log_session_id(data: dict[str, object]) -> str:
+    """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id."""
+    from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+    from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log
+
+    metadata = data["metadata"]
+    assert isinstance(metadata, dict)
+    litellm_params = get_litellm_params(
+        litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None,
+        litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None,
+        metadata=metadata,
+    )
+    trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(
+        logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"),
+        litellm_params=litellm_params,
+    )
+    return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("request_correlation_in_logs", [False, True])
+async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree(
+    monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool
+):
+    """Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs
+    must be the same generated id, so cross-referencing the two by session_id works. The id is marked
+    as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it."""
+    monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs)
+    data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
+
+    updated = await add_litellm_data_to_request(
+        data=data,
+        request=_request_for("/v1/chat/completions"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": "generate"},
+    )
+
+    callback_session_id = updated["metadata"]["session_id"]
+    assert isinstance(callback_session_id, str) and len(callback_session_id) == 36
+    assert _spend_log_session_id(updated) == callback_session_id
+    assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True
+    assert get_fireworks_session_id(
+        {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]}
+    ) is None
+
+
+@pytest.mark.asyncio
+async def test_missing_session_id_unset_keeps_legacy_divergence():
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o", "messages": []},
+        request=_request_for("/v1/chat/completions"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={},
+    )
+
+    assert "session_id" not in updated["metadata"]
+    assert "litellm_session_id" not in updated
+    assert _spend_log_session_id(updated) == "per-call-random-trace-id"
+
+
+@pytest.mark.asyncio
+async def test_missing_session_id_generate_reuses_traceparent_trace_id():
+    """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it."""
+    request = _request_for("/v1/chat/completions")
+    request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
+
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o", "messages": []},
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": "generate"},
+    )
+
+    assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
+    assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("policy", ["generate", "reject"])
+async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str):
+    request = _request_for("/v1/chat/completions")
+    request.headers = {"x-litellm-session-id": "client-session-1"}
+
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o", "messages": []},
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": policy},
+    )
+
+    assert updated["litellm_session_id"] == "client-session-1"
+    assert updated["metadata"]["session_id"] == "client-session-1"
+    assert _spend_log_session_id(updated) == "client-session-1"
+    assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"]
+    assert (
+        get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]})
+        == "client-session-1"
+    )
+
+
+@pytest.mark.asyncio
+async def test_missing_session_id_reject_accepts_body_metadata_session_id():
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}},
+        request=_request_for("/v1/chat/completions"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": "reject"},
+    )
+
+    assert updated["metadata"]["session_id"] == "body-session-1"
+
+
+@pytest.mark.asyncio
+async def test_missing_session_id_reject_returns_400_without_session_id():
+    with pytest.raises(ProxyException) as exc_info:
+        await add_litellm_data_to_request(
+            data={"model": "gpt-4o", "messages": []},
+            request=_request_for("/v1/chat/completions"),
+            user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+            proxy_config=MagicMock(),
+            general_settings={"missing_session_id": "reject"},
+        )
+
+    assert exc_info.value.code == "400"
+    assert exc_info.value.param == "session_id"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"])
+async def test_missing_session_id_policy_skips_non_inference_routes(path: str):
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o"},
+        request=_request_for(path),
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": "reject"},
+    )
+
+    assert "session_id" not in updated["metadata"]
+
+
+@pytest.mark.asyncio
+async def test_missing_session_id_unknown_value_is_ignored():
+    updated = await add_litellm_data_to_request(
+        data={"model": "gpt-4o", "messages": []},
+        request=_request_for("/v1/chat/completions"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={"missing_session_id": "typo"},
+    )
+
+    assert "session_id" not in updated["metadata"]

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -16,7 +16,7 @@ import litellm
 from litellm import Router
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
-from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
+from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.router_strategy.complexity_router.complexity_router import (
     _CLASSIFICATION_CURRENT_MESSAGE_ONLY,
     _CLASSIFICATION_WITH_CONVERSATION,
@@ -4274,6 +4274,26 @@ class TestSessionAffinity:
         assert first.model == "o1-preview"
         assert second.model == "gpt-4o-mini"
 
+    @pytest.mark.asyncio
+    async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config):
+        """A session id the proxy generated for a request that had none is per request, so
+        it must not create a pin even with session_affinity enabled."""
+        mock_router_instance.cache = DualCache()
+        router = ComplexityRouter(
+            model_name="test-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config=session_affinity_config,
+        )
+        request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}}
+        first = await router.async_pre_routing_hook(
+            model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
+        )
+        second = await router.async_pre_routing_hook(
+            model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
+        )
+        assert first.model == "o1-preview"
+        assert second.model == "gpt-4o-mini"
+
     @pytest.mark.asyncio
     async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config):
         """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the
@@ -10777,6 +10797,9 @@ class TestModalityRouting:
             ("custom_tiers_walk", "premium-model", "modality_escalation"),
             ("pin_kept_bypasses", "text-cheap", "session_affinity_pin"),
             ("pin_replacement_gated", "vision-big", "modality_escalation"),
+            ("pin_override_escalates", "vision-mid", "modality_pin_override"),
+            ("pin_override_same_tier", "vision-cheap", "modality_pin_override"),
+            ("pin_override_inert_without_modality_routing", "text-cheap", "session_affinity_pin"),
             ("adaptive_pick_rewritten", "vision-mid", "modality_escalation"),
         ],
     )
@@ -10827,7 +10850,7 @@ class TestModalityRouting:
             messages = [
                 {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]}
             ]
-        elif path in ("pin_kept_bypasses", "pin_replacement_gated"):
+        elif path.startswith(("pin_kept", "pin_replacement", "pin_override")):
             cache = AsyncMock()
             cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
             mock_router_instance.cache = cache
@@ -10839,6 +10862,13 @@ class TestModalityRouting:
                 messages = [
                     {"role": "user", "content": [{"type": "text", "text": "LITELLM ESCALATE describe this"}, IMG_PART]}
                 ]
+            elif path == "pin_override_same_tier":
+                config["modality_pin_override"] = True
+                config["tiers"]["SIMPLE"] = ["text-cheap", "vision-cheap"]
+                vision["vision-cheap"] = True
+            elif path == "pin_override_inert_without_modality_routing":
+                config["modality_routing"] = False
+            config["modality_pin_override"] = path.startswith("pin_override")
         elif path == "adaptive_pick_rewritten":
             config["adaptive"] = True
             mock_router_instance.model_list = []
@@ -11005,4 +11035,60 @@ class TestModalityRouting:
         from litellm.router_strategy.complexity_router.complexity_router import _decision_is_pinnable
 
         assert _decision_is_pinnable({"cause": "modality_escalation"}) is False
+        assert _decision_is_pinnable({"cause": "modality_pin_override"}) is False
         assert _decision_is_pinnable({"cause": "heuristic_scorer"}) is True
+
+    @pytest.mark.asyncio
+    async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance):
+        """The override is for one request: the session keeps the model it was pinned to."""
+        cache = AsyncMock()
+        cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
+        mock_router_instance.cache = cache
+        router = self._router(
+            mock_router_instance,
+            {
+                "tiers": dict(self.BASE_TIERS),
+                "modality_routing": True,
+                "modality_pin_override": True,
+                "session_affinity": True,
+            },
+            dict(self.BASE_VISION),
+        )
+        request_kwargs = {"metadata": {"session_id": "s1"}}
+
+        image_turn = await router.async_pre_routing_hook(
+            model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE
+        )
+        assert image_turn.model == "vision-mid"
+        assert image_turn.routing_decision["cause"] == "modality_pin_override"
+        assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"]
+
+        assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}
+
+        text_turn = await router.async_pre_routing_hook(
+            model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}]
+        )
+        assert text_turn.model == "text-cheap"
+        assert text_turn.routing_decision["cause"] == "session_affinity_pin"
+
+    @pytest.mark.asyncio
+    async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance):
+        """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin."""
+        cache = AsyncMock()
+        cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
+        mock_router_instance.cache = cache
+        router = self._router(
+            mock_router_instance,
+            {
+                "tiers": {"SIMPLE": "text-cheap", "COMPLEX": "text-big"},
+                "modality_routing": True,
+                "modality_pin_override": True,
+                "session_affinity": True,
+            },
+            {"text-cheap": False, "text-big": False},
+        )
+        with pytest.raises(litellm.BadRequestError, match="no model"):
+            await router.async_pre_routing_hook(
+                model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE
+            )
+        assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}