From c9a5777739a38ef36281c161b69a98752cccdcc8 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:24:20 -0700 Subject: [PATCH] fix(azure): share one Entra token request per worker and audit #47348 (#48769) * test(azure): audit #47348 Add proxy integration tests for #47348's cached Entra token providers and off-loop token refresh, and fix the token herd the audit found: once the fetch moved to a worker thread, concurrent cold requests each fetched their own token. Every cached bearer token provider now serializes its calls so the first fetch fills the cache and the rest reuse it * fix(azure): build each cached Entra token provider once under concurrent first calls lru_cache does not hold back concurrent misses, so worker threads that miss together each built their own credential and token provider, and each fetched its own token. A lock around each provider cache makes the first caller build it and the rest reuse it * fix(azure): share Entra token fetches only while no valid token is held Concurrent callers now share one result (or one error) from the cold fetch, while callers inside the refresh window keep the still-valid token. The oidc/azure provider build and the client-assertion exchange are single-flighted per audience, and the shared helpers no longer call subscripted generic aliases, which crashed on Python 3.10. * fix(azure): refresh Entra tokens once while other requests keep the valid token MSAL treats a token with less than five minutes left as expired, so every request in that window fetched a new one. Now one request refreshes while the others keep using the still-valid token. Requests that waited on a fetch that failed retry once through a fresh fetch instead of all failing with the first error. Credentials that only implement get_token keep working, and OIDC token exchanges are cached per scope. * fix(azure): bound Entra token waits during an outage Requests waiting on a failed token fetch now retry only when the fetch before it succeeded, so during an Entra outage later bursts fail after one fetch instead of two. Requests during a refresh only get the held token while it has more than 60 seconds left, and a failed refresh falls back to that token instead of failing its request. * fix(azure): share one retry per failed Entra token fetch and raise on a failed refresh Waiters of a failed fetch now retry the flight they joined, so a fast-failing Entra makes one retry instead of one per waiter. A failed refresh raises again instead of returning the held token, which Postgres token auth read as a fresh one. The DefaultAzureCredential fallback reuses one provider per scope. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/single_flight.py | 92 +++ litellm/llms/azure/common_utils.py | 131 ++-- .../get_azure_ad_token_provider.py | 100 ++- litellm/secret_managers/main.py | 4 +- .../test_azure_entra_token_reuse_wire.py | 625 ++++++++++++++++++ tests/unit/llms/azure/test_azure.py | 8 +- .../llms/azure/test_azure_common_utils.py | 140 +++- .../test_get_azure_ad_token_provider.py | 270 +++++++- .../test_secret_managers_main.py | 40 ++ 9 files changed, 1335 insertions(+), 75 deletions(-) create mode 100644 litellm/litellm_core_utils/single_flight.py create mode 100644 tests/integration/providers/test_azure_entra_token_reuse_wire.py diff --git a/litellm/litellm_core_utils/single_flight.py b/litellm/litellm_core_utils/single_flight.py new file mode 100644 index 00000000000..44a84d05925 --- /dev/null +++ b/litellm/litellm_core_utils/single_flight.py @@ -0,0 +1,92 @@ +import threading +from collections.abc import Callable +from concurrent.futures import Future +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Final, Generic, ParamSpec, TypeVar + +_P: Final = ParamSpec("_P") +_R: Final = TypeVar("_R") + + +@dataclass(slots=True) +class SingleFlight(Generic[_R]): + lock: threading.Lock = field(default_factory=threading.Lock) + in_flight: Future[_R] | None = None + flights_started: int = 0 + last_flight_failed: bool = False + + def run(self, fetch: Callable[[], _R], while_busy: _R | None = None) -> _R: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + joined_flight: Final = self.flights_started + retry_if_it_fails: Final = not self.last_flight_failed + if in_flight is None: + self._take_off(leader) + if in_flight is None: + return self._lead(leader, fetch) + if while_busy is not None: + return while_busy + try: + return in_flight.result() + except Exception: + retry: Final = self._claim_retry(joined_flight) if retry_if_it_fails else None + if retry is None: + raise + retry_flight, leads_retry = retry + return self._lead(retry_flight, fetch) if leads_retry else retry_flight.result() + + def _take_off(self, leader: Future[_R]) -> None: + self.in_flight = leader + self.flights_started += 1 + + def _claim_retry(self, failed_flight: int) -> tuple[Future[_R], bool] | None: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + if in_flight is not None: + return in_flight, False + if self.flights_started > failed_flight and self.last_flight_failed: + return None + self._take_off(leader) + return leader, True + + def _lead(self, flight: Future[_R], fetch: Callable[[], _R]) -> _R: + try: + result: Final = fetch() + except BaseException as error: + self._land(failed=True) + flight.set_exception(error) + raise + self._land(failed=False) + flight.set_result(result) + return result + + def _land(self, failed: bool) -> None: + with self.lock: + self.in_flight = None + self.last_flight_failed = failed + + +@dataclass(frozen=True, slots=True) +class SingleFlightCache(Generic[_P, _R]): + cached: Callable[..., _R] + clear: Callable[[], None] + lock: threading.Lock = field(default_factory=threading.Lock) + + def __call__( + self, + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwarded verbatim to the cached builder, typed by its own signature + ) -> _R: + with self.lock: + return self.cached(*args, **kwargs) + + def cache_clear(self) -> None: + self.clear() + + +def single_flight_lru_cache(build: Callable[_P, _R]) -> SingleFlightCache[_P, _R]: + cached: Final = lru_cache(maxsize=128)(build) + return SingleFlightCache(cached, cached.cache_clear) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 23c5aa35f4b..4bd19889bcc 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -5,7 +5,6 @@ import os import re from collections.abc import Awaitable, Callable, Mapping from contextlib import closing -from functools import lru_cache from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast @@ -18,15 +17,18 @@ import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache 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.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, + single_flight_bearer_token_provider, ) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType from litellm.utils import add_path_to_api_base azure_ad_cache: Final = DualCache() @@ -94,7 +96,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: return {**llm_response_headers, **openai_headers} -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_entra_id_token_provider( tenant_id: str, client_id: str, @@ -107,33 +109,38 @@ def _cached_entra_id_token_provider( to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one per request. """ - from azure.identity import ClientSecretCredential, get_bearer_token_provider + from azure.identity import ClientSecretCredential - return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) + return single_flight_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_username_password_token_provider( client_id: str, azure_username: str, azure_password: str, scope: str, ) -> Callable[[], str]: - from azure.identity import UsernamePasswordCredential, get_bearer_token_provider + from azure.identity import UsernamePasswordCredential credential: Final = UsernamePasswordCredential( client_id=client_id, username=azure_username, password=azure_password, ) - return get_bearer_token_provider(credential, scope) + return single_flight_bearer_token_provider(credential, scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=scope) +@single_flight_lru_cache +def _cached_default_azure_credential_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope, azure_credential=AzureCredentialType.DefaultAzureCredential) + + def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return _cached_azure_ad_token_refresh_provider(scope) @@ -240,6 +247,58 @@ def get_azure_ad_token_from_username_password( return token_provider +@single_flight_lru_cache +def _oidc_token_exchange(cache_key: str) -> SingleFlight[str]: + return SingleFlight() + + +def _cached_oidc_token(cache_key: str) -> str | None: + cached: Final = azure_ad_cache.get_cache(cache_key) + return cached if isinstance(cached, str) else None + + +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, + }, + ) + + if req_token.status_code != 200: + raise AzureOpenAIError( + status_code=req_token.status_code, + message=req_token.text, + ) + + azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() + azure_ad_token_access_token: Final = azure_ad_token_json.get("access_token", None) + azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) + + if azure_ad_token_access_token is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") + + if azure_ad_token_expires_in is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") + + azure_ad_cache.set_cache( + key=cache_key, + value=azure_ad_token_access_token, + ttl=azure_ad_token_expires_in, + ) + + return azure_ad_token_access_token + + def get_azure_ad_token_from_oidc( azure_ad_token: str, azure_client_id: str | None = None, @@ -282,50 +341,25 @@ def get_azure_ad_token_from_oidc( "azure_client_id": azure_client_id, "azure_tenant_id": azure_tenant_id, "azure_authority_host": azure_authority_host, + "scope": scope, "oidc_token": oidc_token, } ) - azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key) - if azure_ad_token_access_token is not None: - return azure_ad_token_access_token + cached_access_token: Final = _cached_oidc_token(azure_ad_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( - f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", - data={ - "client_id": azure_client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, + return _oidc_token_exchange(azure_ad_token_cache_key).run( + lambda: _exchange_oidc_token( + azure_ad_token_cache_key, + token_url=f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", + client_id=azure_client_id, + scope=scope, + oidc_token=oidc_token, ) - - if req_token.status_code != 200: - raise AzureOpenAIError( - status_code=req_token.status_code, - message=req_token.text, - ) - - azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() - azure_ad_token_access_token = azure_ad_token_json.get("access_token", None) - azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) - - if azure_ad_token_access_token is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") - - if azure_ad_token_expires_in is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") - - azure_ad_cache.set_cache( - key=azure_ad_token_cache_key, - value=azure_ad_token_access_token, - ttl=azure_ad_token_expires_in, ) - return azure_ad_token_access_token - def select_azure_base_url_or_endpoint(azure_client_params: dict): azure_endpoint: Final = azure_client_params.get("azure_endpoint", None) @@ -503,17 +537,10 @@ class BaseAzureLLM(BaseOpenAILLM): Returns: Token provider callable if DefaultAzureCredential is enabled and available, None otherwise """ - from litellm.types.secret_managers.get_azure_ad_token_provider import ( - AzureCredentialType, - ) - verbose_logger.debug("Attempting to use DefaultAzureCredential for Azure Auth") try: - azure_ad_token_provider: Final = get_azure_ad_token_provider( - azure_scope=scope, - azure_credential=AzureCredentialType.DefaultAzureCredential, - ) + azure_ad_token_provider: Final = _cached_default_azure_credential_provider(scope) verbose_logger.debug("Successfully obtained Azure AD token provider using DefaultAzureCredential") return azure_ad_token_provider except Exception as e: diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 2b707f179d0..4967e199218 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,14 +1,27 @@ import os +import time from collections.abc import Callable +from dataclasses import dataclass +from types import TracebackType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.single_flight import SingleFlight from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, ) +_HELD_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 + if TYPE_CHECKING: - from azure.core.credentials import TokenProvider + from azure.core.credentials import ( + AccessToken, + AccessTokenInfo, + SupportsTokenInfo, + TokenCredential, + TokenRequestOptions, + ) + from typing_extensions import TypeIs # noqa: TID251 # tells the two azure-core credential protocols apart def infer_credential_type_from_environment() -> AzureCredentialType: @@ -37,6 +50,86 @@ def infer_credential_type_from_environment() -> AzureCredentialType: return AzureCredentialType.DefaultAzureCredential +@dataclass(slots=True) +class _HeldToken: + token: "AccessToken | AccessTokenInfo | None" = None + + def still_valid(self) -> str | None: + held: Final = self.token + if held is None or held.expires_on - time.time() <= _HELD_TOKEN_EXPIRY_MARGIN_SECONDS: + return None + return held.token + + +@dataclass(slots=True) +class _TokenInfoWatcher: + credential: "SupportsTokenInfo" + held: _HeldToken + + def get_token_info(self, *scopes: str, options: "TokenRequestOptions | None" = None) -> "AccessTokenInfo": + token: Final = self.credential.get_token_info(*scopes, options=options) + self.held.token = token + return token + + def close(self) -> None: + self.credential.close() + + def __enter__(self) -> "SupportsTokenInfo": + return self.credential.__enter__() + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.credential.__exit__(exc_type, exc_value, traceback) + + +@dataclass(slots=True) +class _TokenWatcher: + credential: "TokenCredential" + held: _HeldToken + + def get_token( + self, + *scopes: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, # kwargs-ok: azure-core forwards its request options verbatim to the wrapped credential + ) -> "AccessToken": + token: Final = self.credential.get_token( + *scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs + ) + self.held.token = token + return token + + +def _supports_token_info( + credential: "TokenCredential | SupportsTokenInfo", +) -> "TypeIs[SupportsTokenInfo]": + return hasattr(credential, "get_token_info") + + +def single_flight_bearer_token_provider( + credential: "TokenCredential | SupportsTokenInfo", scope: str +) -> Callable[[], str]: + from azure.identity import get_bearer_token_provider + + held: Final = _HeldToken() + watched: Final = ( + _TokenInfoWatcher(credential, held) if _supports_token_info(credential) else _TokenWatcher(credential, held) + ) + provider: Final = get_bearer_token_provider(watched, scope) + token_request: Final[SingleFlight[str]] = SingleFlight() + + def bearer_token() -> str: + return token_request.run(provider, while_busy=held.still_valid()) + + return bearer_token + + def get_azure_ad_token_provider( azure_scope: str | None = None, azure_credential: AzureCredentialType | None = None, @@ -65,7 +158,6 @@ def get_azure_ad_token_provider( DefaultAzureCredential, ManagedIdentityCredential, WorkloadIdentityCredential, - get_bearer_token_provider, ) if azure_scope is None: @@ -77,7 +169,7 @@ def get_azure_ad_token_provider( else None or os.environ.get("AZURE_CREDENTIAL") or infer_credential_type_from_environment() ) verbose_logger.info("For Azure AD Token Provider, choosing credential type: %s", cred) - credential: TokenProvider | None = None + credential: TokenCredential | SupportsTokenInfo | None = None if cred == AzureCredentialType.ClientSecretCredential: credential = ClientSecretCredential( client_id=os.environ["AZURE_CLIENT_ID"], @@ -133,4 +225,4 @@ def get_azure_ad_token_provider( if credential is None: raise ValueError("No credential provided") - return get_bearer_token_provider(credential, azure_scope) + return single_flight_bearer_token_provider(credential, azure_scope) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index b33cf0e6b7e..588bdb1409c 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -5,7 +5,6 @@ import time import traceback from collections.abc import Callable from dataclasses import dataclass -from functools import lru_cache from typing import Final import httpx @@ -14,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.single_flight import single_flight_lru_cache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -27,7 +27,7 @@ _PARSED_LITERAL: Final = TypeAdapter(object) _SECRET_MANAGER_CLIENT: Final = TypeAdapter(object) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_oidc_token_provider(audience: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=audience) diff --git a/tests/integration/providers/test_azure_entra_token_reuse_wire.py b/tests/integration/providers/test_azure_entra_token_reuse_wire.py new file mode 100644 index 00000000000..85aef110124 --- /dev/null +++ b/tests/integration/providers/test_azure_entra_token_reuse_wire.py @@ -0,0 +1,625 @@ +import asyncio +import itertools +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path +from queue import SimpleQueue +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import httpx +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.ports import allocate_port +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import Peer, completed, responses_peer, session +from integration._support.tls import write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_DEPLOYMENT: Final = "gpt-5.4-nano" +_DATED_API_VERSION: Final = "2024-10-21" +_REALTIME_API_VERSION: Final = "2024-10-01-preview" +_IDENTITY_PATH: Final = "/msi/token" +_EXCHANGE_SUFFIX: Final = "/oauth2/v2.0/token" +_FOUNDRY_PROJECT: Final = "/api/projects/audit" +_BING_TOOL: Final = "audit-bing-entra" +_BING_RESOURCE: Final = "https://ai.azure.com" +_BING_CITATION: Final = "https://result.invalid/bing" +_TENANT: Final = "audit-tenant" +_APP_CLIENT: Final = "audit-app-client" +_HOLD_SECONDS: Final = 60 +_LIVELINESS_SECONDS: Final = 3 +_HERD: Final = 6 +_HERD_WINDOW_SECONDS: Final = 3 +_TOKEN_LIFETIME_SECONDS: Final = 3600 +_INSIDE_REFRESH_WINDOW_SECONDS: Final = 120 +_REMOVED_ENVIRONMENT: Final = ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_FEDERATED_TOKEN_FILE", + "AZURE_USERNAME", + "AZURE_PASSWORD", + "AZURE_CREDENTIAL", + "AZURE_SCOPE", + "MSI_ENDPOINT", + "MSI_SECRET", + "IMDS_ENDPOINT", + "IDENTITY_SERVER_THUMBPRINT", + "BING_GROUNDING_TOKEN", + "HTTP_PROXY", + "HTTPS_PROXY", + "http_proxy", + "https_proxy", +) + + +class _Surface(StrEnum): + LEGACY_CHAT = "legacy-chat" + V1_CHAT = "v1-chat" + RESPONSES = "responses" + RESPONSES_WEBSOCKET = "responses-websocket" + REALTIME = "realtime" + REALTIME_HEALTH_CHECK = "realtime-health-check" + + +_HTTP_SURFACES: Final = frozenset({_Surface.LEGACY_CHAT, _Surface.V1_CHAT, _Surface.RESPONSES}) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + port: int + cert: tuple[Path, Path] + + @property + def edge(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +@dataclass(frozen=True, slots=True) +class _Hold: + resource: str + arrivals: SimpleQueue[str] = field(default_factory=SimpleQueue) + release: threading.Event = field(default_factory=threading.Event) + + +def _config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "azure-entra-token-reuse.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "litellm_settings": {**object_value(stock["litellm_settings"]), "enable_azure_ad_token_refresh": True}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "search_tools": [ + {"search_tool_name": _BING_TOOL, "litellm_params": {"search_provider": "bing_grounding"}} + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("azure-entra-token-reuse") + port: Final = allocate_port() + cert: Final = write_self_signed_cert(directory) + edge: Final = f"http://127.0.0.1:{port}" + environment: Final = { + "AZURE_CLIENT_ID": "audit-managed-identity", + "IDENTITY_ENDPOINT": f"{edge}{_IDENTITY_PATH}", + "IDENTITY_HEADER": "audit-identity-header", + "AZURE_AUTHORITY_HOST": edge, + "BING_GROUNDING_PROJECT_ENDPOINT": f"{edge}{_FOUNDRY_PROJECT}", + "BING_GROUNDING_MODEL": _DEPLOYMENT, + "BING_GROUNDING_CONNECTION_ID": "audit-bing-connection", + "SSL_CERT_FILE": str(cert[0]), + } + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, directory, environment, config=_config(directory), remove_environment=_REMOVED_ENVIRONMENT + ) as proxy, + ): + yield _Rig(proxy, port, cert) + + +def _json(payload: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def _identity_token(resource: str) -> str: + return f"mi-token:{resource}" + + +def _identity_resource(request: Request) -> str | None: + target: Final = urlsplit(request.target) + if target.path != _IDENTITY_PATH: + return None + return parse_qs(target.query)["resource"][-1] + + +def _is_token_endpoint(request: Request) -> bool: + return _identity_resource(request) is not None or request.target.endswith(_EXCHANGE_SUFFIX) + + +def _managed_identity( + resource: str, + hold: _Hold | None, + lifetime_seconds: int = _TOKEN_LIFETIME_SECONDS, + access_token: str | None = None, +) -> Reply: + if hold is not None and resource == hold.resource: + hold.arrivals.put(resource) + hold.release.wait(_HOLD_SECONDS) + return _json( + { + "access_token": access_token or _identity_token(resource), + "expires_on": str(int(time.time()) + lifetime_seconds), + "resource": resource, + "token_type": "Bearer", + } + ) + + +def _entra_exchange(request: Request) -> Reply: + assertion: Final = parse_qs(request.body.decode())["client_assertion"][-1] + return _json({"access_token": f"aad:{assertion}", "expires_in": 3600, "token_type": "Bearer"}) + + +def _edge(upstream: Callable[[Request], Reply], hold: _Hold | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + resource: Final = _identity_resource(request) + if resource is not None: + return _managed_identity(resource, hold) + if request.target.endswith(_EXCHANGE_SUFFIX): + return _entra_exchange(request) + return upstream(request) + + return respond + + +def _azure_openai(request: Request) -> Reply: + if urlsplit(request.target).path.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + return responses_reply(f"resp_{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + + +def _identity_fetches(requests: tuple[Request, ...], resource: str) -> int: + return sum(1 for request in requests if _identity_resource(request) == resource) + + +def _provider_calls(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if not _is_token_endpoint(request)) + + +def _marker() -> str: + return f"entra audit {uuid.uuid4().hex}" + + +def _model(scenario: Scenario, surface: _Surface, api_base: str, scope: str) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.RESPONSES: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_DATED_API_VERSION, + azure_scope=scope, + ) + case _Surface.V1_CHAT: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, api_version="v1", azure_scope=scope + ) + case _Surface.RESPONSES_WEBSOCKET: + return scenario.model(model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, azure_scope=scope) + case _Surface.REALTIME | _Surface.REALTIME_HEALTH_CHECK: + return scenario.model( + model_info={"mode": "realtime"}, + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_REALTIME_API_VERSION, + azure_scope=scope, + ) + + +def _post(surface: _Surface, gateway: Gateway, key: str, model: str) -> httpx.Response: + if surface is _Surface.RESPONSES: + return gateway.request("POST", "/v1/responses", {"model": model, "input": _marker()}, key=key) + return gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": _marker()}]}, key=key + ) + + +async def _realtime_handshake(gateway: Gateway, key: str, model: str, peer: Peer) -> str | None: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=_HOLD_SECONDS, + ): + credentials: Final = await asyncio.to_thread(peer.credentials.get, True, _HOLD_SECONDS) + return credentials.authorization + + +def _upstream_authorization(surface: _Surface, gateway: Gateway, key: str, model: str, wire: Wire, peer: Peer) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.V1_CHAT | _Surface.RESPONSES: + response: Final = _post(surface, gateway, key, model) + assert response.status_code == 200, response.text + (call,) = _provider_calls(wire.drain()) + return call.headers.get("authorization", "") + case _Surface.RESPONSES_WEBSOCKET: + (events,) = asyncio.run(session(str(gateway.client.base_url), key, model, (_marker(),))) + assert completed(events)["status"] == "completed", events + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + case _Surface.REALTIME: + return asyncio.run(_realtime_handshake(gateway, key, model, peer)) or "" + case _Surface.REALTIME_HEALTH_CHECK: + gateway.request("GET", "/health", params={"model": model}) + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + + +def _liveliness(gateway: Gateway) -> int | None: + try: + return httpx.get( + f"{str(gateway.client.base_url).rstrip('/')}/health/liveliness", timeout=_LIVELINESS_SECONDS + ).status_code + except httpx.TimeoutException: + return None + + +def _liveliness_while_held(gateway: Gateway, hold: _Hold) -> int | None: + try: + assert hold.arrivals.get(timeout=_HOLD_SECONDS) == hold.resource + return _liveliness(gateway) + finally: + hold.release.set() + + +@pytest.mark.parametrize("surface", tuple(_Surface)) +def test_a_slow_managed_identity_fetch_leaves_the_proxy_answering_other_requests(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-hold-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + responses_peer(rig.cert) as peer, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=1) as pool, + ): + api_base: Final = rig.edge if surface in _HTTP_SURFACES else peer.url.removesuffix("/v1") + model: Final = _model(scenario, surface, api_base, f"{resource}/.default") + key: Final = scenario.key() + call: Final = pool.submit(_upstream_authorization, surface, rig.gateway, key, model, wire, peer) + liveliness: Final = _liveliness_while_held(rig.gateway, hold) + authorization: Final = call.result(timeout=_HOLD_SECONDS) + + assert liveliness == 200, f"/health/liveliness did not answer within {_LIVELINESS_SECONDS}s during the fetch" + assert authorization == f"Bearer {_identity_token(resource)}" + + +@pytest.mark.parametrize("surface", (_Surface.LEGACY_CHAT, _Surface.RESPONSES)) +def test_concurrent_cold_requests_share_one_managed_identity_fetch(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, surface, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, surface, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, resource) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(resource)}" + ] * _HERD + + +def _scripted_identity(resource: str, fetches: tuple[Callable[[], Reply], ...]) -> Callable[[Request], Reply]: + attempts: Final = itertools.count() + edge: Final = _edge(_azure_openai) + + def respond(request: Request) -> Reply: + if _identity_resource(request) != resource: + return edge(request) + return fetches[min(next(attempts), len(fetches) - 1)]() + + return respond + + +def _held_failure(hold: _Hold) -> Reply: + hold.arrivals.put(hold.resource) + hold.release.wait(_HOLD_SECONDS) + return Reply(body=b"not json") + + +def _herd( + rig: _Rig, edge: Callable[[Request], Reply], resource: str, hold: _Hold, warm_up: bool +) -> tuple[tuple[httpx.Response, ...], tuple[Request, ...]]: + with ( + wire_server(edge, port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, _Surface.LEGACY_CHAT, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + warm_up_responses: Final = (_post(_Surface.LEGACY_CHAT, rig.gateway, key, model),) if warm_up else () + calls: Final = tuple(pool.submit(_post, _Surface.LEGACY_CHAT, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + return warm_up_responses + tuple(call.result(timeout=_HOLD_SECONDS) for call in calls), wire.drain() + + +def test_concurrent_requests_inside_the_refresh_window_share_one_refresh_and_keep_the_valid_token(rig: _Rig) -> None: + resource: Final = f"api://audit-refresh-{uuid.uuid4().hex}" + refreshed: Final = f"mi-token-refreshed:{resource}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity( + resource, + ( + lambda: _managed_identity(resource, None, _INSIDE_REFRESH_WINDOW_SECONDS), + lambda: _managed_identity(resource, hold, access_token=refreshed), + ), + ) + + responses, requests = _herd(rig, edge, resource, hold, warm_up=True) + + assert [response.status_code for response in responses] == [200] * (_HERD + 1), [r.text for r in responses] + assert _identity_fetches(requests, resource) == 2 + assert sorted(call.headers.get("authorization", "") for call in _provider_calls(requests)) == sorted( + [f"Bearer {_identity_token(resource)}"] * _HERD + [f"Bearer {refreshed}"] + ) + + +def test_a_failed_cold_fetch_fails_only_the_request_that_made_it(rig: _Rig) -> None: + resource: Final = f"api://audit-failure-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity(resource, (lambda: _held_failure(hold), lambda: _managed_identity(resource, None))) + + responses, _ = _herd(rig, edge, resource, hold, warm_up=False) + + assert sorted(response.status_code == 200 for response in responses) == [False] + [True] * (_HERD - 1), [ + r.text for r in responses + ] + + +def _oidc_model(scenario: Scenario, rig: _Rig, audience: str) -> str: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=rig.edge, + api_key=None, + api_version=_DATED_API_VERSION, + azure_ad_token=f"oidc/azure/{audience}", + tenant_id=_TENANT, + client_id=_APP_CLIENT, + ) + + +def _exchanges(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if request.target.endswith(_EXCHANGE_SUFFIX)) + + +def test_an_oidc_federated_deployment_fetches_its_assertion_once_across_responses_calls(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-{uuid.uuid4().hex}" + with wire_server(_edge(_azure_openai), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + responses: Final = tuple(_post(_Surface.RESPONSES, rig.gateway, key, model) for _ in range(3)) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + exchanges: Final = _exchanges(requests) + assert [request.target for request in exchanges] == [f"/{_TENANT}{_EXCHANGE_SUFFIX}"] + assert {name: values[-1] for name, values in parse_qs(exchanges[0].body.decode()).items()} == { + "client_id": _APP_CLIENT, + "grant_type": "client_credentials", + "scope": "https://cognitiveservices.azure.com/.default", + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": _identity_token(audience), + } + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * 3 + + +def test_concurrent_cold_oidc_requests_share_one_assertion_fetch_and_one_exchange(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(audience) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, _Surface.RESPONSES, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + assert len(_exchanges(requests)) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * _HERD + + +def _foundry(request: Request) -> Reply: + assert urlsplit(request.target).path == f"{_FOUNDRY_PROJECT}/openai/v1/responses", request.target + text: Final = "Audit result." + return _json( + { + "status": "completed", + "error": None, + "incomplete_details": None, + "output": [ + { + "type": "message", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [ + { + "type": "url_citation", + "url": _BING_CITATION, + "title": "Bing result", + "start_index": len(text), + "end_index": len(text), + } + ], + } + ], + } + ], + } + ) + + +class _SearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + url: str + + +class _SearchResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + results: tuple[_SearchResult, ...] + + +def test_bing_grounding_search_mints_its_entra_token_once_across_searches(rig: _Rig) -> None: + with wire_server(_edge(_foundry), port=rig.port) as wire: + responses: Final = tuple( + rig.gateway.request("POST", f"/v1/search/{_BING_TOOL}", {"query": _marker(), "max_results": 1}) + for _ in range(3) + ) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert [ + tuple(result.url for result in _SearchResponse.model_validate_json(response.content).results) + for response in responses + ] == [(_BING_CITATION,)] * 3 + assert _identity_fetches(requests, _BING_RESOURCE) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(_BING_RESOURCE)}" + ] * 3 + + +def _lifecycle(upstream_id: str) -> Callable[[Request], Reply]: + response: Final[dict[str, JsonValue]] = { + "id": upstream_id, + "object": "response", + "created_at": 1, + "model": _DEPLOYMENT, + "output": [], + } + + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + match request.method: + case "DELETE": + return _json({"id": upstream_id, "object": "response", "deleted": True}) + case "GET" if path.endswith("/input_items"): + return _json({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}) + case "POST" if path.endswith("/cancel"): + return _json({**response, "status": "cancelled"}) + case _: + return _json({**response, "status": "completed"}) + + return respond + + +_LIFECYCLE_ROUTES: Final = ( + ("GET", "/v1/responses/{}", ""), + ("GET", "/v1/responses/{}/input_items", "/input_items"), + ("POST", "/v1/responses/{}/cancel", "/cancel"), + ("DELETE", "/v1/responses/{}", ""), +) + + +def test_every_responses_route_on_an_entra_deployment_reaches_azure_with_the_managed_identity_token( + rig: _Rig, +) -> None: + resource: Final = f"api://audit-routes-{uuid.uuid4().hex}" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(_edge(_lifecycle(upstream_id)), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _model(scenario, _Surface.RESPONSES, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + created: Final = _post(_Surface.RESPONSES, rig.gateway, key, model) + assert created.status_code == 200, created.text + response_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"]) + followups: Final = tuple( + rig.gateway.request(method, route.format(response_id), key=key) for method, route, _ in _LIFECYCLE_ROUTES + ) + requests: Final = wire.drain() + + assert [response.status_code for response in followups] == [200] * len(_LIFECYCLE_ROUTES), [ + response.text for response in followups + ] + calls: Final = _provider_calls(requests) + assert (calls[0].method, urlsplit(calls[0].target).path.endswith("/responses")) == ("POST", True) + assert [ + (call.method, urlsplit(call.target).path.endswith(f"/{upstream_id}{suffix}")) + for call, (_, _, suffix) in zip(calls[1:], _LIFECYCLE_ROUTES, strict=True) + ] == [(method, True) for method, _, _ in _LIFECYCLE_ROUTES] + assert {call.headers.get("authorization") for call in calls} == {f"Bearer {_identity_token(resource)}"} + assert _identity_fetches(requests, resource) == 1 diff --git a/tests/unit/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7d5afde68ef..c52321baf02 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -6,7 +6,7 @@ import os import time from collections.abc import Mapping from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -291,11 +291,9 @@ def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provide mock_credential.assert_called_once_with(client_id=client_id, username=username, password=password) - mock_get_bearer_token_provider.assert_called_once_with( - mock_credential.return_value, "https://cognitiveservices.azure.com/.default" - ) + mock_get_bearer_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") - assert result == mock_token_provider + assert result() == "mock-token" def test_azure_openai_gpt_4o_naming(monkeypatch): diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 1dc7ff74323..f7e2e73d426 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -2,9 +2,12 @@ import json import os import threading import traceback +import uuid +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Callable, Final, Optional from urllib.parse import parse_qs -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -13,6 +16,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, _cached_azure_ad_token_refresh_provider, + _cached_default_azure_credential_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -41,6 +45,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() with ( patch( @@ -86,6 +91,7 @@ def setup_mocks(monkeypatch): "select_url": mock_select_url, } _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): @@ -1436,6 +1442,33 @@ def test_get_azure_ad_token_with_token_refresh_reuses_provider_per_scope(setup_m ] +def test_get_azure_ad_token_reuses_one_default_azure_credential_provider_per_scope(setup_mocks): + setup_mocks["litellm"].enable_azure_ad_token_refresh = True + default_scope: Final = "https://cognitiveservices.azure.com/.default" + other_scope: Final = "api://other-resource/.default" + + def provider_for(azure_scope: str, azure_credential: AzureCredentialType | None = None) -> Callable[[], str]: + if azure_credential != AzureCredentialType.DefaultAzureCredential: + raise ValueError("No credential provided") + return lambda: f"token-for-{azure_scope}" + + setup_mocks["token_provider"].side_effect = provider_for + + tokens: Final = ( + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams(azure_scope=other_scope)), + ) + + default_credential_builds: Final = tuple( + call.kwargs["azure_scope"] + for call in setup_mocks["token_provider"].call_args_list + if call.kwargs.get("azure_credential") == AzureCredentialType.DefaultAzureCredential + ) + assert tokens == (f"token-for-{default_scope}", f"token-for-{default_scope}", f"token-for-{other_scope}") + assert default_credential_builds == (default_scope, other_scope) + + def test_get_azure_ad_token_with_token_refresh_error(setup_mocks): """Test get_azure_ad_token with token refresh enabled but raising an error.""" # Reset mocks to ensure clean state @@ -1536,9 +1569,7 @@ def test_get_azure_ad_token_provider_with_default_azure_credential(): mock_default_cred.assert_called_once_with() # Verify get_bearer_token_provider was called with the right parameters - mock_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" - ) + mock_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") # Verify the returned token provider works token = token_provider() @@ -2137,6 +2168,42 @@ class TestEntraIdTokenProviderCache: assert first is second assert mock_credential.call_count == 1 + def test_concurrent_first_callers_build_one_credential(self): + callers: Final = 8 + start: Final = threading.Barrier(callers) + built: Final = SimpleQueue[None]() + everyone_built: Final = threading.Event() + release: Final = threading.Event() + + def held_credential(*_: str) -> MagicMock: + built.put(None) + if built.qsize() >= callers: + everyone_built.set() + release.wait(5) + return MagicMock() + + def first_call() -> Callable[[], str]: + start.wait(5) + return get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + + with ( + patch("azure.identity.ClientSecretCredential", side_effect=held_credential), + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(first_call) for _ in range(callers)) + everyone_built.wait(0.5) + release.set() + providers: Final = tuple(call.result(timeout=5) for call in calls) + + assert built.qsize() == 1 + assert all(provider is providers[0] for provider in providers) + @pytest.mark.parametrize( "second_call_kwargs", [ @@ -2291,6 +2358,71 @@ def test_oidc_token_exchange_closes_its_own_http_client(monkeypatch: pytest.Monk assert token_endpoint.closed +def test_concurrent_oidc_callers_share_one_token_exchange(monkeypatch: pytest.MonkeyPatch) -> None: + callers: Final = 8 + start: Final = threading.Barrier(callers) + exchanges: Final = SimpleQueue[None]() + everyone_exchanging: Final = threading.Event() + release: Final = threading.Event() + + class HeldTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + exchanges.put(None) + if exchanges.qsize() >= callers: + everyone_exchanging.set() + release.wait(5) + return httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", HeldTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_BURST_ASSERTION", "burst-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def exchange() -> str: + start.wait(5) + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_BURST_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + ) + + with ThreadPoolExecutor(max_workers=callers) as pool: + calls: Final = tuple(pool.submit(exchange) for _ in range(callers)) + exchanges.get(timeout=5) + everyone_exchanging.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + + assert tokens == ("exchanged-token",) * callers + assert exchanges.qsize() == 0 + + +def test_oidc_tokens_for_different_scopes_are_exchanged_and_cached_separately(monkeypatch: pytest.MonkeyPatch) -> None: + exchanged_scopes: Final = SimpleQueue[str]() + + class ScopedTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + scope: Final = parse_qs(request.read().decode())["scope"][-1] + exchanged_scopes.put(scope) + return httpx.Response(200, json={"access_token": f"token-for:{scope}", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", ScopedTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_SCOPED_ASSERTION", "scoped-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def token_for(scope: str) -> str: + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_SCOPED_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + scope=scope, + ) + + tokens: Final = tuple(token_for(scope) for scope in ("api://first/.default", "api://second/.default") * 2) + + assert tokens == ("token-for:api://first/.default", "token-for:api://second/.default") * 2 + assert exchanged_scopes.qsize() == 2 + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. diff --git a/tests/unit/secret_managers/test_get_azure_ad_token_provider.py b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py index 4bc7c21d8d2..141072635bc 100644 --- a/tests/unit/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py @@ -1,16 +1,25 @@ +import itertools import json import os -from typing import Optional -from unittest.mock import MagicMock, patch +import threading +import time +from collections.abc import Callable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor, wait +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Optional +from unittest.mock import ANY, MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules import pytest +from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions from azure.core.exceptions import ClientAuthenticationError from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, infer_credential_type_from_environment, + single_flight_bearer_token_provider, ) from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, @@ -190,7 +199,7 @@ class TestGetAzureAdTokenProvider: tenant_id="test-tenant-id", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -228,7 +237,7 @@ class TestGetAzureAdTokenProvider: client_id="test-client-id" ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -270,7 +279,7 @@ class TestGetAzureAdTokenProvider: certificate_path="/path/to/cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -314,7 +323,7 @@ class TestGetAzureAdTokenProvider: password="pwd4cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -348,7 +357,7 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_certificate_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -417,9 +426,254 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_default_azure_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works token = result() assert token == "mock-default-token" + + +_CONCURRENT_CALLERS: Final = 4 +_SCOPE: Final = "api://audit/.default" +_YEAR_2100: Final = 4_102_444_800 +_LONG_EXPIRED: Final = 1 +_JOIN_WINDOW_SECONDS: Final = 0.5 +_FAST_FAILING_CALLERS: Final = 16 + + +@dataclass(frozen=True, slots=True) +class _HeldCredential: + expires_on: int = _YEAR_2100 + failing_attempts: int = 0 + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[tuple[str, ...]] = field(default_factory=SimpleQueue) + everyone_asked: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def get_token_info(self, *scopes: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(scopes) + if self.requests.qsize() >= _CONCURRENT_CALLERS: + self.everyone_asked.set() + self.release.wait(5) + attempt: Final = next(self.attempts) + if attempt < self.failing_attempts: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", self.expires_on) + + +def _call_together(provider: Callable[[], str], credential: _HeldCredential) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + assert credential.requests.get(timeout=5) == (_SCOPE,) + credential.everyone_asked.wait(0.5) + credential.release.set() + return calls + + +def test_concurrent_cold_callers_share_one_token_request() -> None: + credential: Final = _HeldCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-0",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_concurrent_callers_after_the_token_expired_share_one_token_request() -> None: + credential: Final = _HeldCredential(expires_on=_LONG_EXPIRED) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.release.set() + provider() + credential.requests.get_nowait() + credential.release.clear() + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-1",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_a_failed_cold_token_request_fails_only_the_caller_that_made_it() -> None: + credential: Final = _HeldCredential(failing_attempts=1) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + failed: Final = sorted(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert failed == [False] * (_CONCURRENT_CALLERS - 1) + [True] + + +@dataclass(frozen=True, slots=True) +class _ScriptedCredential: + expiries: Iterator[int | None] = field(default_factory=lambda: itertools.repeat(None)) + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + answers: threading.Semaphore = field(default_factory=lambda: threading.Semaphore(0)) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(None) + self.answers.acquire(timeout=5) + attempt: Final = next(self.attempts) + expires_on: Final = next(self.expiries) + if expires_on is None: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", expires_on) + + +def _answer_together( + provider: Callable[[], str], credential: _ScriptedCredential, token_requests: int +) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + for _ in range(token_requests): + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release() + return calls + + +def test_an_entra_outage_retries_only_the_first_failed_token_request() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + first_burst: Final = _answer_together(provider, credential, token_requests=2) + later_burst: Final = _answer_together(provider, credential, token_requests=1) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in first_burst + later_burst) + assert credential.requests.qsize() == 0 + + +def test_waiting_callers_retry_a_failed_token_request_again_once_entra_recovers() -> None: + credential: Final = _ScriptedCredential(iter((None, None, _LONG_EXPIRED, None, _LONG_EXPIRED))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + _answer_together(provider, credential, token_requests=2) + credential.answers.release() + assert provider() == "token-2" + credential.requests.get_nowait() + + recovered: Final = _answer_together(provider, credential, token_requests=2) + + assert sorted(call.exception() is None for call in recovered) == [False] + [True] * (_CONCURRENT_CALLERS - 1) + assert {call.result() for call in recovered if call.exception() is None} == {"token-4"} + + +def test_a_failed_refresh_fails_the_request_that_made_it_even_while_a_valid_token_is_held() -> None: + credential: Final = _ScriptedCredential(iter((int(time.time()) + 120, None))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.answers.release(2) + + assert provider() == "token-0" + with pytest.raises(ClientAuthenticationError): + provider() + + +def test_callers_waiting_on_a_failed_token_request_share_one_retry_when_entra_fails_fast() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + start: Final = threading.Barrier(_FAST_FAILING_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_FAST_FAILING_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_FAST_FAILING_CALLERS)) + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release(_FAST_FAILING_CALLERS) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert credential.requests.qsize() == 1 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokens: + first_lifetime_seconds: int = 120 + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + refresh_started: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def next_token(self) -> tuple[str, int]: + self.requests.put(None) + if self.requests.qsize() == 1: + return "expiring-token", int(time.time()) + self.first_lifetime_seconds + self.refresh_started.set() + self.release.wait(5) + return "refreshed-token", _YEAR_2100 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenInfoCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + return AccessTokenInfo(*self.tokens.next_token()) + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenOnlyCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token( + self, + *_: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, + ) -> AccessToken: + return AccessToken(*self.tokens.next_token()) + + +@pytest.mark.parametrize("credential_type", (_ExpiringSoonTokenInfoCredential, _ExpiringSoonTokenOnlyCredential)) +def test_callers_during_a_refresh_keep_the_valid_token_while_one_of_them_refreshes( + credential_type: type[_ExpiringSoonTokenInfoCredential | _ExpiringSoonTokenOnlyCredential], +) -> None: + credential: Final = credential_type() + tokens: Final = credential.tokens + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert tuple(other.result(timeout=2) for other in others) == ("expiring-token",) * (_CONCURRENT_CALLERS - 1) + finally: + tokens.release.set() + + assert (refresh.result(), provider()) == ("refreshed-token", "refreshed-token") + assert tokens.requests.qsize() == 2 + + +def test_callers_during_a_refresh_wait_for_the_new_token_when_the_valid_one_is_about_to_expire() -> None: + tokens: Final = _ExpiringSoonTokens(first_lifetime_seconds=30) + provider: Final = single_flight_bearer_token_provider(_ExpiringSoonTokenInfoCredential(tokens), _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert not wait(others, timeout=_JOIN_WINDOW_SECONDS).done + finally: + tokens.release.set() + + assert tuple(call.result() for call in (refresh, *others)) == ("refreshed-token",) * _CONCURRENT_CALLERS + assert tokens.requests.qsize() == 2 diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index 7c6e1de6949..862658d9169 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -2,7 +2,11 @@ import base64 import json import logging import os +import threading import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Final from unittest.mock import Mock, patch @@ -12,6 +16,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.secret_managers.main import ( + _cached_azure_oidc_token_provider, get_secret, normalize_nonempty_secret_str, resolve_local_oidc_reference, @@ -313,6 +318,41 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider, monkeypat mock_token_provider.assert_called_once_with() +def test_concurrent_azure_oidc_reads_build_one_token_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AZURE_FEDERATED_TOKEN_FILE", raising=False) + callers: Final = 8 + start: Final = threading.Barrier(callers) + builds: Final = SimpleQueue[str | None]() + everyone_building: Final = threading.Event() + release: Final = threading.Event() + + def held_build(azure_scope: str | None = None) -> Callable[[], str]: + builds.put(azure_scope) + if builds.qsize() >= callers: + everyone_building.set() + release.wait(5) + return lambda: "identity-token" + + def read() -> str | None: + start.wait(5) + return get_secret("oidc/azure/api://concurrent-audience") + + _cached_azure_oidc_token_provider.cache_clear() + with ( + patch("litellm.secret_managers.main.get_azure_ad_token_provider", side_effect=held_build), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(read) for _ in range(callers)) + assert builds.get(timeout=5) == "api://concurrent-audience" + everyone_building.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + _cached_azure_oidc_token_provider.cache_clear() + + assert tokens == ("identity-token",) * callers + assert builds.qsize() == 0 + + def test_oidc_file_success(tmp_path, monkeypatch): token_file = tmp_path / "token.txt" token_file.write_text("file_token")