From 66ebc722d6751a7a8d0aef751e8d2e5f6a2efe49 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 8 Sep 2026 19:33:09 -0700
Subject: [PATCH] perf(azure): reuse the token refresh credential across image
 generation requests

With enable_azure_ad_token_refresh, every keyless image request built a new DefaultAzureCredential and fetched a token. Cache the provider per scope like the Entra ID one.
---
 litellm/llms/azure/common_utils.py            |  9 ++--
 .../test_azure_image_generation_init.py       | 47 +++++++++++++++++++
 .../llms/azure/test_azure_common_utils.py     |  3 ++
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py
index f276d8b18d1..9f70761514b 100644
--- a/litellm/llms/azure/common_utils.py
+++ b/litellm/llms/azure/common_utils.py
@@ -95,6 +95,11 @@ def _cached_entra_id_token_provider(
     return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope)
 
 
+@lru_cache(maxsize=128)
+def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]:
+    return get_azure_ad_token_provider(azure_scope=scope)
+
+
 def get_azure_ad_token_from_entra_id(
     tenant_id: str,
     client_id: str,
@@ -649,9 +654,7 @@ class BaseAzureLLM(BaseOpenAILLM):
                 "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth"
             )
             try:
-                azure_ad_token_provider = get_azure_ad_token_provider(
-                    azure_scope=scope,
-                )
+                azure_ad_token_provider = _cached_azure_ad_token_refresh_provider(scope)
             except ValueError:
                 verbose_logger.debug("Azure AD Token Provider could not be used.")
         if api_version is None:
diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
index a30aa277f3d..cfde1760389 100644
--- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
+++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
@@ -11,6 +11,7 @@ import litellm
 from litellm.caching.llm_caching_handler import LLMClientCache
 from litellm.llms.azure.azure import AzureChatCompletion
 from litellm.llms.azure.common_utils import (
+    _cached_azure_ad_token_refresh_provider,
     _cached_entra_id_token_provider,
     get_azure_request_auth_headers,
     redact_azure_auth_headers,
@@ -775,6 +776,52 @@ def test_azure_image_generation_with_api_key_keeps_api_key_header(
     assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***"
 
 
+@pytest.fixture
+def fake_default_azure_credential(monkeypatch: pytest.MonkeyPatch):
+    built_credentials = []
+
+    class FakeDefaultAzureCredential:
+        def __init__(self) -> None:
+            built_credentials.append(self)
+
+    for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", "AZURE_AD_TOKEN"):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setattr("azure.identity.DefaultAzureCredential", FakeDefaultAzureCredential)
+    monkeypatch.setattr(
+        "azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "default-credential-token"
+    )
+    monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True)
+    _cached_azure_ad_token_refresh_provider.cache_clear()
+    yield built_credentials
+    _cached_azure_ad_token_refresh_provider.cache_clear()
+
+
+def test_azure_image_generation_token_refresh_reuses_credential_across_requests(
+    respx_mock: respx.MockRouter, fake_default_azure_credential: list
+):
+    api_base = "https://my-resource.openai.azure.com"
+    api_version = "2025-04-01-preview"
+    route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1")
+
+    for _ in range(3):
+        AzureChatCompletion().image_generation(
+            prompt="a cat",
+            timeout=60.0,
+            optional_params={"n": 1, "size": "1024x1024"},
+            logging_obj=MagicMock(),
+            headers={"Content-Type": "application/json"},
+            model="gpt-image-1",
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params={"api_base": api_base, "api_version": api_version},
+        )
+
+    assert route.call_count == 3
+    assert all(call.request.headers["Authorization"] == "Bearer default-credential-token" for call in route.calls)
+    assert len(fake_default_azure_credential) == 1
+
+
 @pytest.mark.parametrize(
     "caller_auth_header",
     [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}],
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 f000abb4c9a..7189be7c052 100644
--- a/tests/test_litellm/llms/azure/test_azure_common_utils.py
+++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py
@@ -9,6 +9,7 @@ import pytest
 import litellm
 from litellm.llms.azure.common_utils import (
     BaseAzureLLM,
+    _cached_azure_ad_token_refresh_provider,
     _cached_entra_id_token_provider,
     get_azure_ad_token,
     get_azure_ad_token_from_entra_id,
@@ -34,6 +35,7 @@ def setup_mocks(monkeypatch):
     monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
     monkeypatch.delenv("AZURE_SCOPE", raising=False)
     monkeypatch.delenv("AZURE_AD_TOKEN", raising=False)
+    _cached_azure_ad_token_refresh_provider.cache_clear()
 
     with (
         patch(
@@ -78,6 +80,7 @@ def setup_mocks(monkeypatch):
             "logger": mock_logger,
             "select_url": mock_select_url,
         }
+    _cached_azure_ad_token_refresh_provider.cache_clear()
 
 
 def test_initialize_with_api_key(setup_mocks):