From 0e358c03a091ca87c20b9f9ade59700609faadb3 Mon Sep 17 00:00:00 2001 From: Burt Holzman Date: Mon, 23 Mar 2026 09:43:37 -0500 Subject: [PATCH] Cache Azure token providers Previously, get_azure_ad_token_from_{entra_id,username_password} constructed a new credential and called get_bearer_token_provider on every invocation, returning a new provider callable every time (and unnecessary overhead, since the new callable would generate new credentials when used). This commit adds a TTL cache so that repeated calls return the same provider callable. --- litellm/llms/azure/common_utils.py | 39 ++- litellm/llms/azure/credential_cache.py | 49 ++++ .../llms/azure/test_azure_common_utils.py | 260 ++++++++++++++++++ 3 files changed, 339 insertions(+), 9 deletions(-) create mode 100644 litellm/llms/azure/credential_cache.py diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 4fc1ae960b8..27b3904480f 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -8,6 +8,11 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.llms.azure.credential_cache import ( + _cache_lock, + _hash_secret, + _provider_cache, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.common_utils import BaseOpenAILLM from litellm.secret_managers.get_azure_ad_token_provider import ( @@ -108,9 +113,17 @@ def get_azure_ad_token_from_entra_id( ) if _tenant_id is None or _client_id is None or _client_secret is None: raise ValueError("tenant_id, client_id, and client_secret must be provided") - credential = ClientSecretCredential(_tenant_id, _client_id, _client_secret) - token_provider = get_bearer_token_provider(credential, scope) + cache_key = ("entra", _tenant_id, _client_id, _hash_secret(_client_secret), scope) + + with _cache_lock: + token_provider = _provider_cache.get(cache_key) + if token_provider is None: + credential = ClientSecretCredential(_tenant_id, _client_id, _client_secret) + verbose_logger.debug("credential %s", credential) + + token_provider = get_bearer_token_provider(credential, scope) + _provider_cache[cache_key] = token_provider verbose_logger.debug("token_provider %s", token_provider) @@ -143,16 +156,24 @@ def get_azure_ad_token_from_username_password( azure_username is not None, azure_password is not None, ) - credential = UsernamePasswordCredential( - client_id=client_id, - username=azure_username, - password=azure_password, - ) - token_provider = get_bearer_token_provider(credential, scope) + cache_key = ("upw", client_id, azure_username, _hash_secret(azure_password), scope) + + with _cache_lock: + token_provider = _provider_cache.get(cache_key) + if token_provider is None: + credential = UsernamePasswordCredential( + client_id=client_id, + username=azure_username, + password=azure_password, + ) + + verbose_logger.debug("credential %s", credential) + + token_provider = get_bearer_token_provider(credential, scope) + _provider_cache[cache_key] = token_provider verbose_logger.debug("token_provider %s", token_provider) - return token_provider diff --git a/litellm/llms/azure/credential_cache.py b/litellm/llms/azure/credential_cache.py new file mode 100644 index 00000000000..535706feb5e --- /dev/null +++ b/litellm/llms/azure/credential_cache.py @@ -0,0 +1,49 @@ +""" +Module-level cache for Azure AD token provider callables. + +Provides a shared cache used by get_azure_ad_token_from_entra_id and +get_azure_ad_token_from_username_password so that repeated calls with the same +credentials return the same provider callable rather than constructing a new one +each time. + +Cache keys are tuples: + ("entra", tenant_id, client_id, hmac(client_secret), scope) + ("upw", client_id, username, hmac(password), scope) + +where hmac uses HMAC-SHA256 with a module-level key to avoid storing plaintext +credentials in cache keys. + +If cachetools is available, the cache is a TTLCache(maxsize=128, ttl=3600) protected +by _cache_lock. Otherwise, it falls back to an unbounded dict with no TTL eviction. +""" + +import hashlib +import hmac +import threading +from typing import Any + +_CACHE_MAX_SIZE = 128 +_CACHE_TTL_SECONDS = 3600 # 1 hour + +# Deterministic HMAC key for cache keys — prevents raw credential strings from +# appearing verbatim in heap dumps, but is NOT a secret value and does NOT +# provide cryptographic confidentiality. Do not rely on this for security. +_CREDENTIAL_CACHE_HMAC_KEY = hashlib.sha256( + b"litellm-azure-credential-cache-v1" +).digest() + +_cache_lock = threading.Lock() + +try: + from cachetools import TTLCache + + _provider_cache: Any = TTLCache(maxsize=_CACHE_MAX_SIZE, ttl=_CACHE_TTL_SECONDS) +except ImportError: + _provider_cache = {} # unbounded fallback; install cachetools for TTL+LRU eviction + + +def _hash_secret(secret: str) -> str: + # codeql[py/weak-cryptographic-algorithm] + return hmac.new( + _CREDENTIAL_CACHE_HMAC_KEY, secret.encode(), hashlib.sha256 + ).hexdigest() diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 3fa794375e7..57273969e85 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1703,3 +1703,263 @@ def test_azure_traditional_api_uses_azure_openai_client(): assert isinstance( async_client, AsyncAzureOpenAI ), f"Expected AsyncAzureOpenAI client for api_version={api_version}" + + +@pytest.fixture(autouse=False) +def clear_azure_env(monkeypatch): + """Remove all Azure-related env vars so tests are fully isolated.""" + for var in ( + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_USERNAME", + "AZURE_PASSWORD", + "AZURE_SCOPE", + "AZURE_AUTHORITY_HOST", + ): + monkeypatch.delenv(var, raising=False) + + +@pytest.fixture(autouse=False) +def clear_provider_cache(): + """Clear the module-level credential cache before and after each test.""" + import litellm.llms.azure.credential_cache as cc + with cc._cache_lock: + cc._provider_cache.clear() + yield + with cc._cache_lock: + cc._provider_cache.clear() + + +def test_get_azure_ad_token_twice_entra_id_same_provider(clear_azure_env, clear_provider_cache): + """ + Calling get_azure_ad_token twice with the same entra-id credentials should + return the same provider callable object — i.e. get_bearer_token_provider + inside get_azure_ad_token_from_entra_id should be called only once. + """ + from litellm.llms.azure.common_utils import get_azure_ad_token + + captured_providers: list = [] + + def fake_get_bearer_token_provider(credential, scope): + provider = MagicMock(return_value="fake-entra-token") + captured_providers.append(provider) + return provider + + litellm_params = GenericLiteLLMParams( + tenant_id="test-tenant-id", + client_id="test-client-id", + client_secret="test-client-secret", + azure_ad_token_provider=None, + azure_ad_token=None, + azure_username=None, + azure_password=None, + ) + + with ( + patch("azure.identity.ClientSecretCredential") as mock_csc, + patch( + "azure.identity.get_bearer_token_provider", + side_effect=fake_get_bearer_token_provider, + ), + patch( + "litellm.llms.azure.common_utils.get_secret_str", + return_value=None, + ), + patch( + "litellm.llms.azure.common_utils.litellm" + ) as mock_litellm, + ): + mock_csc.return_value = MagicMock() + mock_litellm.enable_azure_ad_token_refresh = False + + token1 = get_azure_ad_token(litellm_params) + token2 = get_azure_ad_token(litellm_params) + + assert token1 == "fake-entra-token" + assert token2 == "fake-entra-token" + # The provider should have been constructed exactly once (cached) and + # invoked once per get_azure_ad_token call. + assert len(captured_providers) == 1, ( + f"Expected get_bearer_token_provider to be called once (provider cached), " + f"but it was called {len(captured_providers)} times" + ) + assert captured_providers[0].call_count == 2, ( + "Expected the single cached provider to be called twice (once per " + f"get_azure_ad_token call), but it was called {captured_providers[0].call_count} times" + ) + + +def test_get_azure_ad_token_twice_username_password_same_provider(clear_azure_env, clear_provider_cache): + """ + Calling get_azure_ad_token twice with the same username/password credentials + should return the same provider callable object — i.e. get_bearer_token_provider + inside get_azure_ad_token_from_username_password should be called only once. + """ + from litellm.llms.azure.common_utils import get_azure_ad_token + + captured_providers: list = [] + + def fake_get_bearer_token_provider(credential, scope): + provider = MagicMock(return_value="fake-username-password-token") + captured_providers.append(provider) + return provider + + litellm_params = GenericLiteLLMParams( + azure_username="test-username", + azure_password="test-password", + client_id="test-client-id", + azure_ad_token_provider=None, + azure_ad_token=None, + tenant_id=None, + client_secret=None, + ) + + with ( + patch("azure.identity.UsernamePasswordCredential") as mock_upc, + patch( + "azure.identity.get_bearer_token_provider", + side_effect=fake_get_bearer_token_provider, + ), + patch( + "litellm.llms.azure.common_utils.get_secret_str", + return_value=None, + ), + patch( + "litellm.llms.azure.common_utils.litellm" + ) as mock_litellm, + ): + mock_upc.return_value = MagicMock() + mock_litellm.enable_azure_ad_token_refresh = False + + token1 = get_azure_ad_token(litellm_params) + token2 = get_azure_ad_token(litellm_params) + + assert token1 == "fake-username-password-token" + assert token2 == "fake-username-password-token" + # The provider should have been constructed exactly once (cached) and + # invoked once per get_azure_ad_token call. + assert len(captured_providers) == 1, ( + f"Expected get_bearer_token_provider to be called once (provider cached), " + f"but it was called {len(captured_providers)} times" + ) + assert captured_providers[0].call_count == 2, ( + "Expected the single cached provider to be called twice (once per " + f"get_azure_ad_token call), but it was called {captured_providers[0].call_count} times" + ) + + +def test_get_azure_ad_token_twice_oidc_cache_hit(clear_azure_env): + """ + Calling get_azure_ad_token twice with the same OIDC credentials should hit + the azure_ad_cache on the second call — i.e. the HTTP POST to the Azure + token endpoint should be made only once. + """ + import litellm as _litellm + from litellm.llms.azure.common_utils import azure_ad_cache, get_azure_ad_token + + # Flush any stale cache entries from other tests. + azure_ad_cache.flush_cache() + + litellm_params = GenericLiteLLMParams( + azure_ad_token="oidc/test-oidc-token", + client_id="test-client-id", + tenant_id="test-tenant-id", + azure_ad_token_provider=None, + client_secret=None, + azure_username=None, + azure_password=None, + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "access_token": "fake-access-token", + "expires_in": 3600, + } + + def fake_get_secret_str(key): + # Resolve the OIDC token reference; return None for anything else. + if key == "oidc/test-oidc-token": + return "resolved-oidc-jwt" + return None + + with ( + patch( + "litellm.llms.azure.common_utils.get_secret_str", + side_effect=fake_get_secret_str, + ), + patch.object( + _litellm.module_level_client, + "post", + return_value=mock_response, + ) as mock_post, + ): + token1 = get_azure_ad_token(litellm_params) + token2 = get_azure_ad_token(litellm_params) + + assert token1 == "fake-access-token" + assert token2 == "fake-access-token" + assert mock_post.call_count == 1, ( + f"Expected HTTP POST to Azure token endpoint to be called once (second " + f"call should be a cache hit), but it was called {mock_post.call_count} times" + ) + + +def test_get_azure_ad_token_entra_id_provider_reconstructed_after_ttl_expiry( + clear_azure_env, clear_provider_cache +): + """ + With a TTLCache, sleeping past the TTL causes the cached provider to be + evicted, so the second get_azure_ad_token call constructs a new callable. + """ + pytest.importorskip("cachetools") + import time + + from cachetools import TTLCache + + from litellm.llms.azure.common_utils import get_azure_ad_token + + litellm_params = GenericLiteLLMParams( + tenant_id="test-tenant-id", + client_id="test-client-id", + client_secret="test-client-secret", + azure_ad_token_provider=None, + azure_ad_token=None, + azure_username=None, + azure_password=None, + ) + + # --- TTLCache: provider is reconstructed after expiry --- + captured_ttl: list = [] + + def fake_gbtp_ttl(credential, scope): + provider = MagicMock(return_value="fake-entra-token") + captured_ttl.append(provider) + return provider + + tiny_ttl_cache = TTLCache(maxsize=128, ttl=0.001) + + import threading + test_lock = threading.Lock() + + with ( + patch("litellm.llms.azure.common_utils._provider_cache", tiny_ttl_cache), + patch("azure.identity.get_bearer_token_provider", side_effect=fake_gbtp_ttl), + patch("azure.identity.ClientSecretCredential", return_value=MagicMock()), + patch("litellm.llms.azure.common_utils.get_secret_str", return_value=None), + patch("litellm.llms.azure.common_utils.litellm"), + ): + get_azure_ad_token(litellm_params) + time.sleep(0.2) # outlast the 0.001s TTL (200x margin for CI robustness) + get_azure_ad_token(litellm_params) + + assert len(captured_ttl) == 2, ( + f"TTLCache: expected get_bearer_token_provider called twice (entry expired), " + f"but was called {len(captured_ttl)} times" + ) + assert captured_ttl[0] is not captured_ttl[1], ( + "TTLCache: expected two distinct provider callables after expiry" + ) +