From 784754ca4837b1cd4c27afb44aa221c7029dc549 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:44:55 -0800 Subject: [PATCH] fix(azure)!: answer 401 when Entra rejects the deployment's credential (#49863) * fix(azure): answer 401 when Entra rejects the deployment's credential A ClientAuthenticationError from the Entra token provider (wrong, expired or rotated client secret) used to surface as a retryable 500 carrying the app's client id and Entra trace ids. Map it to a 401 AuthenticationError that keeps the AADSTS code and first sentence with GUIDs redacted, on both the direct token fetch and the provider handed to the OpenAI SDK client. Network failures stay 500. * fix(azure): keep network failures to Entra a retryable 500 azure-identity wraps a transport failure in ClientAuthenticationError, so the 401 mapping now skips errors whose cause chain holds ServiceRequestError or ServiceResponseError * fix(azure): cover async token providers and Entra 5xx, and scope caller-sent Entra credentials * fix(azure): map only AADSTS-coded Entra rejections to 401 and keep caller credential failures out of cooldown An Entra ClientAuthenticationError now answers 401 only when its message carries an AADSTS code, its HTTP status is not 429 or 5xx, and the code is not the AADSTS50196 throttle, so a chained credential's outage or throttle stays a retryable error. Caller-sent Entra fields no longer change the deployment id; a 401 from a value the caller sent skips the shared deployment's cooldown instead. The Azure AI Foundry agent token fetch maps the same rejection, and the A2A endpoint answers it with HTTP 401. * fix(router): count router default_litellm_params as configured Entra credentials when deciding cooldown * fix(azure): keep transient Entra codes retryable and scope 401 handling to Entra rejections ClientSecretCredential drops Entra's HTTP status, so throttling and transient codes (AADSTS50033, 50196, 90006, 90024, 90033, 90055) now stay retryable instead of mapping to 401. The A2A endpoint answers 401 only for an AuthenticationError caused by an Entra token rejection, and the router skips cooldown for a caller-sent Entra field only on Azure-family deployments. The azure/ Decisions route and the Bing grounding search tool now map an environment-credential rejection to the same 401. * fix(azure): map only known Entra credential rejections to 401 and cover the OIDC token exchange Unlisted AADSTS codes keep the retryable 500, the router skips the 401 cooldown only when an Entra rejection caused it, and a federated token Entra rejects answers the same 401 on Azure chat, Responses and A2A agents * fix(azure): answer 401 for Entra rejections on the /azure passthrough and token counting The fallback hop skips cooldown when Entra rejected a credential the caller sent, Prometheus counts a deployment's own Entra rejection as a deployment failure again, and an OIDC exchange error keeps Entra's headers so Retry-After reaches the router. * fix(azure): keep a no-op caller Entra field from skipping cooldown and redact the passthrough 401 message A caller-sent Entra field now skips the deployment's cooldown only when it differs from the deployment, its named credential, the router defaults, the AZURE_* environment variables and the default scope. An empty value counts as unset, so a deployment whose own secret is wrong still cools down when the caller's field changes nothing about the token request. The /azure passthrough's Entra 401 message goes through the same internal hostname redaction as other client-facing errors. * fix(router): compare a caller's Entra field with the value the deployment uses A caller-sent Entra field skipped the cooldown only when it matched no configured source, so a value equal to an environment default the deployment overrides (a stale AZURE_CLIENT_SECRET, say) counted as the deployment's own and let a caller's rejected credential cool down a working deployment. The check now resolves the value the deployment would use without the caller, in order from the router defaults, the deployment, its named credential, the AZURE_* environment variables and the default scope, and skips the cooldown only when the caller's value differs from it. * fix(azure): follow the settings merge for empty Entra fields and keep non-rejection OIDC exchange errors unchanged * fix(azure): count only the failed attempt's Entra rejection and ignore a non-callable token provider * fix(azure): keep the streaming Entra rejection inside the failed attempt --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/prometheus.py | 3 +- litellm/llms/azure/common_utils.py | 211 +++++++- .../llms/azure/decisions/transformation.py | 3 +- litellm/llms/azure/search/transformation.py | 5 +- .../anthropic/count_tokens/handler.py | 3 + litellm/llms/azure_ai/common_utils.py | 67 ++- .../proxy/agent_endpoints/a2a_endpoints.py | 3 + .../llm_passthrough_endpoints.py | 13 +- litellm/router_api.py | 3 + litellm/router_python_core.py | 16 +- litellm/router_utils/cooldown_handlers.py | 39 ++ .../router_utils/fallback_event_handlers.py | 9 + .../test_prometheus_invalid_key_filtering.py | 28 + .../azure/decisions/test_transformation.py | 23 + ...st_bing_grounding_search_transformation.py | 36 ++ .../llms/azure/test_azure_common_utils.py | 493 +++++++++++++++++- .../count_tokens/test_token_counter.py | 73 +++ .../llms/azure_ai/test_azure_ai_entra_auth.py | 51 ++ .../agent_endpoints/test_a2a_endpoints.py | 147 ++++++ .../test_llm_pass_through_endpoints.py | 172 ++++++ .../test_fallback_event_handlers.py | 57 ++ tests/unit/test_router/test_router.py | 435 ++++++++++++++++ 22 files changed, 1858 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1c9a443b546..fc1a601f9da 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -54,6 +54,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_cust from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.models.end_user import LiteLLM_EndUserTable @@ -448,7 +449,7 @@ def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str, return not ( isinstance(exception, CallerCredentialAuthenticationError) or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)) - or logged_call.get("api_call_start_time") is None + or (logged_call.get("api_call_start_time") is None and not caused_by_entra_rejection(exception)) or _used_a_caller_credential(logged_call) ) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6bb6d584385..16c8bf6f9f1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,15 +1,16 @@ import asyncio import hashlib +import itertools import json import os import re -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import closing from types import MappingProxyType -from typing import Final, Literal, NamedTuple, cast +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 @@ -20,7 +21,7 @@ from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler, MaskedHTTPStatusError from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -52,6 +53,13 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object): _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams) +_AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitable[str]] | None]( + Callable[[], str | Awaitable[str]] | None +) + + +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): @@ -81,6 +89,160 @@ class AzureOpenAIError(BaseLLMException): ) +class EntraRejectionError(AzureOpenAIError): + pass + + +_GUID: Final = re.compile(r"[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}") +_AADSTS_REASON: Final = re.compile(r"\bAADSTS(\d+)\b.*?(?:\.(?=[\s'\"\\]|$)|$)", re.MULTILINE) +_ENTRA_CREDENTIAL_REJECTION_CODES: Final = frozenset( + { + 50012, + 50013, + 50027, + 50034, + 50053, + 50055, + 50057, + 50059, + 50076, + 50079, + 50126, + 50128, + 53003, + 65001, + 70001, + 70002, + 70011, + 70021, + 70025, + 70052, + 90002, + 90094, + 500011, + 500014, + 700016, + 700023, + 700024, + 700025, + 700027, + 700211, + 700213, + 700223, + 700238, + 7000112, + 7000215, + 7000218, + 7000222, + } +) + + +def _entra_rejection_reason(message: str) -> str | None: + reason: Final = _AADSTS_REASON.search(message) + if reason is None or int(reason.group(1)) not in _ENTRA_CREDENTIAL_REJECTION_CODES: + return None + return _GUID.sub("", reason.group(0)) + + +def _entra_rejection(reason: str) -> EntraRejectionError: + return EntraRejectionError(status_code=401, message=f"Failed to get Azure AD token: {reason}") + + +def entra_rejection_as_auth_error(error: Exception) -> EntraRejectionError | None: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return None + if not isinstance(error, ClientAuthenticationError): + return None + status_code: Final = error.status_code + if status_code is not None and (status_code == 429 or status_code >= 500): + return None + reason: Final = _entra_rejection_reason(str(error)) + return None if reason is None else _entra_rejection(reason) + + +def _raise_entra_rejection_as_auth_error(error: Exception) -> NoReturn: + auth_error: Final = entra_rejection_as_auth_error(error) + if auth_error is None: + raise error + raise auth_error from error + + +def entra_token(provider: Callable[[], str]) -> str: + try: + return provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def entra_rejection_types() -> tuple[type[Exception], ...]: + try: + from azure.core.exceptions import ClientAuthenticationError + except ImportError: + return () + return (ClientAuthenticationError,) + + +_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16 + + +def _exception_chain(error: BaseException) -> Iterator[BaseException]: + link: BaseException | None = error # rebind-ok: advances one link of the cause chain per iteration + for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH): + if link is None: + return + yield link + link = link.__cause__ or link.__context__ + + +_REQUEST_FAILURE_TYPES: Final = (APIError, httpx.HTTPError) + + +def _ends_the_failed_attempt(link: BaseException, failed_attempt: BaseException) -> bool: + return not isinstance(link, EntraRejectionError) and ( + isinstance(link, _REQUEST_FAILURE_TYPES) or type(link) is type(failed_attempt) + ) + + +def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]: + yield error + yield from itertools.takewhile( + lambda link: not _ends_the_failed_attempt(link, error), itertools.islice(_exception_chain(error), 1, None) + ) + + +def entra_rejection_in(error: BaseException) -> 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: + return entra_rejection_in(error) is not None + + +async def _awaited_token(token: Awaitable[str]) -> str: + try: + return await token + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + + +def _entra_rejections_as_auth_errors( + provider: Callable[[], str | Awaitable[str]], +) -> Callable[[], str | Awaitable[str]]: + def call_provider() -> str | Awaitable[str]: + try: + token: Final = provider() + except entra_rejection_types() as error: + _raise_entra_rejection_as_auth_error(error) + return token if isinstance(token, str) else _awaited_token(token) + + return call_provider + + def process_azure_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "x-ratelimit-limit-requests" in headers: @@ -91,7 +253,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} @@ -257,22 +419,34 @@ def _cached_oidc_token(cache_key: str) -> str | None: return cached if isinstance(cached, str) else None +def _oidc_exchange_rejection_reason(error: MaskedHTTPStatusError) -> str | None: + if error.response.status_code not in (400, 401): + return None + return _entra_rejection_reason(error.text or error.response.text) + + def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str: cached_access_token: Final = _cached_oidc_token(cache_key) if cached_access_token is not None: return cached_access_token with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: - req_token: Final = client.post( - token_url, - data={ - "client_id": client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, - ) + try: + req_token: Final = client.post( + token_url, + data={ + "client_id": client_id, + "grant_type": "client_credentials", + "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": oidc_token, + }, + ) + except MaskedHTTPStatusError as error: + reason: Final = _oidc_exchange_rejection_reason(error) + if reason is None: + raise + raise _entra_rejection(reason) from error if req_token.status_code != 200: raise AzureOpenAIError( @@ -483,6 +657,9 @@ def get_azure_ad_token( raise except Exception as e: verbose_logger.error("Error calling Azure AD token provider: %s", e) + auth_error: Final = entra_rejection_as_auth_error(e) + if auth_error is not None: + raise auth_error from e raise RuntimeError(f"Failed to get Azure AD token: {e}") from e return azure_ad_token @@ -696,7 +873,7 @@ class BaseAzureLLM(BaseOpenAILLM): api_version: str | None, is_async: bool, ) -> dict: - azure_ad_token_provider = 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") @@ -775,7 +952,7 @@ class BaseAzureLLM(BaseOpenAILLM): azure_client_params["timeout"] = timeout if azure_ad_token_provider is not None: - azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider + azure_client_params["azure_ad_token_provider"] = _entra_rejections_as_auth_errors(azure_ad_token_provider) # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client diff --git a/litellm/llms/azure/decisions/transformation.py b/litellm/llms/azure/decisions/transformation.py index 3fb40e39702..bdbd784185f 100644 --- a/litellm/llms/azure/decisions/transformation.py +++ b/litellm/llms/azure/decisions/transformation.py @@ -8,6 +8,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, cached_azure_ad_token_refresh_provider, + entra_token, ) from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig from litellm.secret_managers.main import get_secret_str @@ -61,7 +62,7 @@ class AzureDecisionsConfig(OpenAIDecisionsConfig): scope: Final = ( params.azure_scope or get_secret_str("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default" ) - azure_headers["Authorization"] = f"Bearer {cached_azure_ad_token_refresh_provider(scope)()}" + azure_headers["Authorization"] = f"Bearer {entra_token(cached_azure_ad_token_refresh_provider(scope))}" return azure_headers async def avalidate_environment( diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index e255a667606..6d0241d4d3b 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -36,7 +36,7 @@ import httpx from pydantic import ConfigDict, ValidationError from litellm.litellm_core_utils.url_utils import ensure_api_route -from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider +from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider, entra_rejection_as_auth_error from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.responses import ( ResponsesSearchAnnotation, @@ -270,6 +270,9 @@ class BingGroundingSearchConfig(BaseSearchConfig): try: return minter() except Exception as e: + rejection: Final = entra_rejection_as_auth_error(e) + if rejection is not None: + raise rejection from e raise ValueError( f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} " f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / " diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 97ee44bd3bb..3609871fdf8 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -16,6 +16,7 @@ from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.azure.common_utils import EntraRejectionError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, ) @@ -135,6 +136,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): except AnthropicError: # Re-raise Anthropic exceptions as-is raise + except EntraRejectionError as rejection: + raise AnthropicError(status_code=rejection.status_code, message=rejection.message) from rejection except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code verbose_logger.error("HTTP error in CountTokens handler: %s", e) diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 147bf69e91e..1635b9b5679 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,7 @@ import asyncio -from collections.abc import Mapping +import os +from collections.abc import Callable, Mapping +from functools import partial from types import MappingProxyType from typing import Final, Literal from urllib.parse import urlparse @@ -52,6 +54,28 @@ AZURE_ENTRA_CREDENTIAL_PARAM_KEYS: Final = frozenset({"azure_ad_token", "client_ AZURE_ENTRA_LITELLM_PARAM_KEYS: Final = AZURE_ENTRA_CREDENTIAL_PARAM_KEYS | frozenset( {"tenant_id", "client_id", "azure_username", "azure_scope"} ) +AZURE_DEFAULT_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_AZURE_ENTRA_ENVIRONMENT_VARIABLES: Final = MappingProxyType( + { + "tenant_id": "AZURE_TENANT_ID", + "client_id": "AZURE_CLIENT_ID", + "client_secret": "AZURE_CLIENT_SECRET", + "azure_username": "AZURE_USERNAME", + "azure_password": "AZURE_PASSWORD", + "azure_ad_token": "AZURE_AD_TOKEN", + } +) + + +def azure_entra_environment_defaults() -> Mapping[str, object]: + return MappingProxyType( + { + **{key: os.getenv(variable) for key, variable in _AZURE_ENTRA_ENVIRONMENT_VARIABLES.items()}, + "azure_scope": os.getenv("AZURE_SCOPE", AZURE_DEFAULT_SCOPE), + } + ) + + AZURE_ENTRA_CREDENTIAL_HELP: Final = ( "Set `tenant_id` + `client_id` + `client_secret`, `azure_ad_token` (an `oidc/` token also needs " "`tenant_id` + `client_id`), or `client_id` + `azure_username` + `azure_password` in the agent's `litellm_params`" @@ -70,6 +94,17 @@ def _resolve_config_secret(value: object) -> str | None: return get_secret_str(value) if value.startswith("os.environ/") else value +def _agent_entra_token(provider: Callable[[], str], model: str) -> str: + from litellm.llms.azure.common_utils import EntraRejectionError, entra_token + + try: + return entra_token(provider) + except EntraRejectionError as rejection: + raise litellm.AuthenticationError( + message=rejection.message, llm_provider="azure_ai", model=model + ) from rejection + + def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: """Mints the Entra bearer from the agent's own litellm_params, never from process-wide AZURE_* env vars.""" from litellm.llms.azure.common_utils import ( @@ -88,18 +123,32 @@ def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str: azure_username: Final = resolved["azure_username"] azure_password: Final = resolved["azure_password"] azure_ad_token: Final = resolved["azure_ad_token"] + model: Final = str(litellm_params.get("model") or "") if tenant_id and client_id and client_secret: - return get_azure_ad_token_from_entra_id( - tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_entra_id( + tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope + ), + model, + ) if client_id and azure_username and azure_password: - return get_azure_ad_token_from_username_password( - client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope - )() + return _agent_entra_token( + get_azure_ad_token_from_username_password( + client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope + ), + model, + ) federated: Final = azure_ad_token is not None and azure_ad_token.startswith("oidc/") if azure_ad_token and federated and tenant_id and client_id: - return get_azure_ad_token_from_oidc( - azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope + return _agent_entra_token( + partial( + get_azure_ad_token_from_oidc, + azure_ad_token=azure_ad_token, + azure_client_id=client_id, + azure_tenant_id=tenant_id, + scope=scope, + ), + model, ) if azure_ad_token and not federated: return azure_ad_token diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 3dfaa7f213e..f0a4d47f1e4 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.azure.common_utils import caused_by_entra_rejection from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy @@ -1202,6 +1203,8 @@ async def invoke_agent_a2a( pass if isinstance(e, litellm.BadRequestError): return _jsonrpc_error(body.get("id"), -32602, e.message, 400) + if isinstance(e, litellm.AuthenticationError) and caused_by_entra_rejection(e): + return _jsonrpc_error(body.get("id"), -32000, e.message, 401) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 64cde072dc5..627164abb93 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,7 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider -from litellm._logging import verbose_proxy_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, @@ -58,6 +58,7 @@ from litellm.llms.anthropic.common_utils import ( anthropic_pass_through_api_base, merge_anthropic_beta_headers, ) +from litellm.llms.azure.common_utils import entra_rejection_in from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, @@ -2381,6 +2382,16 @@ async def _relay_router_model( ) except httpx.HTTPStatusError as upstream_error: return await _relay_upstream_response(upstream_error.response) + except Exception as error: + rejection: Final = entra_rejection_in(error) + if rejection is None: + raise + raise ProxyException( + message=redact_internal_details_from_client_message(rejection.message), + type=ProxyErrorTypes.auth_error, + param=None, + code=rejection.status_code, + ) from error if not is_streaming_request: return await _relay_upstream_response(cast(httpx.Response, result)) diff --git a/litellm/router_api.py b/litellm/router_api.py index 3007b5345c7..f22e9b1f351 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1402,6 +1402,9 @@ class RoutingContext: def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: return self._router._params_sent_to_deployment(request, deployment) + def default_litellm_params(self) -> Mapping[str, object]: + return self._router._active_default_litellm_params() + def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e251248d675..c0c5d74a88a 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -306,6 +306,7 @@ from litellm.router_utils.cooldown_handlers import ( get_cooldown_deployment_info, get_cooldown_deployments, is_background_response_poll_not_found, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, @@ -10309,8 +10310,9 @@ class PythonCore(DeprecatedRouterMembers): return False exception_status: Final = getattr(exception, "status_code", "") + failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs) - if is_caller_timeout_408(kwargs, exception_status): + if is_caller_timeout_408(failure_call_details, exception_status): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " "A timeout the caller set caused this 408, not the deployment's health." @@ -10334,6 +10336,18 @@ class PythonCore(DeprecatedRouterMembers): "Caller-scoped OAuth authentication failed, not the deployment." ) return False + if is_caller_entra_credential_failure( + caller_failure_deployment, + self._active_default_litellm_params(), + failure_call_details, + exception_status, + exception, + ): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "Entra ID rejected a credential the caller sent, not the deployment's own." + ) + return False exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e50d5bb47a6..7a8a645ad13 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -847,6 +847,45 @@ def is_stream_stall(exception: BaseException) -> bool: return isinstance(exception, litellm.Timeout) and isinstance(exception.__cause__, MidStreamFallbackError) +def is_caller_entra_credential_failure( + deployment: Deployment | None, + router_defaults: Mapping[str, object], + model_call_details: Mapping[str, object], + exception_status: str | int, + exception: BaseException, +) -> bool: + if cast_exception_status_to_int(exception_status) != 401 or deployment is None: + return False + + from litellm.llms.azure.common_utils import caused_by_entra_rejection + + if not caused_by_entra_rejection(exception): + return False + from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, azure_entra_environment_defaults + + call_params: Final = _CREDENTIAL_VALUES_ADAPTER.validate_python(model_call_details.get("litellm_params") or {}) + + credential_name: Final = deployment.litellm_params.litellm_credential_name + configured_by_precedence: Final[tuple[Mapping[str, object], ...]] = ( + router_defaults, + _CREDENTIAL_VALUES_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)), + _CREDENTIAL_VALUES_ADAPTER.validate_python( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ), + ) + environment: Final = azure_entra_environment_defaults() + return any( + call_params.get(key) is not None + and (call_params.get(key) or environment.get(key)) + != (_first_configured(configured_by_precedence, key) or environment.get(key)) + for key in AZURE_ENTRA_LITELLM_PARAM_KEYS + ) + + +def _first_configured(sources: tuple[Mapping[str, object], ...], key: str) -> object: + return next((value for source in sources if (value := source.get(key)) is not None), None) + + def is_caller_scoped_auth_failure(deployment: Deployment | None, exception_status: str | int) -> bool: if cast_exception_status_to_int(exception_status) not in (401, 403) or deployment is None: return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index c929cca44be..e8f832b4b2d 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -28,6 +28,7 @@ from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, is_advisor_orchestration_failure, + is_caller_entra_credential_failure, is_caller_scoped_auth_failure, is_caller_timeout_408, is_stream_stall, @@ -132,6 +133,14 @@ def _trigger_cooldown_for_failed_deployment( deployment_id, ) return + if is_caller_entra_credential_failure( + deployment, ctx.default_litellm_params(), model_call_details, exception_status, exception + ): + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment %s: Entra ID rejected a credential the caller sent.", + deployment_id, + ) + return # Priority: deployment config > response header > router default, matching # Router.deployment_callback_on_failure's precedence for the primary path. diff --git a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py index 453fd28636d..3f3e2e0d465 100644 --- a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py @@ -512,6 +512,20 @@ def _provider_auth_failure_kwargs() -> dict[str, object]: } +def _deployment_entra_rejection_kwargs() -> dict[str, object]: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + return { + key: value for key, value in _provider_auth_failure_kwargs().items() if key not in CALLED_THE_PROVIDER + } | {"exception": raised} + + def _metric_values(name: str, model_id: str) -> list[float]: return [ sample.value @@ -541,6 +555,20 @@ class TestProviderAuthFailures: assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0] + @pytest.mark.asyncio + async def test_log_failure_event_records_entra_rejecting_the_deployment_credential_before_the_call( + self, prometheus_logger + ): + await prometheus_logger.async_log_failure_event( + kwargs=_deployment_entra_rejection_kwargs(), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] + def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger): prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs()) diff --git a/tests/unit/llms/azure/decisions/test_transformation.py b/tests/unit/llms/azure/decisions/test_transformation.py index b6de6084b9c..3f5e97ee000 100644 --- a/tests/unit/llms/azure/decisions/test_transformation.py +++ b/tests/unit/llms/azure/decisions/test_transformation.py @@ -125,6 +125,29 @@ async def test_default_credential_refreshes_without_opt_in( ) +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("systemone", (False, True)) +async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + monkeypatch: pytest.MonkeyPatch, is_async: bool, systemone: bool +) -> None: + from azure.core.exceptions import ClientAuthenticationError + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + await call_decisions(is_async, systemone, {}) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + + @pytest.mark.parametrize("kind", ("service-principal", "oidc", "provider")) async def test_azure_decisions_reuses_existing_azure_credential_resolution( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, kind: str diff --git a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py index dfd1474211f..53f426d8505 100644 --- a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py @@ -466,6 +466,12 @@ def _mock_response(): response.content = json.dumps(MOCK_BING_GROUNDING_RESPONSE).encode() return response +async def _search_what_is_litellm(is_async: bool) -> object: + if is_async: + return await litellm.asearch(query="what is litellm", search_provider="bing_grounding") + return litellm.search(query="what is litellm", search_provider="bing_grounding") + + @pytest.mark.usefixtures("_vcr_outcome_gate") class TestBingGroundingSearchTransformation: """ @@ -546,6 +552,36 @@ class TestBingGroundingSearchTransformation: } ] + @pytest.mark.parametrize("is_async", (False, True)) + async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason( + self, monkeypatch: pytest.MonkeyPatch, is_async: bool + ): + from azure.core.exceptions import ClientAuthenticationError + + from litellm.llms.azure import common_utils + + app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" + credential = Mock(spec=["get_token"]) + credential.get_token.side_effect = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'." + ) + monkeypatch.delenv("BING_GROUNDING_TOKEN") + for name in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential)) + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + try: + with pytest.raises(litellm.AuthenticationError) as excinfo: + await _search_what_is_litellm(is_async) + finally: + common_utils._cached_azure_ad_token_refresh_provider.cache_clear() + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert app_client_id not in str(excinfo.value) + @pytest.mark.asyncio async def test_bing_grounding_asearch(self): with patch( # test-quality-ok: litellm.asearch has no client injection seam diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 0a41861e088..d66ba3f324e 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1,16 +1,23 @@ +import asyncio +import io import json import os +import sys import threading import traceback import uuid +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from queue import SimpleQueue -from typing import Callable, Final, Optional +from typing import Awaitable, Callable, Final, Optional from urllib.parse import parse_qs from unittest.mock import ANY, MagicMock, patch import httpx +import openai import pytest +import requests +from azure.core.pipeline.transport import RequestsTransport import litellm from litellm.llms.azure.common_utils import ( @@ -2500,3 +2507,487 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m closer.reap() assert wrapper.is_closed() is True + + +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" +_ENTRA_REJECTIONS: Final = { + "wrong_secret": ( + "AADSTS7000215", + "Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent in the " + "request is the client secret value, not the client secret ID, for a secret added to app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID} Correlation ID: ecfd2c73-ec0a-474c-bba7-205be903159d " + "Timestamp: 2026-10-29 07:19:40Z", + ), + "expired_secret": ( + "AADSTS7000222", + f"Authentication failed: AADSTS7000222: The provided client secret keys for app '{_ENTRA_APP_CLIENT_ID}' " + f"are expired. Visit the Azure portal to create new keys for your app. Trace ID: {_ENTRA_TRACE_ID} " + "Correlation ID: 354f9b57-5fed-47e4-a67d-89a7c56f6ddc Timestamp: 2026-10-29 07:19:50Z", + ), +} + + +def _raising_token_provider(error: Exception) -> Callable[[], str]: + def token_provider() -> str: + raise error + + return token_provider + + +def _azure_chat(token_provider: Callable[[], str]) -> object: + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + max_retries=0, + ) + + +def _azure_async_chat(token_provider: Callable[[], str]) -> object: + return asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=token_provider, + ) + ) + + +def _azure_responses(token_provider: Callable[[], str]) -> object: + return litellm.responses( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + input="hi", + azure_ad_token_provider=token_provider, + ) + + +def _azure_ai_decisions(token_provider: Callable[[], str]) -> object: + return litellm.systemone( + model="azure_ai/decision-1-qa", + api_base="https://litellm-qa.services.ai.azure.com", + state={"ticket": "The login page returns a 500 error after deploy"}, + questions={"is_bug": {"type": "noul", "instructions": "Is this a software bug report?"}}, + azure_ad_token_provider=token_provider, + ) + + +_AZURE_CALLS: Final = { + "chat": _azure_chat, + "async_chat": _azure_async_chat, + "responses": _azure_responses, + "decisions": _azure_ai_decisions, +} + + +@pytest.fixture +def no_azure_api_key(monkeypatch): + for env_var in ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "azure_key", None) + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize(("aadsts_code", "entra_message"), _ENTRA_REJECTIONS.values(), ids=_ENTRA_REJECTIONS.keys()) +def test_entra_rejecting_the_credential_answers_401_with_only_the_aadsts_reason( + no_azure_api_key, call, aadsts_code, entra_message +): + from azure.core.exceptions import ClientAuthenticationError + + with pytest.raises(litellm.AuthenticationError) as excinfo: + call(_raising_token_provider(ClientAuthenticationError(message=entra_message))) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +@pytest.mark.parametrize("api_version", ["2025-04-01-preview", "v1"]) +def test_entra_rejecting_an_async_token_provider_answers_401(no_azure_api_key, api_version): + from azure.core.exceptions import ClientAuthenticationError + + aadsts_code, entra_message = _ENTRA_REJECTIONS["wrong_secret"] + + async def token_provider() -> str: + raise ClientAuthenticationError(message=entra_message) + + async_token_provider: Final[Callable[[], Awaitable[str]]] = token_provider + + with pytest.raises(litellm.AuthenticationError) as excinfo: + asyncio.run( + litellm.acompletion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version=api_version, + messages=[{"role": "user", "content": "hi"}], + azure_ad_token_provider=async_token_provider, + ) + ) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def _azure_identity_wrapped_token_provider(inner_error: Exception) -> Callable[[], str]: + from azure.core.exceptions import ClientAuthenticationError + + def token_provider() -> str: + try: + raise inner_error + except Exception as error: + raise ClientAuthenticationError(message=f"Authentication failed: {error}") from error + + return token_provider + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("failed_leg", ["request", "response"]) +def test_unreachable_entra_endpoint_stays_a_retryable_server_error(no_azure_api_key, call, failed_leg): + from azure.core.exceptions import ServiceRequestError, ServiceResponseError + + transport_error: Final = {"request": ServiceRequestError, "response": ServiceResponseError}[failed_leg]( + "('Unable to connect to proxy', NewConnectionError(\"Failed to establish a new connection: " + "[Errno 111] Connection refused\"))" + ) + + with pytest.raises(openai.APIError) as excinfo: + call(_azure_identity_wrapped_token_provider(transport_error)) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +_ENTRA_TENANT_ID: Final = "5b3e1c7a-2f4d-4e8b-9a61-0c2d7f8e9b14" +_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +_ENTRA_OPENID_CONFIGURATION: Final = { + "token_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token", + "authorization_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/authorize", + "issuer": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/v2.0", +} + + +class _FakeEntraAdapter(requests.adapters.BaseAdapter): + def __init__(self, status_code: int | None, body: Mapping[str, object] | str) -> None: + super().__init__() + self.status_code: Final = status_code + self.body: Final = body + + def send(self, request: requests.PreparedRequest, *args: object, **kwargs: object) -> requests.Response: + if "openid-configuration" in str(request.url): + return self._answer(request, 200, _ENTRA_OPENID_CONFIGURATION) + if self.status_code is None: + raise requests.exceptions.ConnectionError("Failed to establish a new connection: Connection refused") + return self._answer(request, self.status_code, self.body) + + def close(self) -> None: + pass + + @staticmethod + def _answer( + request: requests.PreparedRequest, status_code: int, body: Mapping[str, object] | str + ) -> requests.Response: + response: Final = requests.Response() + response.status_code = status_code + response.raw = io.BytesIO((body if isinstance(body, str) else json.dumps(body)).encode()) + response.headers = requests.structures.CaseInsensitiveDict( + {"Content-Type": "application/json", "Retry-After": "0"} + ) + response.url = str(request.url) + response.request = request + return response + + +def _entra_answering(status_code: int | None, body: Mapping[str, object] | str = "") -> RequestsTransport: + session: Final = requests.Session() + session.mount("https://", _FakeEntraAdapter(status_code, body)) + return RequestsTransport(session=session, session_owner=False) + + +def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]: + return { + "error": "invalid_client", + "error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}", + "error_codes": [int(aadsts_code.removeprefix("AADSTS"))], + } + + +def _real_credential_token_provider(credential_kind: str, entra: RequestsTransport) -> Callable[[], str]: + from azure.identity import ( + ChainedTokenCredential, + ClientAssertionCredential, + ClientSecretCredential, + get_bearer_token_provider, + ) + + client_secret: Final = ClientSecretCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, "wrong-secret", transport=entra, disable_instance_discovery=True + ) + credentials: Final = { + "client_secret": client_secret, + "default_azure_credential_chain": ChainedTokenCredential(client_secret), + "workload_identity": ClientAssertionCredential( + _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, lambda: "federated-token", transport=entra, retry_total=0 + ), + } + return get_bearer_token_provider(credentials[credential_kind], _ENTRA_SCOPE) + + +_REAL_CREDENTIAL_REJECTIONS: Final = { + "wrong_secret": ("client_secret", "AADSTS7000215", "Invalid client secret provided."), + "wrong_secret_through_a_credential_chain": ( + "default_azure_credential_chain", + "AADSTS7000215", + "Invalid client secret provided.", + ), + "unknown_app_through_workload_identity": ( + "workload_identity", + "AADSTS700016", + f"Application with identifier '{_ENTRA_APP_CLIENT_ID}' was not found in the directory.", + ), + "unknown_federated_subject_through_workload_identity": ( + "workload_identity", + "AADSTS700213", + "No matching federated identity record found for presented assertion subject 'system:serviceaccount:qa:proxy'.", + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "aadsts_code", "description"), + _REAL_CREDENTIAL_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_REJECTIONS.keys(), +) +def test_entra_rejection_through_a_real_azure_identity_credential_answers_401( + no_azure_api_key, credential_kind, aadsts_code, description +): + entra: Final = _entra_answering(401, _entra_error(aadsts_code, description)) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, entra)) + + assert excinfo.value.status_code == 401 + assert aadsts_code in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +_REAL_CREDENTIAL_NON_REJECTIONS: Final = { + "unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""), + "html_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + "Service Unavailable", + ), + "json_503_through_a_credential_chain": ( + "default_azure_credential_chain", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."), + ), + "tenant_throttled_client_secret": ( + "client_secret", + 429, + _entra_error("AADSTS90055", "There are too many incoming requests."), + ), + "request_budget_exceeded_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90024", "A transient error has occurred. Try again."), + ), + "retryable_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50033", "A transient error has occurred."), + ), + "external_server_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90006", "The service is temporarily unavailable."), + ), + "graph_retryable_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS90090", "The service is temporarily unavailable."), + ), + "token_issuance_error_client_secret": ( + "client_secret", + 400, + _entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."), + ), + "throttled_without_a_code_through_a_credential_chain": ( + "default_azure_credential_chain", + 429, + {"error": "temporarily_unavailable", "error_description": "Too many requests."}, + ), + "json_503_through_workload_identity": ( + "workload_identity", + 503, + _entra_error("AADSTS90033", "A transient error has occurred."), + ), + "throttled_workload_identity_naming_a_rejection_code": ( + "workload_identity", + 429, + _entra_error("AADSTS7000215", "Invalid client secret provided."), + ), +} + + +@pytest.mark.parametrize( + ("credential_kind", "status_code", "body"), + _REAL_CREDENTIAL_NON_REJECTIONS.values(), + ids=_REAL_CREDENTIAL_NON_REJECTIONS.keys(), +) +def test_entra_outage_or_throttle_through_a_real_azure_identity_credential_stays_retryable( + no_azure_api_key, credential_kind, status_code, body +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat(_real_credential_token_provider(credential_kind, _entra_answering(status_code, body))) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + entra_status_code: int, + entra_body: Mapping[str, object], + entra_headers: Mapping[str, str] | None = None, +): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post(f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token").mock( + return_value=httpx.Response(entra_status_code, json=entra_body, headers=entra_headers) + ) + return litellm.completion( + model="azure/gpt-4.1", + api_base="https://litellm-qa.openai.azure.com", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "hi"}], + azure_ad_token="oidc/env/LITELLM_QA_FEDERATED_TOKEN", + tenant_id=_ENTRA_TENANT_ID, + client_id=_ENTRA_APP_CLIENT_ID, + max_retries=0, + ) + + +def test_entra_rejecting_an_oidc_token_answers_401_with_only_the_aadsts_reason(no_azure_api_key, respx_mock, monkeypatch): + rejection: Final = _entra_error( + "AADSTS700213", f"No matching federated identity record found for app '{_ENTRA_APP_CLIENT_ID}'." + ) + + with pytest.raises(litellm.AuthenticationError) as excinfo: + _azure_chat_with_oidc_token(respx_mock, monkeypatch, 400, rejection) + + assert excinfo.value.status_code == 401 + assert "AADSTS700213" in str(excinfo.value) + assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value) + assert _ENTRA_TRACE_ID not in str(excinfo.value) + + +def test_entra_answering_an_oidc_exchange_with_a_transient_code_keeps_its_status( + no_azure_api_key, respx_mock, monkeypatch +): + with pytest.raises(openai.APIError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, monkeypatch, 400, _entra_error("AADSTS90024", "A transient error has occurred. Try again.") + ) + + assert excinfo.value.status_code == 400 + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + +def test_entra_throttling_an_oidc_exchange_keeps_its_retry_after(no_azure_api_key, respx_mock, monkeypatch): + with pytest.raises(litellm.RateLimitError) as excinfo: + _azure_chat_with_oidc_token( + respx_mock, + monkeypatch, + 429, + _entra_error("AADSTS90055", "TenantThrottlingError: There are too many pending requests."), + {"Retry-After": "7"}, + ) + + assert excinfo.value.response.headers.get("retry-after") == "7" + + +class _EntraServiceUnavailableAnswer: + status_code: Final = 503 + reason: Final = "Service Unavailable" + + def text(self, encoding: str | None = None) -> str: + return "Service Unavailable" + + +@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys()) +@pytest.mark.parametrize("credential_kind", ["client_secret", "managed_identity"]) +def test_entra_answering_a_server_error_stays_a_retryable_server_error(no_azure_api_key, call, credential_kind): + from azure.core.exceptions import ClientAuthenticationError + from msal.exceptions import MsalServiceError + + token_providers: Final = { + "client_secret": _azure_identity_wrapped_token_provider( + MsalServiceError("HTTP Error: 503", error=None, error_description=None) + ), + "managed_identity": _raising_token_provider( + ClientAuthenticationError(message="No token received.", response=_EntraServiceUnavailableAnswer()) + ), + } + + with pytest.raises(openai.APIError) as excinfo: + call(token_providers[credential_kind]) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + assert excinfo.value.status_code != 401 + + +def test_token_provider_error_without_azure_core_installed_keeps_the_generic_token_error(monkeypatch): + monkeypatch.setitem(sys.modules, "azure.core.exceptions", None) + + with pytest.raises(RuntimeError, match="Failed to get Azure AD token: token service down"): + 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=" diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index faa33e559ca..7ce5d8b6a02 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -171,6 +171,79 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ ) +@pytest.mark.asyncio +async def test_entra_rejecting_the_deployment_oidc_token_becomes_a_401_error_response_without_counting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 401, + json={ + "error": "invalid_client", + "error_description": "AADSTS700211: No matching federated identity record found for presented " + "assertion issuer 'https://oidc.litellm-qa.invalid'. Trace ID: 0f5c2b1e", + "error_codes": [700211], + }, + ) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "qa-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 401) + assert "AADSTS700211" in (result.error_message or "") + assert not count_route.called + + +@pytest.mark.asyncio +async def test_entra_throttling_the_deployment_oidc_token_exchange_keeps_its_429( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response(429, json={"error": "temporarily_unavailable", "error_codes": [90055]}) + ) + count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={ + "litellm_params": { + "api_base": API_BASE, + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + }, + request_model="azure_ai/claude-sonnet-4-5", + ) + + assert result is not None + assert (result.error, result.status_code) == (True, 429) + assert not count_route.called + + @pytest.mark.asyncio async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py index d0c1c4a6fb2..a8958b301db 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py @@ -254,6 +254,57 @@ def test_agent_entra_values_resolve_os_environ_references(monkeypatch): assert get_azure_ai_agent_entra_token({"azure_ad_token": "os.environ/FOUNDRY_AGENT_AD_TOKEN"}) == "token-from-env" +_AGENT_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_AGENT_CREDENTIALS: Final = { + "service_principal": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", + {"tenant_id": "tenant", "client_id": _AGENT_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + ), + "username_password": ( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password", + {"client_id": _AGENT_APP_CLIENT_ID, "azure_username": "agent@example.com", "azure_password": "wrong"}, + ), +} + + +def _raising(error: Exception): + def token() -> str: + raise error + + return token + + +@pytest.mark.parametrize(("token_fetch", "agent_params"), _AGENT_CREDENTIALS.values(), ids=_AGENT_CREDENTIALS.keys()) +def test_agent_entra_rejection_raises_401_with_only_the_aadsts_reason(token_fetch, agent_params): + from azure.core.exceptions import ClientAuthenticationError + + rejection: Final = ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent " + f"in the request is the client secret value, for a secret added to app '{_AGENT_APP_CLIENT_ID}'." + ) + + with patch(token_fetch, return_value=_raising(rejection)): # test-quality-ok: stands in for the Entra token endpoint answering a wrong credential, the shape azure-identity raises + with pytest.raises(litellm.AuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert excinfo.value.status_code == 401 + assert "AADSTS7000215" in str(excinfo.value) + assert _AGENT_APP_CLIENT_ID not in str(excinfo.value) + + +def test_agent_entra_endpoint_unreachable_is_not_an_auth_error(): + from azure.core.exceptions import ClientAuthenticationError + + unreachable: Final = ClientAuthenticationError(message="Authentication failed: Connection refused") + token_fetch, agent_params = _AGENT_CREDENTIALS["service_principal"] + + with patch(token_fetch, return_value=_raising(unreachable)): # test-quality-ok: stands in for an unreachable Entra token endpoint, the shape azure-identity raises + with pytest.raises(ClientAuthenticationError) as excinfo: + get_azure_ai_agent_entra_token(agent_params) + + assert not isinstance(excinfo.value, litellm.AuthenticationError) + + def test_agent_entra_token_failure_names_the_credential_fields(): with pytest.raises(ValueError, match="client_secret"): get_azure_ai_agent_entra_token({"azure_scope": "https://ai.azure.com/.default"}) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 88da07c612b..0aacc1f8bda 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -1268,6 +1268,153 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): assert failure_data.get("agent_id") == "test-agent" +_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2" +_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500" + + +class _JsonRpcError(BaseModel): + model_config = ConfigDict(frozen=True) + + code: int + message: str + + +class _JsonRpcErrorBody(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _JsonRpcError + + +async def _send_message_to_entra_agent( + litellm_params: Mapping[str, str], *patches: AbstractContextManager[object] +) -> tuple[int, _JsonRpcError]: + agent = _make_agent_mock() + agent.litellm_params = dict(litellm_params) + mock_request = _make_request_mock( + "message/send", + {"message": {"role": "user", "parts": [{"kind": "text", "text": "Hello"}], "messageId": "msg-1"}}, + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in (*_base_patches(agent), *patches): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + return response.status_code, _JsonRpcErrorBody.model_validate_json(response.body).error + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_credential_returns_401_with_only_the_aadsts_reason(): + from azure.core.exceptions import ClientAuthenticationError + + def rejected_token() -> str: + raise ClientAuthenticationError( + message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being " + f"sent in the request is the client secret value, for a secret added to app '{_ENTRA_APP_CLIENT_ID}'." + ) + + status_code, error = await _send_message_to_entra_agent( + {"tenant_id": "tenant", "client_id": _ENTRA_APP_CLIENT_ID, "client_secret": "wrong-secret"}, + patch( # test-quality-ok: stands in for the Entra token endpoint answering a wrong secret, the shape azure-identity raises + "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=rejected_token + ), + ) + + assert status_code == 401 + assert "AADSTS7000215" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + + +@pytest.mark.asyncio +async def test_entra_rejecting_the_agent_oidc_token_returns_401_with_only_the_aadsts_reason(monkeypatch, respx_mock): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + respx_mock.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 400, + json={ + "error": "invalid_client", + "error_description": "AADSTS700213: No matching federated identity record found for app " + f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID}", + }, + ) + ) + + status_code, error = await _send_message_to_entra_agent( + { + "tenant_id": "qa-tenant", + "client_id": _ENTRA_APP_CLIENT_ID, + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + } + ) + + assert status_code == 401 + assert error.code == -32000 + assert "AADSTS700213" in error.message + assert _ENTRA_APP_CLIENT_ID not in error.message + assert _ENTRA_TRACE_ID not in error.message + + +@pytest.mark.asyncio +async def test_upstream_authentication_error_without_an_entra_rejection_stays_an_internal_error(): + agent = _make_agent_mock() + mock_handler = MagicMock() + mock_handler.post = AsyncMock( + side_effect=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider="openai", model="gpt-4.1" + ) + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: stands in for the upstream agent rejecting the provider key it was sent + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=mock_handler + ) + ) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=_make_request_mock("tasks/get", {"id": "task-1"}), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"), + ) + + assert response.status_code == 500 + assert json.loads(response.body.decode())["error"]["code"] == -32603 + + @pytest.mark.asyncio async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400(): from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4a76aa77f14..11b07472eaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7241,6 +7241,178 @@ class TestRouterModelRelayUpstreamContract: assert json.loads(result.body) == upstream_body assert result.headers["x-ms-request-id"] == "req-1" + @pytest.mark.asyncio + async def test_azure_relay_answers_401_when_entra_rejects_the_deployment_credential(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + try: + raise EntraRejectionError(status_code=401, message=reason) + except EntraRejectionError as rejection: + raise BaseLLMException(status_code=401, message=str(rejection)) from rejection + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) 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 (excinfo.value.code, excinfo.value.message) == ("401", reason) + + @pytest.mark.asyncio + async def test_azure_relay_entra_rejection_hides_an_internal_issuer_host(self, monkeypatch): + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = ( + "Failed to get Azure AD token: AADSTS700211: No matching federated identity record found for " + "presented assertion issuer 'https://oidc.build.internal:8443/realms/ci'." + ) + + class EntraRejectingRouter: + async def allm_passthrough_route(self, **kwargs): + raise EntraRejectionError(status_code=401, message=reason) + + self._install_router(monkeypatch, EntraRejectingRouter(), {"model": "gpt-5", "stream": False}) + + with pytest.raises(ProxyException) 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 excinfo.value.code == "401" + assert "AADSTS700211" in excinfo.value.message + assert "oidc.build.internal" not in excinfo.value.message + + @pytest.mark.asyncio + async def test_azure_relay_keeps_a_failure_entra_did_not_cause(self, monkeypatch): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + upstream_failure: Final = BaseLLMException(status_code=500, message="Connection reset by Azure") + + class FailingRouter: + async def allm_passthrough_route(self, **kwargs): + raise upstream_failure + + self._install_router(monkeypatch, FailingRouter(), {"model": "gpt-5", "stream": False}) + + 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 excinfo.value is upstream_failure + + @pytest.mark.asyncio + async def test_azure_relay_keeps_entra_throttling_the_oidc_token_exchange_as_a_429(self, monkeypatch): + monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token") + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + 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", + "tenant_id": "throttled-tenant", + "client_id": "qa-client", + "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN", + }, + } + ], + num_retries=0, + ) + self._install_router(monkeypatch, router, {"model": "gpt-5", "stream": False}) + + with respx.mock: + respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock( + return_value=httpx.Response( + 429, + headers={"Retry-After": "7"}, + json={"error": "temporarily_unavailable", "error_codes": [90055]}, + ) + ) + deployment_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock( + return_value=httpx.Response(200, json={"ok": True}) + ) + response: Final = 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 (response.status_code, response.headers.get("retry-after")) == (429, "7") + 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": [ diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 26dbb96c60a..e9dd91d4362 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -302,6 +302,63 @@ def _two_deployment_router() -> Router: ) +def _entra_fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +def _raised_from_an_entra_rejection(failed_deployment_id: str) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided." + try: + raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from ( + EntraRejectionError(status_code=401, message=reason) + ) + except litellm.AuthenticationError as raised: + raised.failed_deployment_id = failed_deployment_id + return raised + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_params", "cooled_down"), + [ + pytest.param({"client_secret": "not-the-credential"}, [], id="caller-secret"), + pytest.param({"client_secret": "configured-secret"}, ["fallback-dep-1"], id="deployment-secret"), + ], +) +async def test_fallback_hop_cools_down_only_a_deployment_entra_rejected_its_own_credential( + call_params: Mapping[str, object], cooled_down: list[str] +) -> None: + router: Final = _entra_fallback_router() + + _trigger_cooldown_for_failed_deployment( + ctx=RoutingContext(router), + kwargs={}, + exception=_raised_from_an_entra_rejection("fallback-dep-1"), + model_call_details={"litellm_params": call_params}, + ) + + assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == cooled_down + + @pytest.mark.asyncio @pytest.mark.parametrize("failure_already_logged", [False, True]) async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index f815b107648..c16c1c62202 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -25836,6 +25836,441 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} +def _caused_by_an_entra_rejection(error: litellm.AuthenticationError) -> litellm.AuthenticationError: + from litellm.llms.azure.common_utils import EntraRejectionError + + try: + raise error from EntraRejectionError(status_code=401, message=error.message) + except litellm.AuthenticationError as raised: + return raised + + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize("caller_credential", ["client_secret", "azure_password", "azure_ad_token"]) +async def test_router_caller_entra_credential_rejection_leaves_the_shared_deployment_routable(caller_credential): + 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", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **{caller_credential: "not-the-credential"}, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_field", + [ + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + pytest.param({"client_id": "environment-client"}, id="environment_client_id"), + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"client_secret": ""}, id="blank_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_the_deployment_credential_leaves_it_routable( + caller_field, monkeypatch +): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + monkeypatch.setenv("AZURE_CLIENT_ID", "environment-client") + monkeypatch.delenv("AZURE_SCOPE", raising=False) + 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", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + "azure_scope": "api://configured-app/.default", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_field, + ) + + response: Final = await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert response.choices[0].message.content == "hi" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("deployment_secret", "router_default_secret"), + [ + pytest.param({"client_secret": "configured-secret"}, None, id="deployment"), + pytest.param({}, {"client_secret": "configured-secret"}, id="router_default_litellm_params"), + ], +) +async def test_router_configured_entra_credential_rejection_cools_down_the_deployment( + deployment_secret, router_default_secret +): + 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", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + **deployment_secret, + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params=router_default_secret, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000222: The provided client secret keys are expired.", + llm_provider="azure", + model="gpt-4.1", + ), + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + assert excinfo.value.all_deployments_in_cooldown is True + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + "caller_fields", + [ + pytest.param({"azure_scope": "https://cognitiveservices.azure.com/.default"}, id="default_scope"), + pytest.param({"azure_username": ""}, id="empty_username"), + pytest.param({"client_secret": "environment-secret"}, id="environment_secret"), + ], +) +async def test_router_caller_entra_field_that_changes_nothing_still_cools_down_a_rejected_deployment( + caller_fields, monkeypatch +): + monkeypatch.delenv("AZURE_SCOPE", raising=False) + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + 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", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + }, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + **caller_fields, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + 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_caller_repeating_the_secret_an_empty_router_default_selects_still_cools_down(monkeypatch): + monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") + 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", + "tenant_id": "configured-tenant", + "client_id": "configured-client", + "client_secret": "configured-secret", + }, + "model_info": {"id": "configured-deployment"}, + } + ], + default_litellm_params={"client_secret": ""}, + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=_caused_by_an_entra_rejection( + litellm.AuthenticationError( + message="Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided.", + llm_provider="azure", + model="gpt-4.1", + ) + ), + client_secret="environment-secret", + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + 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_caller_entra_field_keeps_the_shared_deployment(): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "configured-key"}, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + response: Final = await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + client_id="caller-client", + ) + + assert response._hidden_params["model_id"] == "shared-deployment" + +@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") +@pytest.mark.parametrize( + ("litellm_params", "provider"), + [ + pytest.param({"model": "openai/gpt-4.1-mini", "api_key": "revoked-key"}, "openai", id="openai"), + pytest.param( + { + "model": "azure/gpt-4.1", + "api_base": "https://litellm-qa.openai.azure.com", + "api_version": "2025-04-01-preview", + "api_key": "revoked-key", + }, + "azure", + id="azure_api_key", + ), + ], +) +@pytest.mark.parametrize("caller_entra_field", ["client_id", "azure_scope"]) +async def test_router_auth_failure_entra_did_not_cause_cools_down_despite_a_caller_entra_field( + litellm_params, provider, caller_entra_field +): + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": litellm_params, + "model_info": {"id": "shared-deployment"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError( + message="Incorrect API key provided.", llm_provider=provider, model=litellm_params["model"] + ), + **{caller_entra_field: "caller-value"}, + ) + + with pytest.raises(litellm.RouterRateLimitError) as excinfo: + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + mock_response="hi", + ) + + 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") +@pytest.mark.parametrize( + "stream_request", + [ + pytest.param( + lambda router, **caller_fields: router.acompletion( + model="entra-model", messages=[{"role": "user", "content": "hi"}], stream=True, **caller_fields + ), + id="chat_completions", + ), + pytest.param( + lambda router, **caller_fields: router.aanthropic_messages( + model="entra-model", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + **caller_fields, + ), + id="messages", + ), + ], +) +async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deployment_routable( + monkeypatch: pytest.MonkeyPatch, + stream_request: Callable[..., Awaitable[object]], +) -> None: + 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: Final = 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"}, + } + ] + ) + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router, client_secret="not-the-credential") + + with pytest.raises(litellm.AuthenticationError): + await stream_request(router) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_async_openai_model_client(): router = Router(