Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_bedrock_bearer_skip_sigv4_chain

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

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/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/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/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/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_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/public_endpoints/autorouter_presets.json

@@ -17,6 +17,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -35,6 +36,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -63,6 +65,7 @@
       "classification_mode": "every_request",
       "session_affinity": false,
       "modality_routing": false,
+      "modality_pin_override": false,
       "deployment_affinity": true
     }
   },
@@ -84,6 +87,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

@@ -751,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")
@@ -760,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")
     )
@@ -2393,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
@@ -2407,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)
         ):
@@ -2449,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",
@@ -2458,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),

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

tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py

@@ -370,6 +370,7 @@ 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,

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/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/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/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/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/router_strategy/test_complexity_router.py

@@ -10797,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"),
         ],
     )
@@ -10847,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
@@ -10859,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 = []
@@ -11025,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"}

tests/test_rust_python_harness.py

@@ -8,14 +8,24 @@ import pytest
 
 catalog = importlib.import_module("tests.rust-python-harness.catalog")
 cli = importlib.import_module("tests.rust-python-harness.cli")
+ledger_module = importlib.import_module("tests.rust-python-harness.shared.parity.ledger")
+mapping_validator = importlib.import_module(
+    "tests.rust-python-harness.strategies.unit_tests.mapping_validator"
+)
 models = importlib.import_module("tests.rust-python-harness.models")
 runner = importlib.import_module("tests.rust-python-harness.runner")
 ui = importlib.import_module("tests.rust-python-harness.ui")
 
 load_catalog = catalog.load_catalog
+load_ledger = ledger_module.load_ledger
+ledger_path_for = mapping_validator.ledger_path_for
+REPO_ROOT = mapping_validator.REPO_ROOT
+audit_ledger = mapping_validator.audit_ledger
+build_function_report = mapping_validator.build_function_report
 _pick_values = cli._pick_values
 _coverage_pytest_args = cli._coverage_pytest_args
 _select = cli._select
+_validate_ledger = cli._validate_ledger
 CaseResult = models.CaseResult
 Coverage = models.Coverage
 HarnessCase = models.HarnessCase
@@ -56,13 +66,14 @@ def _manifest() -> dict[str, object]:
     }
 
 
-def test_should_load_the_three_harness_strategies_in_order() -> None:
+def test_should_load_the_four_harness_strategies_in_order() -> None:
     strategies = load_catalog()
 
     assert [strategy.id for strategy in strategies] == [
         "e2e_fuzz_tests",
         "unit_tests_rust",
         "validate_sub_methods",
+        "existing_e2e_test_sdk",
     ]
     assert all(
         tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS
@@ -94,6 +105,8 @@ def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> Non
             True,
         ),
         ("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False),
+        ("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True),
+        ("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False),
     ],
 )
 def test_should_match_pytest_file_and_node_selectors(
@@ -113,6 +126,13 @@ def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None:
     assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",)
 
 
+def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None:
+    (tmp_path / "tests" / "ocr_tests").mkdir(parents=True)
+    case = _case(selectors=("tests/ocr_tests/",))
+
+    assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",)
+
+
 def test_should_mark_planned_and_not_applicable_cases_without_running() -> None:
     planned = CaseResult(case=_case(coverage=Coverage.PLANNED))
     not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE))
@@ -229,8 +249,53 @@ def test_should_report_confidence_for_each_sdk_section() -> None:
     }
 
     assert scores["responses"].verified_strategies == 1
-    assert scores["responses"].required_strategies == 3
-    assert scores["responses"].percentage == 33
+    assert scores["responses"].required_strategies == 4
+    assert scores["responses"].percentage == 25
     assert scores["responses"].level.value == "MEDIUM"
     assert scores["count_tokens"].percentage == 0
     assert scores["count_tokens"].level.value == "LOW"
+
+
+
+def test_should_report_no_ledger_for_a_function_without_one() -> None:
+    report = build_function_report("messages", repo_root=REPO_ROOT)
+
+    assert report.has_ledger is False
+    assert report.is_clean is True
+
+
+def test_should_report_ocr_ledger_stats_and_a_clean_audit() -> None:
+    ledger = load_ledger(ledger_path_for("ocr"))
+
+    report = build_function_report("ocr", repo_root=REPO_ROOT)
+
+    assert report.has_ledger is True
+    assert report.ledger.mapped_count == ledger.mapped_count
+    assert report.ledger.total_count == ledger.total_count
+    assert report.is_clean is True
+
+
+def test_should_scope_validate_ledger_to_the_requested_function(
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    exit_code = _validate_ledger({"messages"})
+
+    captured = capsys.readouterr()
+    assert exit_code == 0
+    assert "messages" in captured.out
+    assert "no ledger yet" in captured.out
+    assert "ocr" not in captured.out
+
+
+def test_should_have_every_python_and_rust_ocr_test_accounted_for_in_the_ledger() -> None:
+    ledger = load_ledger(ledger_path_for("ocr"))
+
+    report = audit_ledger(ledger, repo_root=REPO_ROOT)
+
+    assert report.is_clean, (
+        "\nOCR test-parity ledger is out of sync with the live test files.\n"
+        f"Ledger references a Python test that no longer exists: {list(report.missing_python_tests)}\n"
+        f"Python test exists but is not tracked in the ledger: {list(report.stale_python_tests)}\n"
+        f"Ledger references a Rust test that no longer exists: {list(report.missing_rust_tests)}\n"
+        f"Rust test exists but is not tracked in the ledger: {list(report.stale_rust_tests)}\n"
+    )

type-discipline-budget.json

@@ -1,6 +1,6 @@
 {
   "LIT001": {
-    "limit": 22334
+    "limit": 22330
   },
   "LIT002": {
     "limit": 26763
@@ -27,7 +27,7 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16480
+    "limit": 16478
   },
   "LIT011": {
     "limit": 5520

ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx

@@ -880,6 +880,44 @@ describe("ComplexityRouterConfig modality panel", () => {
 
     expect(screen.getByRole("switch", { name: "Route image requests to vision-capable models" })).toBeChecked();
   });
+
+  // The backend ignores modality_pin_override unless modality_routing is on, so offering it while
+  // image routing is off would let an operator save a flag that does nothing.
+  it("disables the pin-override switch while image routing is off", () => {
+    const onChange = vi.fn();
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Modality Routing"));
+
+    const override = screen.getByRole("switch", { name: "Override session pin for image requests" });
+    expect(override).toHaveAttribute("aria-disabled", "true");
+    fireEvent.click(override);
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
+  it("writes modality_pin_override through onChange once image routing is on", () => {
+    const onChange = vi.fn();
+    const value = { ...defaultValue, modality_routing: true };
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} value={value} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Modality Routing"));
+
+    const override = screen.getByRole("switch", { name: "Override session pin for image requests" });
+    expect(override).not.toBeChecked();
+    fireEvent.click(override);
+
+    expect(onChange).toHaveBeenCalledWith({ ...value, modality_pin_override: true });
+  });
+
+  it("renders a stored modality_pin_override=true as on", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig
+        {...baseProps}
+        value={{ ...defaultValue, modality_routing: true, modality_pin_override: true }}
+      />,
+    );
+    fireEvent.click(screen.getByText("Advanced: Modality Routing"));
+
+    expect(screen.getByRole("switch", { name: "Override session pin for image requests" })).toBeChecked();
+  });
 });
 
 describe("ComplexityRouterConfig affinity panel", () => {

ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx

@@ -410,6 +410,7 @@ export interface ComplexityRouterConfigValue {
   classification_mode?: ClassificationMode;
   session_affinity?: boolean;
   modality_routing?: boolean;
+  modality_pin_override?: boolean;
   deployment_affinity?: boolean;
   /** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */
   plan_mode_min_tier?: string;

ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx

@@ -7,20 +7,36 @@ import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
 export const ModalityRoutingControls: React.FC<{
   value: ComplexityRouterConfigValue;
   onChange: (value: ComplexityRouterConfigValue) => void;
-}> = ({ value, onChange }) => (
-  <>
-    <div className="flex items-center gap-2 mb-2">
-      <Switch
-        checked={value.modality_routing ?? false}
-        onCheckedChange={(modalityRouting) => onChange({ ...value, modality_routing: modalityRouting })}
-        aria-label="Route image requests to vision-capable models"
-      />
-      <strong className="font-semibold">Route image requests to vision-capable models</strong>
-    </div>
-    <span className="block text-xs text-muted-foreground">
-      Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default
-      model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are replaced,
-      and a kept session pin still wins.
-    </span>
-  </>
-);
+}> = ({ value, onChange }) => {
+  const modalityRouting = value.modality_routing ?? false;
+  return (
+    <>
+      <div className="flex items-center gap-2 mb-2">
+        <Switch
+          checked={modalityRouting}
+          onCheckedChange={(nextModalityRouting) => onChange({ ...value, modality_routing: nextModalityRouting })}
+          aria-label="Route image requests to vision-capable models"
+        />
+        <strong className="font-semibold">Route image requests to vision-capable models</strong>
+      </div>
+      <span className="block text-xs mb-3 text-muted-foreground">
+        Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default
+        model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are
+        replaced, and a kept session pin still wins unless you turn on the override below.
+      </span>
+      <div className="flex items-center gap-2 mb-2">
+        <Switch
+          checked={value.modality_pin_override ?? false}
+          onCheckedChange={(modalityPinOverride) => onChange({ ...value, modality_pin_override: modalityPinOverride })}
+          disabled={!modalityRouting}
+          aria-label="Override session pin for image requests"
+        />
+        <strong className="font-semibold">Override session pin for image requests</strong>
+      </div>
+      <span className="block text-xs text-muted-foreground">
+        Route an image turn to a capable model even when the session is pinned to one that cannot take images. The pin
+        is kept, so the next text turn goes back to it. Needs image routing turned on.
+      </span>
+    </>
+  );
+};

ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx

@@ -590,6 +590,44 @@ describe("AddAutoRouterTab", () => {
     });
   });
 
+  it("writes both modality flags as false into the create payload when the panel stays untouched", async () => {
+    const user = userEvent.setup();
+    vi.mocked(getMissingTiersError).mockReturnValue(null);
+
+    renderWithProviders(<Harness />);
+
+    fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } });
+
+    await user.click(screen.getByRole("button", { name: /add auto router/i }));
+
+    await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
+    expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
+      modality_routing: false,
+      modality_pin_override: false,
+    });
+  });
+
+  it("carries the pin override through to the create payload once image routing unlocks it", async () => {
+    const user = userEvent.setup();
+    vi.mocked(getMissingTiersError).mockReturnValue(null);
+
+    renderWithProviders(<Harness />);
+
+    fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } });
+    expandDetailedConfiguration();
+    await user.click(screen.getByText("Advanced: Modality Routing"));
+    await user.click(await screen.findByRole("switch", { name: "Route image requests to vision-capable models" }));
+    await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" }));
+
+    await user.click(screen.getByRole("button", { name: /add auto router/i }));
+
+    await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
+    expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
+      modality_routing: true,
+      modality_pin_override: true,
+    });
+  });
+
   // Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset
   // rather than first.
   it("lists Custom Configuration after the bundled presets", () => {

ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx

@@ -358,6 +358,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
     classifierFallback: complexityRouterConfig.classifier_fallback,
     sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY,
     modalityRouting: complexityRouterConfig.modality_routing ?? false,
+    modalityPinOverride: complexityRouterConfig.modality_pin_override ?? false,
     deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
     customTechnicalKeywords,
     keywordTierRules,

ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts

@@ -56,6 +56,7 @@ describe("buildComplexityRouterConfig", () => {
       session_affinity: false,
       deployment_affinity: true,
       modality_routing: false,
+      modality_pin_override: false,
       escalation_keywords: ["LITELLM ESCALATE"],
     };
     expect(config).toEqual(expected);
@@ -270,6 +271,14 @@ describe("buildComplexityRouterConfig", () => {
     expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: false }).modality_routing).toBe(false);
   });
 
+  it("writes modality_pin_override explicitly both ways, so the stored config never relies on the backend default", () => {
+    expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: true }).modality_pin_override).toBe(true);
+    expect(buildComplexityRouterConfig(baseParams).modality_pin_override).toBe(false);
+    expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: false }).modality_pin_override).toBe(
+      false,
+    );
+  });
+
   it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => {
     const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true });
     expect(config.session_affinity).toBe(true);

ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts

@@ -111,6 +111,7 @@ export interface BuildComplexityRouterConfigParams {
   classificationMode: ClassificationMode | undefined;
   sessionAffinity: boolean;
   modalityRouting?: boolean;
+  modalityPinOverride?: boolean;
   deploymentAffinity: boolean;
   customTechnicalKeywords: string[];
   keywordTierRules: KeywordTierRule[];
@@ -169,6 +170,7 @@ export interface ComplexityRouterConfigPayload {
   session_affinity: boolean;
   deployment_affinity: boolean;
   modality_routing: boolean;
+  modality_pin_override: boolean;
   custom_technical_keywords?: string[];
   keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[];
   semantic_keyword_matching?: boolean;
@@ -409,6 +411,7 @@ export const buildComplexityRouterConfig = ({
   classificationMode,
   sessionAffinity,
   modalityRouting,
+  modalityPinOverride,
   deploymentAffinity,
   customTechnicalKeywords,
   keywordTierRules,
@@ -472,6 +475,7 @@ export const buildComplexityRouterConfig = ({
     session_affinity: sessionAffinity,
     deployment_affinity: deploymentAffinity,
     modality_routing: modalityRouting ?? false,
+    modality_pin_override: modalityPinOverride ?? false,
     ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
     ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
     escalation_keywords: cleanedEscalationKeywords,

ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts

@@ -257,6 +257,30 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
   });
 });
 
+describe("buildUpdatedComplexityRouterConfig modality pin override", () => {
+  it("writes modality_pin_override explicitly both ways", () => {
+    expect(
+      buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: true }).modality_pin_override,
+    ).toBe(true);
+    expect(
+      buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: false }).modality_pin_override,
+    ).toBe(false);
+  });
+
+  it("re-asserts the backend's off-by-default when the form value is absent, rather than dropping the key", () => {
+    const result = buildUpdatedComplexityRouterConfig({ ...STORED, modality_pin_override: true }, FORM_VALUE);
+    expect(result.modality_pin_override).toBe(false);
+  });
+
+  it("round-trips a stored modality_pin_override=true through hydrate then save", () => {
+    const stored = { ...STORED, modality_routing: true, modality_pin_override: true };
+    const hydrated = hydrateComplexityRouterConfig(stored, undefined);
+
+    expect(hydrated.modality_pin_override).toBe(true);
+    expect(buildUpdatedComplexityRouterConfig(stored, hydrated).modality_pin_override).toBe(true);
+  });
+});
+
 describe("buildUpdatedComplexityRouterConfig classification mode", () => {
   it("round-trips a stored user_turn through hydrate then save", () => {
     const stored = { ...STORED, classification_mode: "user_turn" };
@@ -505,6 +529,8 @@ describe("managed keys survive an untouched open-and-save", () => {
     classifier_fallback: "default_model",
     classification_mode: "user_turn",
     session_affinity: true,
+    modality_routing: true,
+    modality_pin_override: true,
     deployment_affinity: false,
     adaptive: true,
     adaptive_weights: { quality: 0.4, cost: 0.6 },

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts

@@ -51,6 +51,7 @@ const expectedClassifiedTierConfig = {
   session_affinity: false,
   deployment_affinity: true,
   modality_routing: false,
+  modality_pin_override: false,
   adaptive: true,
   adaptive_weights: { quality: 0.4, cost: 0.6 },
   adaptive_eligible: "classified_tier",
@@ -74,6 +75,7 @@ const expectedAdaptiveDisabledConfig = {
   session_affinity: false,
   deployment_affinity: true,
   modality_routing: false,
+  modality_pin_override: false,
 };
 
 describe("buildUpdatedComplexityRouterConfig", () => {
@@ -109,6 +111,26 @@ describe("buildUpdatedComplexityRouterConfig", () => {
     expect(disabled.modality_routing).toBe(false);
   });
 
+  it("hydrates a stored modality_pin_override into form state and defaults absent to off", () => {
+    expect(
+      hydrateComplexityRouterConfig({ ...storedConfig, modality_pin_override: true }, null).modality_pin_override,
+    ).toBe(true);
+    expect(hydrateComplexityRouterConfig(storedConfig, null).modality_pin_override).toBe(false);
+  });
+
+  it("round-trips modality_pin_override explicitly in both directions", () => {
+    const enabled = buildUpdatedComplexityRouterConfig(storedConfig, {
+      ...classifiedTierValue,
+      modality_pin_override: true,
+    });
+    expect(enabled.modality_pin_override).toBe(true);
+    const disabled = buildUpdatedComplexityRouterConfig(
+      { ...storedConfig, modality_pin_override: true },
+      { ...classifiedTierValue, modality_pin_override: false },
+    );
+    expect(disabled.modality_pin_override).toBe(false);
+  });
+
   it("includes return_raw_model_name only when enabled", () => {
     const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, {
       ...classifiedTierValue,

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx

@@ -549,6 +549,50 @@ describe("EditAutoRouterModal deployment affinity", () => {
     await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
     expect(savedConfig().deployment_affinity).toBe(false);
   });
+
+  // modality_pin_override is a managed key, so the modal rewrites it from form state on save. A
+  // hydration gap would silently turn a stored override off on the next untouched save.
+  it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true, modality_pin_override: true });
+
+    await user.click(await screen.findByText("Advanced: Modality Routing"));
+    expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toBeChecked();
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().modality_pin_override).toBe(true);
+  });
+
+  it("persists turning the modality pin override on", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true });
+
+    await user.click(await screen.findByText("Advanced: Modality Routing"));
+    await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" }));
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().modality_pin_override).toBe(true);
+  });
+
+  it("writes modality_pin_override=false for a stored config that never carried the key", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig(STORED_CONFIG);
+
+    await user.click(await screen.findByText("Advanced: Modality Routing"));
+    expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toHaveAttribute(
+      "aria-disabled",
+      "true",
+    );
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().modality_pin_override).toBe(false);
+  });
 });
 
 describe("EditAutoRouterModal custom classifier prompt and fallback", () => {

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx

@@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig {
   reasoning_override_min_score?: unknown;
   session_affinity?: unknown;
   modality_routing?: unknown;
+  modality_pin_override?: unknown;
   deployment_affinity?: unknown;
   adaptive?: boolean;
   adaptive_weights?: AdaptiveRouterWeights;
@@ -181,6 +182,8 @@ export const hydrateComplexityRouterConfig = (
     session_affinity:
       typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
     modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false,
+    modality_pin_override:
+      typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false,
     deployment_affinity:
       typeof parsedConfig.deployment_affinity === "boolean"
         ? parsedConfig.deployment_affinity
@@ -221,6 +224,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
   "classification_mode",
   "session_affinity",
   "modality_routing",
+  "modality_pin_override",
   "deployment_affinity",
   "adaptive",
   "adaptive_weights",
@@ -318,6 +322,7 @@ export const buildUpdatedComplexityRouterConfig = (
     classifierFallback: value.classifier_fallback,
     sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY,
     modalityRouting: value.modality_routing ?? false,
+    modalityPinOverride: value.modality_pin_override ?? false,
     deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
     customTechnicalKeywords: customTechnicalKeywords ?? [],
     keywordTierRules: keywordMatching?.keywordTierRules ?? [],

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx

@@ -186,6 +186,12 @@ describe("RoutingDecisionCard", () => {
     expect(screen.queryByText("housekeeping")).not.toBeInTheDocument();
   });
 
+  it("labels a modality pin override instead of showing the raw cause token", () => {
+    render(<RoutingDecisionCard decision={{ ...heuristic, cause: "modality_pin_override" }} />);
+    expect(screen.getByText("Overrode session pin for image input")).toBeInTheDocument();
+    expect(screen.queryByText("modality_pin_override")).not.toBeInTheDocument();
+  });
+
   it("labels a modality escalation instead of showing the raw cause token", () => {
     render(<RoutingDecisionCard decision={{ ...heuristic, cause: "modality_escalation" }} />);
     expect(screen.getByText("Escalated for image input")).toBeInTheDocument();

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx

@@ -93,6 +93,7 @@ const CONSTANT_CAUSE_LABELS: Record<string, string> = {
   session_affinity_escalation: "Escalated from session pin",
   user_turn_continuation: "Continuation turn, classifier skipped",
   modality_escalation: "Escalated for image input",
+  modality_pin_override: "Overrode session pin for image input",
   quality_tier: "Quality tier mapping",
   bandit: "Adaptive bandit",
   default_fallback: "Default model, no route matched",

ui/litellm-dashboard/src/lib/autorouter_presets.test.ts

@@ -148,6 +148,25 @@ describe("autorouter_presets", () => {
     expect(withoutFlag.complexityRouterConfig.modality_routing).toBe(false);
   });
 
+  it("carries a preset's modality_pin_override into the prefilled form state", () => {
+    const preset = getPresetByKey("anthropic_family")!;
+    const withFlag = { ...preset.complexity_router_config, modality_routing: true, modality_pin_override: true };
+    const prefill = buildPresetPrefill(withFlag, groupsOnly(getRequiredModelsInPreset(preset)));
+    expect(prefill.complexityRouterConfig.modality_pin_override).toBe(true);
+    const withoutFlag = buildPresetPrefill(
+      preset.complexity_router_config,
+      groupsOnly(getRequiredModelsInPreset(preset)),
+    );
+    expect(withoutFlag.complexityRouterConfig.modality_pin_override).toBe(false);
+  });
+
+  it("ships every bundled preset with both modality flags written out, since the payload type requires them", () => {
+    for (const preset of getAllPresets()) {
+      expect(preset.complexity_router_config.modality_routing, preset.key).toBe(false);
+      expect(preset.complexity_router_config.modality_pin_override, preset.key).toBe(false);
+    }
+  });
+
   it("prefills the anthropic preset's effort through to tier_model_params", () => {
     const preset = getPresetByKey("anthropic_family")!;
     const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset)));

ui/litellm-dashboard/src/lib/autorouter_presets.ts

@@ -286,6 +286,7 @@ export const buildPresetPrefill = (
       session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
       deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
       modality_routing: config.modality_routing ?? false,
+      modality_pin_override: config.modality_pin_override ?? false,
       adaptive: config.adaptive,
       adaptive_weights: config.adaptive_weights,
       tier_distance_penalty: config.tier_distance_penalty,

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

@@ -34740,9 +34740,15 @@ export interface components {
              * @default 0.5
              */
             match_threshold: number;
+            /**
+             * Modality Pin Override
+             * @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.
+             * @default false
+             */
+            modality_pin_override: boolean;
             /**
              * Modality Routing
-             * @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, 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.
+             * @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, 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, unless modality_pin_override is also enabled.
              * @default false
              */
             modality_routing: boolean;
@@ -35945,7 +35951,7 @@ export interface components {
              * Cause
              * @enum {string}
              */
-            cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
+            cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
             /** Classifier Cost */
             classifier_cost?: number;
             /** Classifier Model */