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

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/proxy/guardrails/anthropic_sse.py

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

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

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

litellm/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_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,

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/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/llms/base_llm/files/test_azure_blob_storage_backend.py

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

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

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