fix(azure): count only the failed attempt's Entra rejection and ignore a non-callable token provider

litellm/llms/azure/common_utils.py

@@ -1,5 +1,6 @@
 import asyncio
 import hashlib
+import itertools
 import json
 import os
 import re
@@ -9,7 +10,7 @@ from types import MappingProxyType
 from typing import Final, Literal, NamedTuple, NoReturn, cast
 
 import httpx
-from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
+from openai import APIError, AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from pydantic import ConfigDict, TypeAdapter, with_config
 from typing_extensions import ReadOnly, TypedDict
 
@@ -57,6 +58,10 @@ _AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitab
 )
 
 
+def _callable_azure_ad_token_provider(provider: object) -> Callable[[], str | Awaitable[str]] | None:
+    return _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(provider) if callable(provider) else None
+
+
 class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object):
     api_version: ReadOnly[str]
     base_url: ReadOnly[str]
@@ -192,8 +197,22 @@ def _exception_chain(error: BaseException) -> Iterator[BaseException]:
         link = link.__cause__ or link.__context__
 
 
+_REQUEST_FAILURE_TYPES: Final = (BaseLLMException, APIError, httpx.HTTPError)
+
+
+def _belongs_to_the_same_attempt(link: BaseException) -> bool:
+    return isinstance(link, EntraRejectionError) or not isinstance(link, _REQUEST_FAILURE_TYPES)
+
+
+def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]:
+    yield error
+    yield from itertools.takewhile(_belongs_to_the_same_attempt, itertools.islice(_exception_chain(error), 1, None))
+
+
 def entra_rejection_in(error: BaseException) -> EntraRejectionError | None:
-    return next((link for link in _exception_chain(error) if isinstance(link, EntraRejectionError)), None)
+    return next(
+        (link for link in _links_within_the_failed_attempt(error) if isinstance(link, EntraRejectionError)), None
+    )
 
 
 def caused_by_entra_rejection(error: BaseException) -> bool:
@@ -230,7 +249,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict:
         openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"]
     if "x-ratelimit-remaining-tokens" in headers:
         openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"]
-    llm_response_headers: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()}
+    llm_response_headers: Final = {f"llm_provider-{k}": v for k, v in headers.items()}
 
     return {**llm_response_headers, **openai_headers}
 
@@ -850,9 +869,7 @@ class BaseAzureLLM(BaseOpenAILLM):
         api_version: str | None,
         is_async: bool,
     ) -> dict:
-        azure_ad_token_provider = _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(
-            litellm_params.get("azure_ad_token_provider")
-        )
+        azure_ad_token_provider = _callable_azure_ad_token_provider(litellm_params.get("azure_ad_token_provider"))
         # If we have api_key, then we have higher priority
         azure_ad_token = litellm_params.get("azure_ad_token")
 

tests/unit/llms/azure/test_azure_common_utils.py

@@ -2971,3 +2971,23 @@ def test_token_provider_error_without_azure_core_installed_keeps_the_generic_tok
         get_azure_ad_token(
             GenericLiteLLMParams(azure_ad_token_provider=_raising_token_provider(ValueError("token service down")))
         )
+
+
+def test_image_generation_with_an_api_key_ignores_a_non_callable_azure_ad_token_provider():
+    import respx
+
+    with respx.mock:
+        image_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock(
+            return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]})
+        )
+        response: Final = litellm.image_generation(
+            model="azure/gpt-image-1",
+            prompt="a red dot",
+            api_base="https://litellm-qa.openai.azure.com",
+            api_version="2025-04-01-preview",
+            api_key="configured-key",
+            azure_ad_token_provider="not-a-callable",
+        )
+
+    assert image_route.calls.last.request.headers["api-key"] == "configured-key"
+    assert response.data[0].b64_json == "aGk="

tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py

@@ -7360,6 +7360,60 @@ class TestRouterModelRelayUpstreamContract:
         assert not deployment_route.called
 
 
+    @pytest.mark.asyncio
+    async def test_azure_relay_keeps_a_later_deployment_failure_after_an_entra_rejection(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        from azure.core.exceptions import ClientAuthenticationError
+        from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+        def rejected_token() -> str:
+            raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.")
+
+        monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+        router: Final = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "gpt-5",
+                    "litellm_params": {
+                        "model": "azure/gpt-5",
+                        "api_base": "https://litellm-qa.openai.azure.com",
+                        "api_version": "2025-04-01-preview",
+                        "azure_ad_token_provider": rejected_token,
+                        "order": 1,
+                    },
+                },
+                {
+                    "model_name": "gpt-5",
+                    "litellm_params": {
+                        "model": "azure/gpt-5",
+                        "api_base": "https://litellm-qa-unreachable.openai.azure.com",
+                        "api_version": "2025-04-01-preview",
+                        "api_key": "configured-key",
+                        "order": 2,
+                    },
+                },
+            ],
+            num_retries=1,
+        )
+        self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False})
+
+        with respx.mock:
+            unreachable_route: Final = respx.post(url__startswith="https://litellm-qa-unreachable.openai.azure.com").mock(
+                side_effect=httpx.ConnectError("Connection refused")
+            )
+            with pytest.raises(BaseLLMException) as excinfo:
+                await azure_proxy_route(
+                    endpoint="openai/deployments/gpt-5/chat/completions",
+                    request=self._request("application/json"),
+                    fastapi_response=MagicMock(spec=Response),
+                    user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
+                )
+
+        assert unreachable_route.called
+        assert excinfo.value.status_code == 500
+
+
 NIM_INFER_BODY = {
     "input": [
         {"type": "image_url", "url": "data:image/png;base64,AAAA"},

tests/unit/test_router/test_router.py

@@ -25791,6 +25791,63 @@ async def test_router_auth_failure_entra_did_not_cause_cools_down_despite_a_call
 
     assert excinfo.value.all_deployments_in_cooldown is True
 
+
+@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
+async def test_router_fallback_auth_failure_after_an_entra_rejection_still_cools_down_the_fallback(monkeypatch):
+    from azure.core.exceptions import ClientAuthenticationError
+
+    def rejected_token() -> str:
+        raise ClientAuthenticationError(message="AADSTS7000215: Invalid client secret provided.")
+
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    router = Router(
+        model_list=[
+            {
+                "model_name": "entra-model",
+                "litellm_params": {
+                    "model": "azure/gpt-4.1",
+                    "api_base": "https://litellm-qa.openai.azure.com",
+                    "api_version": "2025-04-01-preview",
+                    "azure_ad_token_provider": rejected_token,
+                },
+                "model_info": {"id": "entra-deployment"},
+            },
+            {
+                "model_name": "key-model",
+                "litellm_params": {
+                    "model": "azure/gpt-4.1",
+                    "api_base": "https://litellm-qa-fallback.openai.azure.com",
+                    "api_version": "2025-04-01-preview",
+                    "api_key": "revoked-key",
+                },
+                "model_info": {"id": "key-deployment"},
+            },
+        ],
+        fallbacks=[{"entra-model": ["key-model"]}],
+    )
+
+    with respx.mock:
+        fallback_route: Final = respx.post(url__startswith="https://litellm-qa-fallback.openai.azure.com").mock(
+            return_value=httpx.Response(
+                401, json={"error": {"code": "401", "message": "Access denied due to invalid subscription key."}}
+            )
+        )
+        with pytest.raises(litellm.AuthenticationError):
+            await router.acompletion(
+                model="entra-model",
+                messages=[{"role": "user", "content": "hi"}],
+                client_secret="not-the-credential",
+            )
+
+        with pytest.raises(litellm.RouterRateLimitError) as excinfo:
+            await router.acompletion(
+                model="key-model",
+                messages=[{"role": "user", "content": "hi"}],
+            )
+
+    assert fallback_route.call_count == 1
+    assert excinfo.value.all_deployments_in_cooldown is True
+
 @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
 def test_router_get_async_openai_model_client():
     router = Router(