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

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

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

basedpyright-code-budget.json

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

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/__init__.py

@@ -424,6 +424,10 @@ anthropic_beta_headers_url: str = os.getenv(
     "LITELLM_ANTHROPIC_BETA_HEADERS_URL",
     "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json",
 )
+autorouter_presets_url: str = os.getenv(
+    "LITELLM_AUTOROUTER_PRESETS_URL",
+    "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json",
+)
 suppress_debug_info: bool = False
 dynamodb_table_name: Optional[str] = None
 s3_callback_params: Optional[Dict] = None

litellm/constants.py

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

litellm/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/integrations/websearch_interception/handler.py

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

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

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

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

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

litellm/llms/base_llm/guardrail_translation/base_translation.py

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

litellm/llms/base_llm/guardrail_translation/utils.py

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

litellm/llms/fireworks_ai/common_utils.py

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

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

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

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

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

litellm/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/outbound_credentials/token_endpoint.py

@@ -121,6 +121,23 @@ class TokenEndpointClient:
         return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in))
 
 
+class _KeyGuard:
+    """The per-key single-flight lock plus the invalidation generation that lock protects.
+
+    Both live on one object so their lifetimes cannot diverge. `get_or_compute` binds the guard to
+    a local for its whole critical section, which keeps the weak map's entry alive for as long as
+    that compute could still write; an `invalidate` overlapping the compute therefore reaches the
+    very same object and its bump is guaranteed to be observed. Conversely a guard nobody holds is
+    collectible precisely because no write is outstanding for it to fence.
+    """
+
+    __slots__ = ("__weakref__", "generation", "lock")
+
+    def __init__(self) -> None:
+        self.lock = asyncio.Lock()
+        self.generation = 0
+
+
 class ExchangedTokenCache:
     """Memoizes the final token string per key, single-flighting concurrent misses on one lock."""
 
@@ -129,7 +146,7 @@ class ExchangedTokenCache:
             max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
             default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
         )
-        self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
+        self._guards: weakref.WeakValueDictionary[str, _KeyGuard] = weakref.WeakValueDictionary()
 
     async def get_or_compute(
         self,
@@ -144,28 +161,50 @@ class ExchangedTokenCache:
         guaranteeing the token it gets back was minted for the *current* inputs: a stored entry
         whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction
         addressable without the key having to encode the credential material it protects.
+
+        An `invalidate` landing while `compute` is in flight wins over that compute's write. The
+        token is still returned to the caller it was minted for, but it is not stored, so the next
+        resolution re-mints rather than serving a bearer that predates the invalidation for the
+        rest of its TTL.
         """
         cached = self._get(cache_key, fingerprint)
         if cached is not None:
             return Ok(cached)
-        async with self._lock(cache_key):
+        guard = self._guard(cache_key)
+        async with guard.lock:
             cached = self._get(cache_key, fingerprint)
             if cached is not None:
                 return Ok(cached)
+            generation = guard.generation
             match await compute():
                 case Ok(token):
-                    self._cache.set_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
-                        cache_key,
-                        (fingerprint, token.access_token),
-                        ttl=_cache_ttl_seconds(token.expires_in),
-                    )
+                    if guard.generation == generation:
+                        self._store(cache_key, fingerprint, token)
                     return Ok(token.access_token)
                 case Error(err):
                     return Error(err)
 
     def invalidate(self, cache_key: str) -> None:
-        """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401)."""
+        """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).
+
+        Bumping the guard's generation is what makes the eviction stick against a compute already
+        awaiting the token endpoint: that compute snapshotted the old generation and so skips its
+        write. No guard means no compute is in flight, since an in-flight one pins its own.
+
+        Stays synchronous: callers invalidate from plain `def`s.
+        """
         self._cache.delete_cache(cache_key)  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+        guard = self._guards.get(cache_key)
+        if guard is None:
+            return
+        guard.generation += 1
+
+    def _store(self, cache_key: str, fingerprint: str, token: ExchangedToken) -> None:
+        self._cache.set_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            cache_key,
+            (fingerprint, token.access_token),
+            ttl=_cache_ttl_seconds(token.expires_in),
+        )
 
     def _get(self, cache_key: str, fingerprint: str) -> str | None:
         """The stored token, or None when absent or minted for different inputs.
@@ -180,12 +219,12 @@ class ExchangedTokenCache:
             return None
         return token if stored_fingerprint == fingerprint else None
 
-    def _lock(self, cache_key: str) -> asyncio.Lock:
-        lock = self._locks.get(cache_key)
-        if lock is None:
-            lock = asyncio.Lock()
-            self._locks[cache_key] = lock
-        return lock
+    def _guard(self, cache_key: str) -> _KeyGuard:
+        guard = self._guards.get(cache_key)
+        if guard is None:
+            guard = _KeyGuard()
+            self._guards[cache_key] = guard
+        return guard
 
 
 def _cache_ttl_seconds(expires_in: int | None) -> int:

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
@@ -3855,11 +3861,19 @@ if MCP_AVAILABLE:
                 and server.auth_type == MCPAuth.oauth2_token_exchange
                 and oauth2_headers
                 and len(mcp_servers or []) == 1
+                and server.server_id
+                in frozenset(
+                    allowed.server_id
+                    for allowed in await _get_allowed_mcp_servers(
+                        user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
+                    )
+                )
             ):
                 await global_mcp_server_manager.preflight_token_exchange(
                     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
@@ -4188,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
@@ -4511,6 +4526,7 @@ if MCP_AVAILABLE:
                 user_api_key_auth=user_api_key_auth,
                 client_ip=_sse_client_ip,
                 allowed_server_ids=toolset_allowed_server_ids,
+                raw_headers=raw_headers,
             )
 
             # Pre-flight auth check for pass-through servers: surface upstream

litellm/proxy/_experimental/mcp_server/tool_search.py

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

litellm/proxy/_types.py

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

litellm/proxy/agent_endpoints/endpoints.py

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

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

@@ -100,9 +100,10 @@ class CliPollData(TypedDict, total=False):
 
 
 class CliSsoStartData(TypedDict):
-    login_id: str
-    poll_secret: str
-    user_code: str
+    login_id: ReadOnly[str]
+    poll_secret: ReadOnly[str]
+    user_code: ReadOnly[str]
+    verification_uri_complete: ReadOnly[NotRequired[str]]
 
 
 class CliAuthResult(TypedDict):
@@ -860,11 +861,22 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
         poll_secret: Final = cli_sso_flow["poll_secret"]
         user_code: Final = cli_sso_flow["user_code"]
 
-        sso_url = f"{base_url}/sso/key/generate?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id})
+        browser_prefills_code: Final = isinstance(cli_sso_flow.get("verification_uri_complete"), str)
+        sso_url: Final = f"{base_url}/sso/key/generate?" + urlencode(
+            (
+                ("source", LITELLM_CLI_SOURCE_IDENTIFIER),
+                ("key", key_id),
+                *((("user_code", user_code),) if browser_prefills_code else ()),
+            )
+        )
 
         click.echo(f"Opening browser to: {sso_url}")
         click.echo("Please complete the SSO authentication in your browser...")
-        click.echo(f"Verification code: {user_code}")
+        click.echo(
+            f"Verification code: {user_code} (pre-filled in the browser, check it matches)"
+            if browser_prefills_code
+            else f"Verification code: {user_code}"
+        )
         click.echo(f"Session ID: {key_id}")
 
         # Open browser

litellm/proxy/common_utils/reset_budget_job.py

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

litellm/proxy/guardrails/anthropic_sse.py

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

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -514,12 +514,26 @@ async def update_guardrail(
         guardrail_name: Final = result.get("guardrail_name", "Unknown")
 
         try:
-            IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
-                guardrail_id=guardrail_id, guardrail=cast(Guardrail, result)
-            )
+            IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result))
             verbose_proxy_logger.info(
                 "Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
             )
+        except (ValueError, TypeError) as update_error:
+            # The new config is invalid (a raising guardrail __init__):
+            # reinitialize_guardrail already restored the previous live instance, but
+            # update_guardrail_in_db above already persisted the rejected config to
+            # the DB. Roll that back too, so the DB and the live guardrail never
+            # disagree about what's actually enforcing, and surface the rejection to
+            # the caller instead of a misleading 200.
+            await GUARDRAIL_REGISTRY.update_guardrail_in_db(
+                guardrail_id=guardrail_id,
+                guardrail=existing_guardrail,
+                prisma_client=prisma_client,
+            )
+            raise HTTPException(
+                status_code=422,
+                detail=f"Invalid guardrail configuration, update rejected: {update_error}",
+            ) from update_error
         except Exception as update_error:
             verbose_proxy_logger.warning(
                 "Immediate sync: Failed to update '%s' (ID: %s) in memory: %s",

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py

@@ -1,8 +1,8 @@
 from typing import TYPE_CHECKING, Final
 
-from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
+from litellm.types.guardrails import SupportedGuardrailIntegrations
 
-from .crowdstrike_aidr import CrowdStrikeAIDRHandler
+from .crowdstrike_aidr import CrowdStrikeAIDRHandler, streaming_params_from_litellm_params
 
 if TYPE_CHECKING:
     from litellm.types.guardrails import Guardrail, LitellmParams
@@ -15,17 +15,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
     if not guardrail_name:
         raise ValueError("CrowdStrike AIDR guardrail name is required")
 
+    streaming_params: Final = streaming_params_from_litellm_params(litellm_params)
     _crowdstrike_aidr_callback: Final = CrowdStrikeAIDRHandler(
         guardrail_name=guardrail_name,
         api_base=litellm_params.api_base,
         api_key=litellm_params.api_key,
-        # Exclude during_call to prevent duplicate input events
-        event_hook=[
-            GuardrailEventHooks.pre_call.value,
-            GuardrailEventHooks.post_call.value,
-        ],
+        event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
         fail_on_error=litellm_params.fail_on_error,
+        streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
+        streaming_sampling_rate=streaming_params.streaming_sampling_rate,
     )
     litellm.logging_callback_manager.add_litellm_callback(_crowdstrike_aidr_callback)
 

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py

@@ -24,8 +24,11 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy.common_utils.callback_utils import (
     add_guardrail_to_applied_guardrails_header,
 )
-from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
 from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionToolParam
+from litellm.types.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import (
+    CrowdStrikeAIDRGuardrailConfigModelOptionalParams,
+)
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
@@ -153,6 +156,21 @@ def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] |
     return merged if present else None
 
 
+def streaming_params_from_litellm_params(
+    litellm_params: LitellmParams,
+) -> CrowdStrikeAIDRGuardrailConfigModelOptionalParams:
+    extras: Final[Mapping[str, object]] = litellm_params.model_extra or {}
+    nested: Final = litellm_params.optional_params
+    optional_params: Final[Mapping[str, object]] = {} if nested is None else nested.model_dump()
+    return CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_validate(
+        {
+            name: value
+            for name in CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_fields
+            if (value := optional_params.get(name, extras.get(name))) is not None
+        }
+    )
+
+
 def _messages_since_last_assistant(
     messages: Sequence[AllMessageValues],
 ) -> _FilteredMessages:
@@ -241,6 +259,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         api_key: str | None = None,
         api_base: str | None = None,
         fail_on_error: bool | None = True,
+        streaming_end_of_stream_only: bool | None = None,
+        streaming_sampling_rate: int | None = None,
         **kwargs,
     ) -> None:
         """
@@ -250,10 +270,19 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             guardrail_name (str): The name of the guardrail instance.
             api_key (str | None): The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None.
             api_base (str | None): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None.
+            streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of
+                every streaming_sampling_rate chunks. Defaults to False.
+            streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5.
             **kwargs: Additional arguments passed to the CustomGuardrail base class.
         """
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
         self.fail_on_error = True if fail_on_error is None else fail_on_error
+        self._set_streaming_params(
+            CrowdStrikeAIDRGuardrailConfigModelOptionalParams(
+                streaming_end_of_stream_only=streaming_end_of_stream_only,
+                streaming_sampling_rate=streaming_sampling_rate,
+            )
+        )
 
         self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN")
         if not self.api_key:
@@ -274,6 +303,15 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
             "Initialized CrowdStrike AIDR Guardrail: name=%s, api_base=%s", guardrail_name, self.api_base
         )
 
+    def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None:
+        self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False
+        self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5
+
+    @override
+    def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
+        super().update_in_memory_litellm_params(litellm_params)
+        self._set_streaming_params(streaming_params_from_litellm_params(litellm_params))
+
     async def _call_crowdstrike_aidr_guard(
         self, payload: dict[str, Any], hook_name: str
     ) -> _GuardChatCompletionsResult:

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

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

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

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

litellm/proxy/guardrails/guardrail_registry.py

@@ -826,11 +826,12 @@ class InMemoryGuardrailHandler:
         Removes old callback from litellm.callbacks and creates fresh instance.
 
         If the new config fails to initialize (e.g. an invalid on_flagged
-        combination), the previous instance is restored rather than left
-        deleted: initialize_guardrail's own ValueError/TypeError propagate
-        uncaught, so a caller reaching this point after already deleting the
-        old instance would otherwise leave the guardrail providing no
-        protection at all, not merely "still enforcing the old config."
+        combination or an invalid regex), the previous instance is restored
+        rather than left deleted, and the failure is re-raised as ValueError so
+        every init failure reaches callers as one exception type: a caller
+        reaching this point after already deleting the old instance would
+        otherwise leave the guardrail providing no protection at all, not
+        merely "still enforcing the old config."
         """
         guardrail_id: Final = guardrail.get("guardrail_id")
         if not guardrail_id:
@@ -849,7 +850,7 @@ class InMemoryGuardrailHandler:
         # that was enforcing must never fail open because an update was bad.
         try:
             return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
-        except Exception:
+        except Exception as init_error:
             if previous_guardrail is not None:
                 verbose_proxy_logger.exception(
                     "Reinitializing guardrail %s with updated params failed; restoring the previous configuration",
@@ -861,7 +862,7 @@ class InMemoryGuardrailHandler:
                     )
                 except Exception:  # noqa: BLE001  # the original failure must propagate even if the restore breaks
                     verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id)
-            raise
+            raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error
 
     def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None:
         """

litellm/proxy/litellm_pre_call_utils.py

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

litellm/proxy/openai_files_endpoints/batch_file_validation.py

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

litellm/proxy/openai_files_endpoints/general_upload_validation.py

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

litellm/proxy/public_endpoints/autorouter_presets.json

@@ -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/proxy/public_endpoints/public_endpoints.py

@@ -1,11 +1,13 @@
+import asyncio
 import json
 import os
 import re
-from collections.abc import Awaitable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from importlib.resources import files
 from typing import TYPE_CHECKING, Final, Protocol
 
 from fastapi import APIRouter, HTTPException, Request
+from pydantic import TypeAdapter
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -28,6 +30,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import
 )
 from litellm.types.proxy.public_endpoints.public_endpoints import (
     AgentCreateInfo,
+    AutoRouterPresetRecord,
     ComplexityScorerDefaults,
     ProviderCreateInfo,
     PublicModelHubInfo,
@@ -464,6 +467,86 @@ async def get_litellm_blog_posts():
     return BlogPostsResponse(posts=posts)
 
 
+_AUTOROUTER_PRESETS_ADAPTER: Final = TypeAdapter(dict[str, AutoRouterPresetRecord])
+
+
+def _load_bundled_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]:
+    raw: Final = json.loads(
+        files("litellm.proxy.public_endpoints").joinpath("autorouter_presets.json").read_text(encoding="utf-8")
+    )
+    return _AUTOROUTER_PRESETS_ADAPTER.validate_python(raw)
+
+
+async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]:
+    from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+    from litellm.types.llms.custom_http import httpxSpecialProvider
+
+    client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.UI)
+    response: Final = await client.get(url, timeout=5.0)
+    response.raise_for_status()
+    presets: Final = _AUTOROUTER_PRESETS_ADAPTER.validate_python(response.json())
+    if not presets:
+        raise ValueError("remote auto-router preset catalog is empty")
+    return presets
+
+
+async def _resolve_autorouter_presets(
+    url: str,
+    fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]],
+) -> Mapping[str, AutoRouterPresetRecord]:
+    if os.getenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "").lower() == "true":
+        return _load_bundled_autorouter_presets()
+    try:
+        return await fetch(url)
+    except Exception as e:
+        verbose_logger.warning(
+            "LiteLLM: failed to fetch auto-router presets from %s: %s. Serving the bundled catalog for the life of this process.",
+            url,
+            str(e),
+        )
+        return _load_bundled_autorouter_presets()
+
+
+class _AutoRouterPresetsCache:
+    presets: Mapping[str, AutoRouterPresetRecord] | None = None
+    lock: asyncio.Lock | None = None
+
+
+async def get_autorouter_presets(
+    url: str,
+    fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]] = _fetch_remote_autorouter_presets,
+) -> Mapping[str, AutoRouterPresetRecord]:
+    cached: Final = _AutoRouterPresetsCache.presets
+    if cached is not None:
+        return cached
+    if _AutoRouterPresetsCache.lock is None:
+        _AutoRouterPresetsCache.lock = asyncio.Lock()
+    async with _AutoRouterPresetsCache.lock:
+        held: Final = _AutoRouterPresetsCache.presets
+        if held is not None:
+            return held
+        resolved: Final = await _resolve_autorouter_presets(url=url, fetch=fetch)
+        _AutoRouterPresetsCache.presets = resolved
+        return resolved
+
+
+@router.get(
+    "/public/autorouter_presets",
+    tags=["public", "auto router"],  # mutable-ok: FastAPI route tags take a list
+    response_model=dict[str, AutoRouterPresetRecord],
+)
+async def get_public_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]:
+    """
+    Return the auto-router preset catalog the dashboard's template picker renders.
+
+    Resolved once per process, like the model cost map: fetched from ``litellm.autorouter_presets_url``
+    (override with ``LITELLM_AUTOROUTER_PRESETS_URL``) on the first request, falling back to the
+    catalog bundled with the package on any failure. Set ``LITELLM_LOCAL_AUTOROUTER_PRESETS=True``
+    to serve the bundled catalog only. A restart picks up a newly published catalog.
+    """
+    return await get_autorouter_presets(url=litellm.autorouter_presets_url)
+
+
 @router.get(
     "/public/endpoints",
     tags=["public"],

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -55,6 +55,10 @@ router: Final = APIRouter()
 
 SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
 
+_SESSION_GROUP_KEY_SQL: Final = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
+_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
+_AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
+
 _INTERNAL_HEALTH_CHECK_API_KEYS: Final = (
     LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
     hash_token(token=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME),
@@ -144,21 +148,16 @@ class _DailyTagSpendRow(TypedDict):
     total_spend: float
 
 
-class _SessionCountAggregate(TypedDict):
-    session_id: int
-
-
-class _SessionCountRow(TypedDict):
-    session_id: str
-    _count: _SessionCountAggregate
-
-
 class _SessionSpendRow(TypedDict):
     session_id: str
+    api_key: ReadOnly[str]
+    session_total_count: ReadOnly[int]
     session_total_spend: float
     mcp_tool_call_count: int
     mcp_tool_call_spend: float
     session_cache_hit_count: ReadOnly[int]
+    session_llm_count: ReadOnly[int]
+    session_agent_count: ReadOnly[int]
 
 
 class _SpendSumAggregate(TypedDict, total=False):
@@ -242,18 +241,6 @@ async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, obj
     return await _spend_logs_table(prisma_client).count(where=where)
 
 
-async def _count_logs_per_session(
-    prisma_client: PrismaClient, session_ids: Sequence[str | None]
-) -> Sequence[_SessionCountRow]:
-    """Count spend log rows per session for the given session ids."""
-    rows: Final = await _spend_logs_table(prisma_client).group_by(
-        by=["session_id"],
-        where={"session_id": {"in": session_ids}},
-        count={"session_id": True},
-    )
-    return cast(Sequence[_SessionCountRow], rows)  # cast-ok: group_by(count=) shape is fixed by the by/count args
-
-
 async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None:
     """Read a single team row as a Prisma model instance."""
     return await _team_table(prisma_client).find_unique(where={"team_id": team_id})
@@ -2290,6 +2277,10 @@ async def ui_view_spend_logs(
         default=False,
         description="Exclude LiteLLM internal health check requests from results",
     ),
+    group_by_session: bool = fastapi.Query(
+        default=False,
+        description="Paginate over sessions instead of raw logs: one representative row per session, total counts sessions",
+    ),
 ):
     """
     View spend logs with pagination support.
@@ -2644,12 +2635,16 @@ async def ui_view_spend_logs(
         else:
             _order_expr = order_column
 
+        joined_conditions: Final = " AND ".join(sql_conditions)
+        session_grouping: Final = group_by_session is True
+        count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else ""
         count_query: Final = f"""
             SELECT COUNT(*) AS total_count
             FROM (
                 SELECT 1
                 FROM "LiteLLM_SpendLogs"
-                WHERE {" AND ".join(sql_conditions)}
+                WHERE {joined_conditions}
+                {count_group_clause}
                 LIMIT ${p}
             ) AS bounded_matches
         """
@@ -2660,21 +2655,36 @@ async def ui_view_spend_logs(
         total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
         total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
 
-        sql_query: Final = f"""
-            SELECT
-                request_id, call_type, api_key, spend, total_tokens,
+        select_columns: Final = """request_id, call_type, api_key, spend, total_tokens,
                 prompt_tokens, completion_tokens, "startTime", "endTime",
                 "completionStartTime", model, model_id, model_group,
                 custom_llm_provider, api_base, "user", metadata,
                 cache_hit, cache_key, request_tags, team_id,
                 organization_id, end_user, requester_ip_address,
                 session_id, status, mcp_namespaced_tool_name, agent_id,
-                COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms
+                COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms"""
+        sql_query: Final = (
+            f"""
+                SELECT * FROM (
+                    SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
+                        {select_columns}
+                    FROM "LiteLLM_SpendLogs"
+                    WHERE {joined_conditions}
+                    ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
+                ) AS session_representatives
+                ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}, request_id
+                LIMIT ${p} OFFSET ${p + 1}
+            """
+            if session_grouping
+            else f"""
+            SELECT
+                {select_columns}
             FROM "LiteLLM_SpendLogs"
-            WHERE {" AND ".join(sql_conditions)}
+            WHERE {joined_conditions}
             ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}
             LIMIT ${p} OFFSET ${p + 1}
         """
+        )
         sql_params.extend([page_size, skip])
 
         data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
@@ -4075,11 +4085,12 @@ async def _build_ui_spend_logs_response(
     Build the paginated response for the UI spend-logs endpoint.
 
     When ``enrich_session_counts`` is ``True`` (the default for the v1/UI
-    endpoint), each row is enriched with ``session_total_count`` so the
-    frontend knows which sessions are expandable (multi-call sessions).
-    For every row that carries a ``session_id``, a single ``GROUP BY`` query
-    fetches the total number of logs in each referenced session.  Rows without
-    a ``session_id`` default to ``1``.
+    endpoint), each row is enriched with ``session_total_count`` plus spend
+    and call-type aggregates so the frontend knows which sessions are
+    expandable (multi-call sessions).  One ``GROUP BY (session_id, api_key)``
+    query serves every referenced session, keyed per api key so two callers
+    reusing a session id never see each other's totals.  Rows without a
+    ``session_id`` default to ``1``.
 
     When ``enrich_session_counts`` is ``False`` (v2 endpoint), rows are
     serialised without the extra query.
@@ -4101,7 +4112,6 @@ async def _build_ui_spend_logs_response(
         A dict with ``data`` (enriched rows), ``total``, ``page``,
         ``page_size``, ``total_pages``, and ``total_is_capped``.
     """
-    count_map: dict[str, int] = {}
     if enrich_session_counts:
         session_ids: Final[Sequence[str | None]] = list(
             {
@@ -4110,15 +4120,8 @@ async def _build_ui_spend_logs_response(
                 if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None))
             }
         )
-        if session_ids:
-            # NOTE: This GROUP BY runs on every v1/UI page load. The IN clause
-            # is bounded by page_size (typically 25-50 distinct session IDs).
-            # If performance degrades at scale, consider short-lived caching or
-            # folding the count into the main query via a window function.
-            counts: Final = await _count_logs_per_session(prisma_client, session_ids)
-            count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")}
 
-    session_spend_map: dict[str, dict[str, int | float]] = {}
+    session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {}
     if enrich_session_counts and session_ids:
         from prisma.errors import PrismaError
 
@@ -4130,38 +4133,46 @@ async def _build_ui_spend_logs_response(
                 {
                     (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None))
                     for row in data
-                    if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None))
+                    if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None
                 }
             )
             rows: Final[Sequence[_SessionSpendRow]] = await _query_raw(
                 prisma_client,
-                """
-                SELECT session_id,
+                f"""
+                SELECT session_id, api_key,
+                       COUNT(*)::int AS session_total_count,
                        COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
                        COUNT(*) FILTER (
-                           WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools')
+                           WHERE call_type IN {_MCP_CALL_TYPES_SQL}
                        )::int AS mcp_tool_call_count,
                        COALESCE(SUM(spend) FILTER (
-                           WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools')
+                           WHERE call_type IN {_MCP_CALL_TYPES_SQL}
                        ), 0)::double precision AS mcp_tool_call_spend,
-                       COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count
+                       COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
+                       COUNT(*) FILTER (
+                           WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
+                       )::int AS session_llm_count,
+                       COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
                 FROM "LiteLLM_SpendLogs"
                 WHERE session_id = ANY($1::text[])
                   AND api_key = ANY($2::text[])
-                GROUP BY session_id
+                GROUP BY session_id, api_key
                 """,
                 session_ids,
                 authorized_api_keys,
             )
             session_spend_map = {
-                row["session_id"]: {
+                (row["session_id"], row["api_key"]): {
+                    "session_total_count": int(row.get("session_total_count") or 0),
                     "session_total_spend": float(row.get("session_total_spend") or 0.0),
                     "mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0),
                     "mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0),
                     "session_cache_hit_count": int(row.get("session_cache_hit_count") or 0),
+                    "session_llm_count": int(row.get("session_llm_count") or 0),
+                    "session_agent_count": int(row.get("session_agent_count") or 0),
                 }
                 for row in rows
-                if row.get("session_id")
+                if row.get("session_id") and row.get("api_key") is not None
             }
         except PrismaError:
             verbose_proxy_logger.debug(
@@ -4174,14 +4185,17 @@ async def _build_ui_spend_logs_response(
         for row in data:
             row_dict = dict(row) if isinstance(row, dict) else row.model_dump()
             sid = row_dict.get("session_id")
-            row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1
-            session_stats = session_spend_map.get(sid) if sid else None
+            row_api_key = row_dict.get("api_key")
+            session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None
+            row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1
             if session_stats:
                 row_dict["session_total_spend"] = session_stats["session_total_spend"]
                 if session_stats["mcp_tool_call_count"]:
                     row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"]
                     row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"]
                 row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"]
+                row_dict["session_llm_count"] = session_stats["session_llm_count"]
+                row_dict["session_agent_count"] = session_stats["session_agent_count"]
             enriched.append(row_dict)
         response_data: list = enriched
     else:

litellm/responses/streaming_iterator.py

@@ -17,6 +17,7 @@ from typing_extensions import TypeIs
 
 import litellm
 from litellm.constants import (
+    EMPTY_MAPPING,
     LITELLM_MAX_STREAMING_DURATION_SECONDS,
     STREAM_SSE_DONE_STRING,
 )
@@ -273,6 +274,9 @@ class BaseResponsesAPIStreamingIterator:
         self._hidden_params["additional_headers"] = process_response_headers(
             self.response.headers or {}
         )  # GUARANTEE OPENAI HEADERS IN RESPONSE
+        self._raw_response_headers: Mapping[str, str] = MappingProxyType(
+            dict(self.response.headers or {})  # mutable-ok: immediately frozen by MappingProxyType
+        )
 
     def _check_max_streaming_duration(self) -> None:
         """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
@@ -446,6 +450,7 @@ class BaseResponsesAPIStreamingIterator:
             except Exception:
                 # Fallback to original if serialization fails
                 pass
+        self._restore_provider_response_headers(logging_response)
 
         end_time: Final = datetime.now()
         if is_async:
@@ -480,6 +485,41 @@ class BaseResponsesAPIStreamingIterator:
             )
         self._run_post_success_hooks(end_time=end_time)
 
+    def _restore_provider_response_headers(self, logging_response: object) -> None:
+        """Re-apply the provider's response headers to the copy handed to logging callbacks.
+
+        ``model_validate(model_dump())`` above drops pydantic private attributes, so the
+        ``_hidden_params`` the provider transform set on the nested response are lost. Returns early
+        when that copy fell back to the original event, so logging-only state never lands on the
+        object the caller is iterating.
+        """
+        if logging_response is self.completed_response:
+            return
+        target: Final[object] = getattr(logging_response, "response", None)
+        existing_hidden: Final[object] = getattr(target, "_hidden_params", None)
+        if not isinstance(existing_hidden, Mapping):
+            return
+        existing: Final[Mapping[str, object]] = existing_hidden
+        source_hidden: Final[object] = getattr(
+            getattr(self.completed_response, "response", None), "_hidden_params", None
+        )
+        source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING
+        processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers")
+        raw: Final[object] = source.get("headers") or self._raw_response_headers
+        headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING
+        raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING
+        # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy
+        # splats into the client's HTTP headers, and copying non-header keys would carry response_cost
+        setattr(  # noqa: B010  # target is typed object here, so a plain attribute store does not type check
+            target,
+            "_hidden_params",
+            {  # mutable-ok: the cost calculator writes optional_params into _hidden_params
+                "additional_headers": {**headers},  # mutable-ok: fresh copy, logging callbacks may mutate it
+                "headers": {**raw_headers},  # mutable-ok: fresh copy, logging callbacks may mutate it
+                **existing,
+            },
+        )
+
     def _handle_logging_completed_response(self):
         """Base implementation - should be overridden by subclasses"""
 

litellm/router.py

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

litellm/router_strategy/complexity_router/README.md

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

litellm/router_strategy/complexity_router/complexity_router.py

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

litellm/router_strategy/complexity_router/config.py

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

litellm/router_utils/fallback_event_handlers.py

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

litellm/router_utils/pre_call_checks/deployment_affinity_check.py

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

litellm/types/integrations/datadog_llm_obs.py

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

litellm/types/llms/gemini_audio_transcription.py

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

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

@@ -4,7 +4,18 @@ from .base import GuardrailConfigModel
 
 
 class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel):
-    pass
+    streaming_end_of_stream_only: bool | None = Field(
+        default=None,
+        description="If False (default when unset), post_call scans the accumulated streamed response every "
+        "streaming_sampling_rate chunks and an in-flight block stops the stream. If True, the guard runs once "
+        "over the assembled response at end of stream, so flagged content may already have reached the client.",
+    )
+    streaming_sampling_rate: int | None = Field(
+        default=None,
+        ge=1,
+        description="When streaming_end_of_stream_only is False, scan the accumulated streamed response every Nth "
+        "chunk. Defaults to 5 when unset.",
+    )
 
 
 class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]):

litellm/types/proxy/public_endpoints/public_endpoints.py

@@ -1,7 +1,7 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import Any, Literal
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class PublicModelHubInfo(BaseModel):
@@ -73,6 +73,44 @@ class SupportedEndpointsResponse(BaseModel):
     endpoints: list[SupportedEndpoint]
 
 
+class AutoRouterPresetTiers(BaseModel):
+    """Exactly the four built-in tiers the dashboard's preset prefill can apply.
+
+    extra="forbid" on purpose: a tier name this dashboard cannot apply would grey out or crash the
+    picker, so such a catalog is rejected wholesale and the bundled one serves instead.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    SIMPLE: Sequence[str]
+    MEDIUM: Sequence[str]
+    COMPLEX: Sequence[str]
+    REASONING: Sequence[str]
+
+
+class AutoRouterPresetConfig(BaseModel):
+    """The complexity_router_config a preset prefills.
+
+    Only tiers is validated, because every dashboard consumer dereferences it; everything else
+    passes through verbatim with unknown fields kept (extra="allow"), so a catalog published after
+    this proxy shipped still serves its new fields intact.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    tiers: AutoRouterPresetTiers
+
+
+class AutoRouterPresetRecord(BaseModel):
+    """One auto-router preset as served to the dashboard's template picker."""
+
+    model_config = ConfigDict(extra="allow")
+
+    label: str
+    description: str
+    complexity_router_config: AutoRouterPresetConfig
+
+
 class ComplexityScorerDefaults(BaseModel):
     """The complexity router's shipped heuristic scorer defaults.
 

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/code_coverage_tests/check_py310_typing_imports.py

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

tests/e2e/ui/helpers/traffic.ts

@@ -21,6 +21,8 @@ interface ChatOptions {
   apiKey?: string;
   /** Sent as `user`, which lands in the spend log's end_user column. */
   endUser?: string;
+  /** Sent as `litellm_trace_id`, which lands in the spend log's session_id column. */
+  traceId?: string;
 }
 
 /** POST /v1/chat/completions and return the completion id (the Logs Request ID). */
@@ -34,6 +36,7 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO
       model: opts.model,
       messages: [{ role: "user", content: opts.prompt }],
       ...(opts.endUser ? { user: opts.endUser } : {}),
+      ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}),
     },
   });
   expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true);

tests/e2e/ui/tests/logs/logsPagination.spec.ts

@@ -0,0 +1,150 @@
+import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
+import { Page } from "../../fixtures/pages";
+import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
+
+/**
+ * Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a
+ * session must never straddle pages, and two callers reusing one session id stay separate rows.
+ * All traffic is generated per run behind a unique key alias or session id, so concurrent specs
+ * cannot decide the outcome.
+ */
+
+const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+
+/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */
+const requestLogsRows = (page: PlaywrightPage): Locator =>
+  page.locator("table").filter({ visible: true }).first().locator("tbody tr");
+
+const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
+
+async function openLogs(page: PlaywrightPage): Promise<void> {
+  await navigateToPage(page, Page.Logs);
+  await dismissFeedbackPopup(page);
+  await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
+}
+
+async function openFilterDrawer(page: PlaywrightPage): Promise<Locator> {
+  await visibleTestId(page, "datatable-filters-trigger").click();
+  const drawer = page.getByRole("dialog", { name: "Filters" });
+  await expect(drawer).toBeVisible({ timeout: 10_000 });
+  return drawer;
+}
+
+async function applyKeyAliasFilter(page: PlaywrightPage, drawer: Locator, alias: string): Promise<void> {
+  await drawer.getByRole("combobox", { name: "Search a key alias" }).click();
+  await page.keyboard.type(alias);
+  await page.getByRole("option", { name: alias, exact: true }).first().click();
+  await drawer.getByRole("button", { name: "Apply Filters" }).click();
+  await expect(drawer).not.toBeVisible({ timeout: 10_000 });
+}
+
+async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
+  await visibleTestId(page, "pagination-page-size").click();
+  await page.getByRole("option", { name: size, exact: true }).click();
+}
+
+test.describe("Logs page session-grouped pagination", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH });
+
+  test("a 25-row page renders exactly 25 session rows and no session straddles pages", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const alias = `e2e-logs-pgn-${suffix}`;
+    const mine = await createVirtualKey(request, { key_alias: alias });
+
+    const soloIds: string[] = [];
+    for (let i = 0; i < 26; i++) {
+      soloIds.push(
+        await sendChatCompletion(request, {
+          model: CHAT_MODEL_A,
+          prompt: `logs-pgn-solo-${i}-${suffix}`,
+          apiKey: mine.key,
+        }),
+      );
+    }
+    const sessionA = `sess-pgn-a-${suffix}`;
+    const sessionB = `sess-pgn-b-${suffix}`;
+    let lastSessionCallId = "";
+    for (let i = 0; i < 7; i++) {
+      lastSessionCallId = await sendChatCompletion(request, {
+        model: CHAT_MODEL_A,
+        prompt: `logs-pgn-a-${i}-${suffix}`,
+        apiKey: mine.key,
+        traceId: sessionA,
+      });
+    }
+    for (let i = 0; i < 3; i++) {
+      lastSessionCallId = await sendChatCompletion(request, {
+        model: CHAT_MODEL_A,
+        prompt: `logs-pgn-b-${i}-${suffix}`,
+        apiKey: mine.key,
+        traceId: sessionB,
+      });
+    }
+    await waitForSpendLog(request, lastSessionCallId);
+    await waitForSpendLog(request, soloIds[soloIds.length - 1]);
+
+    // 36 calls in 28 session groups: 26 solos plus sessions of 7 and 3.
+    await openLogs(page);
+    const drawer = await openFilterDrawer(page);
+    await applyKeyAliasFilter(page, drawer, alias);
+    await setRowsPerPage(page, "25");
+
+    await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 28", { timeout: 30_000 });
+    await expect(requestLogsRows(page)).toHaveCount(25);
+    // The sessions are the newest groups, so their single representative rows sit on page 1.
+    await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(1);
+    await expect(requestLogsRows(page).filter({ hasText: sessionA })).toContainText("7");
+    await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(1);
+
+    await visibleTestId(page, "pagination-next").click();
+
+    await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-28 of 28", { timeout: 30_000 });
+    await expect(requestLogsRows(page)).toHaveCount(3);
+    await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(0);
+    await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(0);
+  });
+
+  test("two keys reusing one session id stay separate rows", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const mine = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-mine-${suffix}` });
+    const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-theirs-${suffix}` });
+    const sharedSession = `sess-pgn-shared-${suffix}`;
+
+    let lastId = "";
+    for (let i = 0; i < 2; i++) {
+      lastId = await sendChatCompletion(request, {
+        model: CHAT_MODEL_A,
+        prompt: `logs-pgn-shared-mine-${i}-${suffix}`,
+        apiKey: mine.key,
+        traceId: sharedSession,
+      });
+    }
+    lastId = await sendChatCompletion(request, {
+      model: CHAT_MODEL_A,
+      prompt: `logs-pgn-shared-theirs-${suffix}`,
+      apiKey: theirs.key,
+      traceId: sharedSession,
+    });
+    await waitForSpendLog(request, lastId);
+
+    await openLogs(page);
+    const drawer = await openFilterDrawer(page);
+    await drawer.getByPlaceholder("Enter session ID…").fill(sharedSession);
+    await drawer.getByRole("button", { name: "Apply Filters" }).click();
+    await expect(drawer).not.toBeVisible({ timeout: 10_000 });
+
+    // One row per caller: reusing a session id must not merge two keys' activity into one row.
+    await expect(requestLogsRows(page).filter({ hasText: sharedSession })).toHaveCount(2, { timeout: 30_000 });
+
+    // And each row carries ITS key's totals: two calls badge the first key's row,
+    // while the other key's single call renders as a plain LLM row.
+    const mineRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: mine.token });
+    const theirsRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: theirs.token });
+    await expect(mineRow).toHaveCount(1);
+    await expect(theirsRow).toHaveCount(1);
+    await expect(mineRow.getByText("2", { exact: true })).toBeVisible();
+    await expect(theirsRow.getByText("LLM", { exact: true })).toBeVisible();
+  });
+});

tests/openai_endpoints_tests/test_e2e_openai_responses_api.py

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

tests/router_unit_tests/test_router_anthropic_messages_fallback.py

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

tests/rust-python-harness/AGENTS.md

@@ -0,0 +1,44 @@
+# Expected Structure
+
+```text
+tests/rust-python-harness/
+├── __main__.py
+│
+├── strategies/
+│   ├── e2e_parity/
+│   │   ├── runner.py
+│   │   ├── sdk/
+│   │   │   ├── ocr/
+│   │   │   ├── messages/
+│   │   │   ├── chat_completions/
+│   │   │   └── responses/
+│   │   └── gateway/
+│   │
+│   ├── trace_parity/
+│   │   ├── runner.py
+│   │   ├── sdk/
+│   │   └── gateway/
+│   │
+│   └── unit_tests/
+│       ├── runner.py
+│       ├── mapping_validator.py
+│       ├── python_runner.py
+│       └── rust_runner.py
+│
+└── shared/
+    ├── parity/
+    ├── tracing/
+    └── reporting/
+```
+
+- Run locally only; no CI integration
+- `__main__.py` selects strategies and combines their reports; each strategy also runs independently
+- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
+- `trace_parity/` compares mapped operations, call counts, and required execution ordering
+- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders
+- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs
+- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts
+- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results
+- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation
+- `shared/` contains reusable parity, tracing, and reporting machinery
+- Keep fixtures with their owning API and existing Python tests in their current locations

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/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py

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

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

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

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

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

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

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

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

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

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

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

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py

@@ -7,6 +7,7 @@ cache's hit/single-flight behavior. Each assertion fails under a real mutation o
 """
 
 import asyncio
+import gc
 import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -359,6 +360,110 @@ async def test_cache_invalidate_only_evicts_the_named_key():
     assert calls == 2
 
 
+@pytest.mark.asyncio
+async def test_cache_invalidate_mid_compute_is_not_overwritten_by_that_compute():
+    """A bearer minted before an invalidation must never be served after it.
+
+    The compute is suspended at the token endpoint when the invalidation lands, so its write is
+    the one that would resurrect the evicted bearer for the rest of its TTL. The caller it was
+    minted for still gets it; the *cache* is what the invalidation is about.
+    """
+    cache = ExchangedTokenCache()
+    mint_started, release_mint = asyncio.Event(), asyncio.Event()
+
+    async def slow_mint():
+        mint_started.set()
+        await release_mint.wait()
+        return _ok_token("bearer-minted-before-invalidation")
+
+    async def re_mint():
+        return _ok_token("bearer-minted-after-invalidation")
+
+    in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
+    await mint_started.wait()
+
+    assert not in_flight.done()
+    cache.invalidate("slot")
+    release_mint.set()
+
+    raced = await in_flight
+    assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation"
+
+    after = await cache.get_or_compute("slot", re_mint, fingerprint="fp")
+    assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation"
+
+
+@pytest.mark.asyncio
+async def test_cache_invalidate_mid_compute_survives_garbage_collection():
+    """The record of an invalidation must outlive a collection cycle taken mid-compute.
+
+    Per-key state is held weakly so idle keys do not accumulate. If the state a compute checks
+    before writing were collectible while that compute is suspended, the check would read as
+    "nothing was invalidated" and the stale write would land; the running compute has to pin it.
+    """
+    cache = ExchangedTokenCache()
+    mint_started, release_mint = asyncio.Event(), asyncio.Event()
+
+    async def slow_mint():
+        mint_started.set()
+        await release_mint.wait()
+        return _ok_token("bearer-minted-before-invalidation")
+
+    async def re_mint():
+        return _ok_token("bearer-minted-after-invalidation")
+
+    in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
+    await mint_started.wait()
+
+    assert not in_flight.done()
+    cache.invalidate("slot")
+    gc.collect()
+    release_mint.set()
+    await in_flight
+
+    after = await cache.get_or_compute("slot", re_mint, fingerprint="fp")
+    assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation"
+
+
+@pytest.mark.asyncio
+async def test_cache_stores_a_compute_that_started_after_the_invalidation():
+    """Only the mint that predates the invalidation loses its write.
+
+    A caller queued behind the single-flight lock computes after the eviction, so its token is
+    fresh and must be cached; otherwise the fix would trade one stale bearer for re-minting on
+    every subsequent resolution.
+    """
+    cache = ExchangedTokenCache()
+    mint_started, release_mint = asyncio.Event(), asyncio.Event()
+
+    async def slow_mint():
+        mint_started.set()
+        await release_mint.wait()
+        return _ok_token("bearer-minted-before-invalidation")
+
+    async def re_mint():
+        return _ok_token("bearer-minted-after-invalidation")
+
+    async def must_not_run():
+        pytest.fail("the mint that followed the invalidation should have been cached")
+
+    in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp"))
+    await mint_started.wait()
+    queued = asyncio.create_task(cache.get_or_compute("slot", re_mint, fingerprint="fp"))
+    await asyncio.sleep(0)
+
+    assert not queued.done()
+    cache.invalidate("slot")
+    release_mint.set()
+
+    raced, fresh = await asyncio.gather(in_flight, queued)
+    assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation"
+    assert isinstance(fresh, Ok) and fresh.ok == "bearer-minted-after-invalidation"
+
+    served = await cache.get_or_compute("slot", must_not_run, fingerprint="fp")
+    assert isinstance(served, Ok) and served.ok == "bearer-minted-after-invalidation"
+
+
 @pytest.mark.asyncio
 async def test_cache_does_not_store_a_failed_compute():
     cache = ExchangedTokenCache()

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"},
@@ -8198,6 +8200,81 @@ class TestPreemptive401ModeAware:
         await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
 
 
+def _make_obo_server(alias: str) -> MCPServer:
+    return MCPServer(
+        server_id=f"id-{alias}",
+        name=alias,
+        alias=alias,
+        server_name=alias,
+        url=f"https://{alias}.test/mcp",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2_token_exchange,
+        token_exchange_endpoint="https://idp.test/token",
+        client_id="cid",
+        client_secret="csecret",
+        mcp_info={"server_name": alias},
+    )
+
+
+class TestOboPreflightScopedToAllowedServers:
+    """The connect-time OBO exchange is an outbound IdP call whose result is cached, so it must
+    only run for a server the caller's key resolves to through the allowed set, not for any
+    server the requested path happens to name."""
+
+    SUBJECT_HEADERS = {"Authorization": "Bearer upstream-subject-token"}
+
+    async def _run(self, requested: MCPServer, allowed: list[MCPServer], user_api_key_auth: UserAPIKeyAuth | None):
+        from litellm.proxy._experimental.mcp_server import server as server_module
+
+        allowed_lookup = AsyncMock(return_value=allowed)
+        preflight = AsyncMock()
+        with (
+            patch.object(  # test-quality-ok: route handler reads the module-level manager, no injection seam
+                server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested
+            ),
+            patch.object(  # test-quality-ok: the exchanger is the observable; a real one would call an IdP
+                server_module.global_mcp_server_manager, "preflight_token_exchange", preflight
+            ),
+            patch.object(  # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer
+                server_module, "_get_allowed_mcp_servers", allowed_lookup
+            ),
+        ):
+            await server_module._raise_preemptive_401_for_unauthenticated_servers(
+                scope={"type": "http", "method": "POST", "path": f"/mcp/{requested.alias}", "headers": []},
+                mcp_servers=[requested.alias],
+                oauth2_headers=self.SUBJECT_HEADERS,
+                mcp_server_auth_headers=None,
+                user_api_key_auth=user_api_key_auth,
+                client_ip="10.0.0.7",
+            )
+        return allowed_lookup, preflight
+
+    @pytest.mark.asyncio
+    async def test_unentitled_key_never_reaches_the_exchanger(self):
+        requested = _make_obo_server("obo_tools")
+        key = UserAPIKeyAuth(api_key="sk-plain-only")
+
+        allowed_lookup, preflight = await self._run(
+            requested, allowed=[_make_obo_server("plain_tools")], user_api_key_auth=key
+        )
+
+        preflight.assert_not_awaited()
+        allowed_lookup.assert_awaited_once_with(
+            user_api_key_auth=key, mcp_servers=[requested.alias], client_ip="10.0.0.7"
+        )
+
+    @pytest.mark.asyncio
+    async def test_entitled_key_still_exchanges_at_connect(self):
+        requested = _make_obo_server("obo_tools")
+        key = UserAPIKeyAuth(api_key="sk-obo")
+
+        _, 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, raw_headers=None
+        )
+
+
 @pytest.mark.asyncio
 async def test_post_mcp_call_guardrails_return_the_rewritten_result():
     """The result a post_mcp_call guardrail rewrote must be what the caller sends back."""

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/client/cli/test_auth_commands.py

@@ -59,6 +59,7 @@ def _mock_cli_sso_start_response(
     login_id: str = "cli-session-uuid-456",
     poll_secret: str = "poll-secret",
     user_code: str = "ABCD-EFGH",
+    **extra_fields: object,
 ) -> Mock:
     mock_response = Mock()
     mock_response.status_code = 200
@@ -66,6 +67,7 @@ def _mock_cli_sso_start_response(
         "login_id": login_id,
         "poll_secret": poll_secret,
         "user_code": user_code,
+        **extra_fields,
     }
     mock_response.raise_for_status = Mock()
     return mock_response
@@ -333,7 +335,9 @@ class TestLoginCommand:
             call_args = mock_browser.call_args[0][0]
             assert "https://test.example.com/sso/key/generate" in call_args
             assert "cli-test-uuid-123" in call_args
+            assert "user_code" not in call_args
             assert "Verification code: ABCD-EFGH" in result.output
+            assert "pre-filled in the browser" not in result.output
             mock_post.assert_called_once()
             mock_get.assert_called()
             assert mock_get.call_args.kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"}
@@ -347,6 +351,72 @@ class TestLoginCommand:
             # Verify commands were shown
             mock_show_commands.assert_called_once()
 
+    def test_login_prefills_the_code_in_the_browser_when_the_proxy_advertises_it(
+        self, isolated_home, secret_vault_factory
+    ) -> None:
+        vault = secret_vault_factory()
+        poll_response = Mock()
+        poll_response.status_code = 200
+        poll_response.json.return_value = {
+            "status": "ready",
+            "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt",
+            "user_id": "test-user-123",
+            "team_id": "team-1",
+            "teams": ["team-1"],
+        }
+        start_response = _mock_cli_sso_start_response(
+            login_id="cli-test-uuid-123",
+            verification_uri_complete=(
+                "https://internal-hostname.example.com/sso/key/generate"
+                "?source=litellm-cli&key=cli-test-uuid-123&user_code=ABCD-EFGH"
+            ),
+        )
+
+        with (
+            patch("webbrowser.open") as mock_browser,
+            patch("requests.post", return_value=start_response),
+            patch("requests.get", return_value=poll_response),
+        ):
+            result = self.runner.invoke(login, obj={"base_url": "https://test.example.com", "secret_vault": vault})
+
+        assert result.exit_code == 0, result.output
+        assert json.loads(vault.blob)["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt"
+        assert json.loads((isolated_home / ".litellm" / "token.json").read_text())["user_id"] == "test-user-123"
+        opened_url = mock_browser.call_args[0][0]
+        assert opened_url.startswith("https://test.example.com/sso/key/generate?")
+        assert "internal-hostname" not in opened_url
+        assert "key=cli-test-uuid-123" in opened_url
+        assert "user_code=ABCD-EFGH" in opened_url
+        assert "Verification code: ABCD-EFGH (pre-filled in the browser, check it matches)" in result.output
+
+    def test_login_keeps_the_code_out_of_the_url_when_the_proxy_sends_a_non_url_verification_uri(
+        self, secret_vault_factory
+    ) -> None:
+        poll_response = Mock()
+        poll_response.status_code = 200
+        poll_response.json.return_value = {
+            "status": "ready",
+            "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt",
+            "user_id": "test-user-123",
+            "team_id": "team-1",
+            "teams": ["team-1"],
+        }
+        for advertised in (None, True):
+            start_response = _mock_cli_sso_start_response(verification_uri_complete=advertised)
+
+            with (
+                patch("webbrowser.open") as mock_browser,
+                patch("requests.post", return_value=start_response),
+                patch("requests.get", return_value=poll_response),
+            ):
+                result = self.runner.invoke(
+                    login, obj={"base_url": "https://test.example.com", "secret_vault": secret_vault_factory()}
+                )
+
+            assert result.exit_code == 0, result.output
+            assert "user_code" not in mock_browser.call_args[0][0]
+            assert "pre-filled in the browser" not in result.output
+
     def test_login_timeout(self):
         """Test login timeout scenario"""
         mock_context = Mock()

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

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

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

@@ -3,7 +3,9 @@ from unittest.mock import patch
 import httpx
 import pytest
 from fastapi import HTTPException
+from pydantic import ValidationError
 
+import litellm
 from litellm.exceptions import Timeout
 from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
 from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail
@@ -12,8 +14,8 @@ from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr
     CrowdStrikeAIDRHandler,
 )
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
-from litellm.types.guardrails import Guardrail, LitellmParams
-from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
+from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
+from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream
 
 
 @pytest.fixture
@@ -1578,3 +1580,142 @@ async def test_unparseable_transformed_response_fails_closed_under_fail_open() -
 
     assert exc_info.value.status_code == 500
     assert "failing closed" in exc_info.value.detail["error"]
+
+
+def _initialize_from_config(**litellm_params_kwargs: object) -> CrowdStrikeAIDRHandler:
+    litellm_params = LitellmParams(
+        guardrail="crowdstrike_aidr",
+        api_key="pts_crowdstrike_tokenid",
+        api_base="https://api.crowdstrike.com/aidr/aiguard",
+        default_on=True,
+        **litellm_params_kwargs,
+    )
+    guardrail = Guardrail(guardrail_name="crowdstrike-aidr-guard", litellm_params=litellm_params)
+    return initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail)
+
+
+@pytest.mark.parametrize(
+    ("mode", "runs_pre_call", "runs_post_call"),
+    [("post_call", False, True), ("pre_call", True, False), (["pre_call", "post_call"], True, True)],
+)
+def test_initialize_guardrail_honors_configured_mode(
+    mode: str | list[str], runs_pre_call: bool, runs_post_call: bool
+) -> None:
+    handler = _initialize_from_config(mode=mode)
+
+    assert handler.should_run_guardrail({}, GuardrailEventHooks.pre_call) is runs_pre_call
+    assert handler.should_run_guardrail({}, GuardrailEventHooks.post_call) is runs_post_call
+
+
+def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_hooks() -> None:
+    with pytest.raises(ValueError, match="during_call is not in the supported event hooks"):
+        _initialize_from_config(mode="during_call")
+
+
+def test_initialize_guardrail_defaults_streaming_params() -> None:
+    handler = _initialize_from_config(mode="post_call")
+
+    assert handler.streaming_end_of_stream_only is False
+    assert handler.streaming_sampling_rate == 5
+
+
+@pytest.mark.parametrize(
+    "configured",
+    [
+        {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50},
+        {"optional_params": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}},
+    ],
+)
+def test_initialize_guardrail_forwards_streaming_params(configured: dict[str, object]) -> None:
+    handler = _initialize_from_config(mode="post_call", **configured)
+
+    assert handler.streaming_end_of_stream_only is True
+    assert handler.streaming_sampling_rate == 50
+
+
+def test_initialize_guardrail_rejects_non_positive_sampling_rate() -> None:
+    with pytest.raises(ValidationError):
+        _initialize_from_config(mode="post_call", streaming_sampling_rate=0)
+
+
+def test_update_in_memory_litellm_params_reapplies_streaming_params() -> None:
+    handler = _initialize_from_config(mode="post_call")
+
+    handler.update_in_memory_litellm_params(
+        LitellmParams(
+            guardrail="crowdstrike_aidr",
+            mode="post_call",
+            streaming_end_of_stream_only=True,
+            streaming_sampling_rate=7,
+        )
+    )
+
+    assert handler.streaming_end_of_stream_only is True
+    assert handler.streaming_sampling_rate == 7
+
+
+def _stream_chunk(content: str, finish_reason: str | None) -> ModelResponseStream:
+    return ModelResponseStream(
+        model="gpt-4",
+        choices=[
+            litellm.StreamingChoices(
+                index=0, delta=Delta(role="assistant", content=content), finish_reason=finish_reason
+            )
+        ],
+    )
+
+
+async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: list[str]) -> int:
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
+
+    async def stream():
+        for i, content in enumerate(chunk_texts):
+            yield _stream_chunk(content, "stop" if i == len(chunk_texts) - 1 else None)
+
+    calls = 0
+
+    def _allow(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(
+            status_code=200, json={"result": {"blocked": False, "transformed": False}}, request=request
+        )
+
+    request_data = {
+        "messages": [{"role": "user", "content": "hi"}],
+        "guardrail_to_apply": handler,
+        "metadata": {"guardrails": ["crowdstrike-aidr-guard"]},
+    }
+    async with httpx.AsyncClient(transport=httpx.MockTransport(_allow)) as client:
+        await handler.async_handler.close()
+        handler.async_handler.client = client
+        async for _ in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"),
+            response=stream(),
+            request_data=request_data,
+        ):
+            pass
+    return calls
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("configured", "expected_calls"),
+    [
+        ({}, 2),
+        ({"streaming_sampling_rate": 2}, 5),
+        ({"streaming_end_of_stream_only": True}, 1),
+        ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1),
+    ],
+)
+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, rate 2 samples 5 times, end-of-stream scans once.
+
+    The final pass is skipped because chunk 10 already scanned the complete output.
+    """
+    handler = _initialize_from_config(mode="post_call", **configured)
+
+    assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls

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

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

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

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

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker):
     mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
     mock_handler.get_source.return_value = "config"
     mock_handler.initialize_guardrail = mocker.Mock()
-    mock_handler.update_in_memory_guardrail = mocker.Mock()
+    mock_handler.sync_guardrail_from_db = mocker.Mock()
     mock_handler.delete_in_memory_guardrail = mocker.Mock()
     mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[])
     return mock_handler
@@ -1047,13 +1047,15 @@ async def test_create_guardrail_endpoint(
     "scenario,expected_result,expected_exception",
     [
         ("success_with_sync", "test-db-guardrail", None),
-        ("success_sync_fails", "test-db-guardrail", None),
+        ("success_sync_fails_unexpected_error", "test-db-guardrail", None),
+        ("sync_fails_invalid_config", None, HTTPException),
         ("database_failure", None, HTTPException),
         ("no_prisma_client", None, HTTPException),
     ],
     ids=[
         "success_with_immediate_sync",
-        "success_but_sync_fails",
+        "success_but_sync_fails_with_unexpected_error",
+        "sync_rejects_invalid_config",
         "database_error",
         "missing_prisma_client",
     ],
@@ -1073,6 +1075,7 @@ async def test_update_guardrail_endpoint(
     mock_logger = None
     if scenario == "success_with_sync":
         mock_prisma_client = mocker.Mock()
+        mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock()
         mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
         mocker.patch(
             "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
@@ -1083,10 +1086,13 @@ async def test_update_guardrail_endpoint(
             mock_in_memory_handler,
         )
 
-    elif scenario == "success_sync_fails":
+    elif scenario == "success_sync_fails_unexpected_error":
+        # A non-ValueError/TypeError failure is not a config-rejection signal,
+        # so it keeps the pre-existing swallow-and-warn behavior rather than
+        # rolling back the DB write.
         mock_prisma_client = mocker.Mock()
-        mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception(
-            "Sync failed"
+        mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
+            side_effect=Exception("Sync failed")
         )
         mock_logger = mocker.patch(
             "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
@@ -1102,6 +1108,25 @@ async def test_update_guardrail_endpoint(
             mock_in_memory_handler,
         )
 
+    elif scenario == "sync_fails_invalid_config":
+        # Regression for the PUT half of the fix: a TypeError from the sync (the
+        # in-place update_in_memory_guardrail raised exactly this on every PUT)
+        # must roll back the DB write and surface a 422, not persist the
+        # rejected config with a 200.
+        mock_prisma_client = mocker.Mock()
+        mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
+            side_effect=TypeError("vars() argument must have __dict__ attribute")
+        )
+        mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)  # test-quality-ok: reused pattern
+        mocker.patch(  # test-quality-ok: reused pattern
+            "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
+            mock_guardrail_registry,
+        )
+        mocker.patch(  # test-quality-ok: reused pattern
+            "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
+            mock_in_memory_handler,
+        )
+
     elif scenario == "database_failure":
         mock_prisma_client = mocker.Mock()
         mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
@@ -1130,6 +1155,16 @@ async def test_update_guardrail_endpoint(
             assert "Database error" in str(exc_info.value.detail)
         elif scenario == "no_prisma_client":
             assert "Prisma client not initialized" in str(exc_info.value.detail)
+        elif scenario == "sync_fails_invalid_config":
+            assert exc_info.value.status_code == 422
+            assert "update rejected" in str(exc_info.value.detail)
+            # Rolled back: update_guardrail_in_db is called once for the
+            # rejected write and once more to restore the previous config.
+            assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
+            assert (
+                mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
+                == MOCK_DB_GUARDRAIL
+            )
 
     else:
         result = await update_guardrail(
@@ -1145,11 +1180,11 @@ async def test_update_guardrail_endpoint(
             prisma_client=mocker.ANY,
         )
 
-        mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with(
-            guardrail_id="test-guardrail-id", guardrail=mocker.ANY
+        mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
+            guardrail=mocker.ANY
         )
 
-        if scenario == "success_sync_fails":
+        if scenario == "success_sync_fails_unexpected_error":
             assert mock_logger is not None
             mock_logger.warning.assert_called_once()
             assert "Failed to update" in str(mock_logger.warning.call_args)

tests/test_litellm/proxy/guardrails/test_guardrail_registry.py

@@ -913,3 +913,96 @@ def test_reinitialize_guardrail_restores_previous_on_failure():
         assert restored.guardrail_name == "restore-me"
     finally:
         registry_module.guardrail_initializer_registry.pop("restore_test", None)
+
+
+def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures():
+    """Regression for the LIT-6479 fix's 422 path: a constructor failure that is not
+    already a ValueError/TypeError (re.error from an invalid regex has neither in its
+    MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422
+    catch is exhaustive instead of warn-and-200 persisting a broken config."""
+    import re
+
+    from litellm.proxy.guardrails import guardrail_registry as registry_module
+
+    def _initializer(litellm_params, guardrail):
+        if litellm_params.api_key == "bad-regex":
+            re.compile("([")
+        return CustomGuardrail(
+            guardrail_name=guardrail["guardrail_name"],
+            event_hook=GuardrailEventHooks.pre_call,
+            default_on=True,
+        )
+
+    registry_module.guardrail_initializer_registry["regex_test"] = _initializer
+    try:
+        handler = InMemoryGuardrailHandler()
+        created = handler.initialize_guardrail(
+            guardrail={
+                "guardrail_name": "regex-me",
+                "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"},
+            },
+        )
+        guardrail_id = created["guardrail_id"]
+
+        with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo:
+            handler.reinitialize_guardrail(
+                guardrail={
+                    "guardrail_id": guardrail_id,
+                    "guardrail_name": "regex-me",
+                    "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"},
+                },
+            )
+
+        assert isinstance(excinfo.value.__cause__, re.error)
+        assert guardrail_id in handler.IN_MEMORY_GUARDRAILS
+        restored = handler.guardrail_id_to_custom_guardrail[guardrail_id]
+        assert restored is not None and restored.guardrail_name == "regex-me"
+    finally:
+        registry_module.guardrail_initializer_registry.pop("regex_test", None)
+
+
+def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance():
+    """
+    Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as
+    a plain jsonb dict, and the in-place update_in_memory_guardrail cast it to
+    LitellmParams without constructing one, so vars() raised and the running proxy
+    kept enforcing the stale config forever. The PUT endpoint now routes through
+    sync_guardrail_from_db, which must rebuild the live instance from the dict:
+    new blocked words compiled in, old ones gone, and the event hook re-derived
+    from mode (the base-class setattr path wrote self.mode while dispatch reads
+    self.event_hook, so only a full re-init applies a mode change).
+    """
+    from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
+        ContentFilterGuardrail,
+    )
+
+    handler = InMemoryGuardrailHandler()
+    gid = "66666666-6666-6666-6666-666666666666"
+
+    def db_guardrail(word: str, mode: str) -> Guardrail:
+        return Guardrail(
+            guardrail_id=gid,
+            guardrail_name="cf-put-sync",
+            litellm_params={
+                "guardrail": "litellm_content_filter",
+                "mode": mode,
+                "default_on": True,
+                "blocked_words": [{"keyword": word, "action": "BLOCK"}],
+            },
+        )
+
+    lists = _all_callback_lists()
+    snapshots = [list(cb_list) for cb_list in lists]
+    try:
+        handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call"))
+        handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call"))
+
+        instance = handler.guardrail_id_to_custom_guardrail[gid]
+        assert isinstance(instance, ContentFilterGuardrail)
+        assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None
+        assert instance._check_blocked_words("hello FOOBARBLOCK") is None
+        assert instance.event_hook == GuardrailEventHooks.during_call
+        assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True
+    finally:
+        for cb_list, snapshot in zip(lists, snapshots):
+            cb_list[:] = snapshot

tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py

@@ -1077,3 +1077,243 @@ def test_public_mcp_hub_does_not_expose_upstream_url():
     assert all("url" not in item for item in data)
     assert secret_url not in response.text
     app.dependency_overrides.clear()
+
+
+
+@pytest.fixture
+def reset_autorouter_presets_cache():
+    from litellm.proxy.public_endpoints.public_endpoints import _AutoRouterPresetsCache
+
+    _AutoRouterPresetsCache.presets = None
+    _AutoRouterPresetsCache.lock = None
+    yield
+    _AutoRouterPresetsCache.presets = None
+    _AutoRouterPresetsCache.lock = None
+
+
+def test_get_autorouter_presets_local_mode_serves_bundled_catalog(
+    monkeypatch, reset_autorouter_presets_cache
+):
+    monkeypatch.setenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "True")
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    response = client.get("/public/autorouter_presets")
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert "anthropic_family" in payload
+    for preset in payload.values():
+        assert isinstance(preset["label"], str)
+        assert isinstance(preset["description"], str)
+        assert "tiers" in preset["complexity_router_config"]
+
+
+@pytest.mark.asyncio
+async def test_get_autorouter_presets_fetches_once_per_process(
+    monkeypatch, reset_autorouter_presets_cache
+):
+    from litellm.proxy.public_endpoints.public_endpoints import (
+        _AUTOROUTER_PRESETS_ADAPTER,
+        get_autorouter_presets,
+    )
+
+    monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False)
+    remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+        {
+            "remote_only": {
+                "label": "Remote Only",
+                "description": "from the remote catalog",
+                "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}},
+            }
+        }
+    )
+    calls = []
+
+    async def fake_fetch(url):
+        calls.append(url)
+        return remote
+
+    first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch)
+    second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch)
+
+    assert first == remote
+    assert second == remote
+    assert calls == ["https://example.test/presets.json"]
+
+
+@pytest.mark.asyncio
+async def test_get_autorouter_presets_single_flight_on_concurrent_cold_start(
+    monkeypatch, reset_autorouter_presets_cache
+):
+    import asyncio
+
+    from litellm.proxy.public_endpoints.public_endpoints import (
+        _AUTOROUTER_PRESETS_ADAPTER,
+        get_autorouter_presets,
+    )
+
+    monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False)
+    remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+        {
+            "remote_only": {
+                "label": "Remote Only",
+                "description": "from the remote catalog",
+                "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}},
+            }
+        }
+    )
+    calls = []
+
+    async def slow_fetch(url):
+        calls.append(url)
+        await asyncio.sleep(0.05)
+        return remote
+
+    results = await asyncio.gather(
+        get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch),
+        get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch),
+        get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch),
+    )
+
+    assert all(result == remote for result in results)
+    assert len(calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_get_autorouter_presets_caches_bundled_fallback_on_remote_failure(
+    monkeypatch, reset_autorouter_presets_cache
+):
+    from litellm.proxy.public_endpoints.public_endpoints import get_autorouter_presets
+
+    monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False)
+    calls = []
+
+    async def broken_fetch(url):
+        calls.append(url)
+        raise ValueError("remote catalog unavailable")
+
+    first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch)
+    second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch)
+
+    assert "anthropic_family" in first
+    assert second == first
+    assert len(calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_autorouter_presets_adapter_rejects_wrong_shapes():
+    from pydantic import ValidationError
+
+    from litellm.proxy.public_endpoints.public_endpoints import _AUTOROUTER_PRESETS_ADAPTER
+
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python({"bad": {"label": "no description or config"}})
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python(["not", "a", "mapping"])
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+            {"no_tiers": {"label": "L", "description": "D", "complexity_router_config": {}}}
+        )
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+            {
+                "missing_builtin_tier": {
+                    "label": "L",
+                    "description": "D",
+                    "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"]}},
+                }
+            }
+        )
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+            {
+                "unknown_tier_name": {
+                    "label": "L",
+                    "description": "D",
+                    "complexity_router_config": {
+                        "tiers": {
+                            "SIMPLE": ["m1"],
+                            "MEDIUM": ["m2"],
+                            "COMPLEX": ["m3"],
+                            "REASONING": ["m4"],
+                            "ULTRA": ["m5"],
+                        }
+                    },
+                }
+            }
+        )
+    with pytest.raises(ValidationError):
+        _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+            {
+                "bad_tiers": {
+                    "label": "L",
+                    "description": "D",
+                    "complexity_router_config": {"tiers": "not-a-mapping"},
+                }
+            }
+        )
+
+
+def test_get_autorouter_presets_passes_unknown_catalog_fields_through(
+    monkeypatch, reset_autorouter_presets_cache
+):
+    from litellm.proxy.public_endpoints.public_endpoints import (
+        _AUTOROUTER_PRESETS_ADAPTER,
+        _AutoRouterPresetsCache,
+    )
+
+    monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False)
+    _AutoRouterPresetsCache.presets = _AUTOROUTER_PRESETS_ADAPTER.validate_python(
+        {
+            "future_preset": {
+                "label": "Future",
+                "description": "carries fields this proxy version does not know",
+                "complexity_router_config": {
+                    "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]},
+                    "future_config_knob": 3,
+                },
+                "icon": "sparkles",
+            }
+        }
+    )
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    response = client.get("/public/autorouter_presets")
+
+    assert response.status_code == 200
+    served = response.json()["future_preset"]
+    assert served["icon"] == "sparkles"
+    assert served["complexity_router_config"]["future_config_knob"] == 3
+    assert served["complexity_router_config"]["tiers"]["SIMPLE"] == ["m1"]
+
+
+@pytest.mark.asyncio
+async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypatch):
+    import litellm.llms.custom_httpx.http_handler as http_handler_module
+    from litellm.proxy.public_endpoints.public_endpoints import _fetch_remote_autorouter_presets
+
+    catalog = {
+        "remote_only": {
+            "label": "Remote Only",
+            "description": "from the remote catalog",
+            "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}},
+        }
+    }
+    response = MagicMock()
+    response.raise_for_status = MagicMock()
+    response.json = MagicMock(return_value=catalog)
+    client = MagicMock()
+    client.get = AsyncMock(return_value=response)
+    monkeypatch.setattr(http_handler_module, "get_async_httpx_client", lambda llm_provider: client)
+
+    presets = await _fetch_remote_autorouter_presets("https://example.test/presets.json")
+    assert presets["remote_only"].label == "Remote Only"
+    response.raise_for_status.assert_called_once()
+
+    response.json = MagicMock(return_value={})
+    with pytest.raises(ValueError, match="empty"):
+        await _fetch_remote_autorouter_presets("https://example.test/presets.json")

tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -3959,18 +3959,18 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
     ]
 
     mock_prisma = MagicMock()
-    mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
-        return_value=[
-            {"session_id": session_id, "_count": {"session_id": 2}},
-        ]
-    )
+    mock_prisma.db.litellm_spendlogs.group_by = AsyncMock()
     mock_prisma.db.query_raw = AsyncMock(
         return_value=[
             {
                 "session_id": session_id,
+                "api_key": api_key,
+                "session_total_count": 2,
                 "session_total_spend": 15.0,
                 "mcp_tool_call_count": 1,
                 "mcp_tool_call_spend": 10.0,
+                "session_llm_count": 1,
+                "session_agent_count": 0,
             }
         ]
     )
@@ -3995,6 +3995,8 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
     assert rows[0]["mcp_tool_call_spend"] == 10.0
     assert rows[1]["mcp_tool_call_count"] == 1
     assert rows[1]["mcp_tool_call_spend"] == 10.0
+    assert rows[0]["session_llm_count"] == 1
+    assert rows[0]["session_agent_count"] == 0
 
     # Every row in the session carries the full session spend, not just its own
     assert rows[0]["session_total_spend"] == 15.0
@@ -4003,13 +4005,126 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts():
     # Row without a session_id defaults to 1
     assert rows[2]["session_total_count"] == 1
 
-    # group_by should have been called with the session_id
-    mock_prisma.db.litellm_spendlogs.group_by.assert_called_once_with(
-        by=["session_id"],
-        where={"session_id": {"in": [session_id]}},
-        count={"session_id": True},
+    # The count is folded into the single aggregate query; no separate group_by call.
+    mock_prisma.db.litellm_spendlogs.group_by.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggregates():
+    """
+    Two keys reusing one session id are separate rows under grouped pagination,
+    and each row must carry ITS key's totals, never the combined session's:
+    the aggregate query and its lookup are keyed by (session_id, api_key).
+    """
+    from litellm.proxy.spend_tracking.spend_management_endpoints import (
+        _build_ui_spend_logs_response,
     )
 
+    session_id = "sess-shared"
+    dict_rows = [
+        {"request_id": "req-a", "session_id": session_id, "call_type": "completion", "api_key": "key-a"},
+        {"request_id": "req-b", "session_id": session_id, "call_type": "completion", "api_key": "key-b"},
+    ]
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "session_id": session_id,
+                "api_key": "key-a",
+                "session_total_count": 2,
+                "session_total_spend": 0.2,
+                "mcp_tool_call_count": 0,
+                "mcp_tool_call_spend": 0.0,
+                "session_cache_hit_count": 1,
+                "session_llm_count": 2,
+                "session_agent_count": 0,
+            },
+            {
+                "session_id": session_id,
+                "api_key": "key-b",
+                "session_total_count": 1,
+                "session_total_spend": 0.7,
+                "mcp_tool_call_count": 0,
+                "mcp_tool_call_spend": 0.0,
+                "session_cache_hit_count": 0,
+                "session_llm_count": 1,
+                "session_agent_count": 0,
+            },
+        ]
+    )
+
+    result = await _build_ui_spend_logs_response(
+        prisma_client=mock_prisma,
+        data=dict_rows,
+        total_records=2,
+        page=1,
+        page_size=50,
+        total_pages=1,
+        enrich_session_counts=True,
+    )
+
+    rows = result["data"]
+    assert [(r["session_total_count"], r["session_total_spend"]) for r in rows] == [(2, 0.2), (1, 0.7)]
+    assert [r["session_cache_hit_count"] for r in rows] == [1, 0]
+    assert [r["session_llm_count"] for r in rows] == [2, 1]
+
+    aggregate_sql = mock_prisma.db.query_raw.mock_calls[0][1][0]
+    assert "GROUP BY session_id, api_key" in aggregate_sql
+
+
+@pytest.mark.asyncio
+async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregates():
+    """
+    The spend-log schema defaults api_key to an empty string, which is a real
+    group value and not a missing one: a multi-call session logged under an
+    empty key must keep its count and spend instead of degrading to a plain
+    single-call row.
+    """
+    from litellm.proxy.spend_tracking.spend_management_endpoints import (
+        _build_ui_spend_logs_response,
+    )
+
+    session_id = "sess-keyless"
+    dict_rows = [
+        {"request_id": "req-1", "session_id": session_id, "call_type": "completion", "api_key": ""},
+    ]
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "session_id": session_id,
+                "api_key": "",
+                "session_total_count": 3,
+                "session_total_spend": 0.09,
+                "mcp_tool_call_count": 0,
+                "mcp_tool_call_spend": 0.0,
+                "session_cache_hit_count": 0,
+                "session_llm_count": 3,
+                "session_agent_count": 0,
+            }
+        ]
+    )
+
+    result = await _build_ui_spend_logs_response(
+        prisma_client=mock_prisma,
+        data=dict_rows,
+        total_records=1,
+        page=1,
+        page_size=50,
+        total_pages=1,
+        enrich_session_counts=True,
+    )
+
+    row = result["data"][0]
+    assert row["session_total_count"] == 3
+    assert row["session_total_spend"] == 0.09
+
+    # The empty key must reach the aggregate's authorized-keys filter too.
+    _, call_args, _ = mock_prisma.db.query_raw.mock_calls[0]
+    assert call_args[2] == [""]
+
 
 @pytest.mark.asyncio
 async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
@@ -4033,14 +4148,13 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
     ]
 
     mock_prisma = MagicMock()
-    mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
-        return_value=[{"session_id": session_id, "_count": {"session_id": 3}}]
-    )
     # The raw aggregate query returns the full session spend (0.01 + 0.02 + 0.03).
     mock_prisma.db.query_raw = AsyncMock(
         return_value=[
             {
                 "session_id": session_id,
+                "api_key": api_key,
+                "session_total_count": 3,
                 "session_total_spend": 0.06,
                 "mcp_tool_call_count": 0,
                 "mcp_tool_call_spend": 0.0,
@@ -4089,13 +4203,12 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count():
     ]
 
     mock_prisma = MagicMock()
-    mock_prisma.db.litellm_spendlogs.group_by = AsyncMock(
-        return_value=[{"session_id": session_id, "_count": {"session_id": 2}}]
-    )
     mock_prisma.db.query_raw = AsyncMock(
         return_value=[
             {
                 "session_id": session_id,
+                "api_key": api_key,
+                "session_total_count": 2,
                 "session_total_spend": 0.05,
                 "mcp_tool_call_count": 0,
                 "mcp_tool_call_spend": 0.0,

tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py

@@ -274,6 +274,9 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch):
         "the page query must not carry a window count that forces a full-window "
         f"scan. SQL was:\n{page_sql}"
     )
+    assert "GROUP BY" not in count_sql and "DISTINCT ON" not in page_sql, (
+        "without group_by_session the endpoint must keep raw per-call pagination"
+    )
 
     assert response["total"] == 137
     assert response["total_is_capped"] is False
@@ -499,3 +502,106 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch):
     params = mock_prisma.db.query_raw.call_args[0][1:]
     assert "team_x" in params, "team_id must be forwarded into the DB query params"
     assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}"
+
+
+@pytest.mark.asyncio
+async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
+    """
+    With group_by_session=true, /spend/logs/ui must page and count SESSIONS,
+    not raw calls: the page query returns one representative row per session
+    (DISTINCT ON the session group key, preferring non-MCP calls, newest
+    first) and the bounded count counts groups. Otherwise the UI collapses a
+    server page of N calls into fewer visible rows while the footer still
+    claims N (issue #38060).
+    """
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.proxy.spend_tracking.spend_management_endpoints import (
+        SPEND_LOGS_PAGINATION_COUNT_CAP,
+        ui_view_spend_logs,
+    )
+
+    page_rows = [
+        {"request_id": "req-1", "metadata": "{}", "session_id": None},
+        {"request_id": "req-2", "metadata": "{}", "session_id": None},
+    ]
+    mock_prisma = _make_ui_spend_logs_mock(count_total=12, page_rows=page_rows)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+
+    auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
+    mock_request = MagicMock()
+    mock_request.url.path = "/spend/logs/ui"
+
+    response = await ui_view_spend_logs(
+        request=mock_request,
+        api_key=None,
+        user_id=None,
+        request_id=None,
+        start_date="2026-02-16 00:00:00",
+        end_date="2026-02-16 23:59:59",
+        page=1,
+        page_size=50,
+        sort_by="startTime",
+        sort_order="desc",
+        user_api_key_dict=auth,
+        group_by_session=True,
+    )
+
+    group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
+
+    count_call = mock_prisma.db.query_raw.call_args_list[0]
+    count_sql = count_call[0][0]
+    assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}"
+    assert "COUNT(*) OVER ()" not in count_sql
+    assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
+    assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1
+
+    page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0]
+    assert f"DISTINCT ON ({group_key})" in page_sql, f"page must return one row per session. SQL was:\n{page_sql}"
+    assert f"ORDER BY {group_key}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" in page_sql, (
+        "the session representative must prefer the newest non-MCP call"
+    )
+    assert "COUNT(*) OVER ()" not in page_sql
+
+    assert response["total"] == 12
+    assert response["total_is_capped"] is False
+    assert response["total_pages"] == 1
+
+
+@pytest.mark.asyncio
+async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(monkeypatch):
+    """
+    A request_id lookup with group_by_session=true must still resolve the
+    exact requested row: the filter runs before grouping, so the row is its
+    own group's representative and deep links keep working.
+    """
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs
+
+    target_row = {"request_id": "req-deep-link", "metadata": "{}", "session_id": None}
+    mock_prisma = _make_ui_spend_logs_mock(count_total=1, page_rows=[target_row])
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+
+    auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
+    mock_request = MagicMock()
+    mock_request.url.path = "/spend/logs/ui"
+
+    response = await ui_view_spend_logs(
+        request=mock_request,
+        api_key=None,
+        user_id=None,
+        request_id="req-deep-link",
+        start_date=None,
+        end_date=None,
+        page=1,
+        page_size=1,
+        sort_by="startTime",
+        sort_order="desc",
+        user_api_key_dict=auth,
+        group_by_session=True,
+    )
+
+    page_call = mock_prisma.db.query_raw.call_args_list[1]
+    assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping"
+    assert "req-deep-link" in page_call[0]
+    assert [row["request_id"] for row in response["data"]] == ["req-deep-link"]
+    assert response["total"] == 1

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

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

tests/test_litellm/responses/test_streaming_iterator.py

@@ -6,7 +6,7 @@ completion_start_time = end_time."""
 import json
 from datetime import datetime
 from typing import Optional
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import httpx
 import pytest
@@ -378,3 +378,162 @@ def test_stamp_responses_usage_cost_survives_calculator_failure():
     _stamp_responses_usage_cost(response, logging_obj)
 
     assert getattr(response.usage, "cost", None) is None
+
+
+def _capture_dispatch(logged: list):
+    """Record the object handed to the success handlers.
+
+    ``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect
+    only runs when the coroutine is awaited, so capture with a plain function instead.
+    """
+
+    async def _noop() -> None:
+        return None
+
+    def _dispatch(result, **kwargs):
+        logged.append(result)
+        return _noop()
+
+    return _dispatch
+
+
+def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock:
+    """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy
+    performs a genuine model_dump/model_validate round trip."""
+    mock_config = Mock(spec=BaseResponsesAPIConfig)
+
+    def _transform(model, parsed_chunk, logging_obj):
+        evt_type = parsed_chunk.get("type")
+        if evt_type != "response.completed":
+            stub = Mock()
+            stub.type = evt_type
+            return stub
+        response = ResponsesAPIResponse(
+            id="resp_headers",
+            created_at=1,
+            output=[],
+            parallel_tool_calls=False,
+            tool_choice="auto",
+            tools=[],
+        )
+        if transform_hidden_params is not None:
+            response._hidden_params.update(transform_hidden_params)
+        return ResponseCompletedEvent(
+            type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
+            response=response,
+        )
+
+    mock_config.transform_streaming_response.side_effect = _transform
+    return mock_config
+
+
+def _make_header_iterator(
+    *,
+    headers: dict,
+    config: Mock,
+    logging_obj: LiteLLMLoggingObj,
+) -> ResponsesAPIStreamingIterator:
+    async def aiter_bytes():
+        yield _sse_event({"type": "response.completed"})
+
+    mock_response = Mock()
+    mock_response.headers = headers
+    mock_response.aiter_bytes = aiter_bytes
+
+    return ResponsesAPIStreamingIterator(
+        response=mock_response,
+        model="gpt-4o-mini",
+        responses_api_provider_config=config,
+        logging_obj=logging_obj,
+        litellm_metadata={},
+        custom_llm_provider="azure",
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_logging_response_carries_provider_response_headers():
+    """LIT-6055: the provider headers the iterator captured must reach the logged response, so
+    custom loggers can read Azure's apim-request-id from the callback payload."""
+    logging_obj = _logging_obj_stub()
+    logged: list[object] = []
+    logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
+
+    logging_obj._on_deferred_stream_complete = None
+
+    iterator = _make_header_iterator(
+        headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"},
+        config=_headers_config(),
+        logging_obj=logging_obj,
+    )
+    async for _ in iterator:
+        pass
+
+    assert len(logged) == 1
+    hidden_params = logged[0].response._hidden_params
+    assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1"
+    assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2"
+    assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1"
+    # the proxy builds the client's response headers from the iterator's own dict, so the logged
+    # response must hold copies rather than alias it
+    assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"]
+    assert hidden_params["headers"] is not iterator._raw_response_headers
+
+
+@pytest.mark.asyncio
+async def test_streaming_logging_copy_preserves_transform_hidden_params():
+    """LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a
+    provider transform already set on the response (fake_stream) must be re-applied."""
+    logging_obj = _logging_obj_stub()
+    logged: list[object] = []
+    logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
+
+    logging_obj._on_deferred_stream_complete = None
+
+    iterator = _make_header_iterator(
+        headers={},
+        config=_headers_config(
+            transform_hidden_params={
+                "additional_headers": {"llm_provider-apim-request-id": "from-transform"},
+                "headers": {"apim-request-id": "from-transform"},
+                "response_cost": 0.5,
+            }
+        ),
+        logging_obj=logging_obj,
+    )
+    async for _ in iterator:
+        pass
+
+    assert len(logged) == 1
+    hidden_params = logged[0].response._hidden_params
+    assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform"
+    assert hidden_params["headers"]["apim-request-id"] == "from-transform"
+    assert iterator.completed_response is not logged[0]
+    # only the header keys travel: response_cost would short-circuit the cost calculator
+    assert "response_cost" not in hidden_params
+
+
+@pytest.mark.asyncio
+async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched():
+    """LIT-6055: when the logging copy falls back to the original event, the header restore must
+    not stamp logging-only state onto the object the caller is iterating."""
+    logging_obj = _logging_obj_stub()
+    logged: list[object] = []
+    logging_obj.dispatch_success_handlers = _capture_dispatch(logged)
+    logging_obj._on_deferred_stream_complete = None
+
+    iterator = _make_header_iterator(
+        headers={"apim-request-id": "azure-correlation-1"},
+        config=_headers_config(),
+        logging_obj=logging_obj,
+    )
+    async for _ in iterator:
+        pass
+
+    assert len(logged) == 1
+    iterator._completed_response_logged = False
+    logged.clear()
+    with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")):
+        iterator._log_completed_response(is_async=True)
+
+    assert logged == [iterator.completed_response]
+    assert iterator.completed_response.response._hidden_params == {}

tests/test_litellm/router_strategy/test_complexity_router.py

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

tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py

@@ -7,7 +7,7 @@ import json
 
 import litellm
 from litellm.caching.dual_cache import DualCache
-from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
+from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
     DeploymentAffinityCheck,
 )
@@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key():
     assert filtered[0]["model_info"]["id"] == "deployment-2"
 
 
+@pytest.mark.asyncio
+async def test_proxy_generated_session_id_does_not_pin_a_deployment():
+    """A session id the proxy generated for a request that had none is per request, so a
+    pin stored under it must be ignored and none must be written."""
+    cache = DualCache()
+    callback = DeploymentAffinityCheck(
+        cache=cache,
+        ttl_seconds=123,
+        enable_user_key_affinity=False,
+        enable_responses_api_affinity=False,
+        enable_session_id_affinity=True,
+    )
+    healthy_deployments = [
+        {"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}},
+        {"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}},
+    ]
+    await cache.async_set_cache(
+        DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"),
+        {"model_id": "deployment-2"},
+    )
+    request_kwargs = {
+        "metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}
+    }
+
+    filtered = await callback.async_filter_deployments(
+        model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs
+    )
+    await callback.async_pre_call_deployment_hook(
+        kwargs={
+            "metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"},
+            "model_info": {"id": "deployment-1"},
+        },
+        call_type=None,
+    )
+
+    assert len(filtered) == 2
+    assert await cache.async_get_cache(
+        DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1")
+    ) == {"model_id": "deployment-2"}
+
+
 MOCK_RESPONSES_API_RESPONSE = {
     "id": "resp_mock-resp-456",
     "object": "response",

tests/test_litellm/test_check_py310_typing_imports.py

@@ -0,0 +1,86 @@
+import sys
+from pathlib import Path
+from typing import Final
+
+_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests"
+sys.path.insert(0, str(_CODE_COVERAGE_DIR))  # test-quality-ok: required to import checker from its source directory
+import check_py310_typing_imports as checker  # noqa: E402  # load checker from its source directory
+
+
+def _scan(tmp_path: Path, source: str) -> tuple[object, ...]:
+    file_path = tmp_path / "fixture.py"
+    file_path.write_text(source, encoding="utf-8")
+    return checker.scan_file(file_path)
+
+
+def test_typing_import_flags_python_311_name(tmp_path: Path) -> None:
+    violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n")
+    assert tuple(violation.name for violation in violations) == ("NotRequired",)
+
+
+def test_typing_extensions_import_passes(tmp_path: Path) -> None:
+    assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == ()
+
+
+def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None:
+    violations = _scan(tmp_path, "import typing\nx: typing.Self\n")
+    assert tuple(violation.name for violation in violations) == ("Self",)
+
+
+def test_version_guarded_typing_import_passes(tmp_path: Path) -> None:
+    source = (
+        "import sys\n"
+        "if sys.version_info >= (3, 11):\n"
+        "    from typing import NotRequired\n"
+        "else:\n"
+        "    from typing_extensions import NotRequired\n"
+    )
+    assert _scan(tmp_path, source) == ()
+
+
+def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None:
+    source = (
+        "import sys\n"
+        "if sys.version_info >= (3, 11):\n"
+        "    from typing_extensions import NotRequired\n"
+        "else:\n"
+        "    from typing import NotRequired\n"
+    )
+    violations = _scan(tmp_path, source)
+    assert tuple(violation.name for violation in violations) == ("NotRequired",)
+
+
+def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None:
+    source = (
+        "import sys\n"
+        "if sys.version_info < (3, 11):\n"
+        "    from typing_extensions import NotRequired\n"
+        "else:\n"
+        "    from typing import NotRequired\n"
+    )
+    assert _scan(tmp_path, source) == ()
+
+
+def test_nearest_if_controls_version_guard(tmp_path: Path) -> None:
+    source = (
+        "if sys.version_info >= (3, 11):\n"
+        "    from typing import Self\n"
+        "    x = 1\n"
+        "if True:\n"
+        "    from typing import Self\n"
+    )
+    violations = _scan(tmp_path, source)
+    assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),)
+
+
+def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None:
+    file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py"
+    file_path.parent.mkdir(parents=True)
+    file_path.write_text("from typing import NotRequired\n", encoding="utf-8")
+
+    violations = checker.scan_directory(tmp_path)
+    assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),)
+
+
+def test_python_310_typing_name_passes(tmp_path: Path) -> None:
+    assert _scan(tmp_path, "from typing import Optional\n") == ()