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

@@ -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
 
 

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()

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"
+    )
+