fix(azure): preserve AD token refresh in v1 OpenAI client path (#28627)

* fix(azure): preserve AD token refresh in v1 OpenAI client path

The /openai/v1/ code path (api_version in {"v1", "latest", "preview"})
constructs a plain OpenAI/AsyncOpenAI client, but only forwarded
`api_key` from `azure_client_params`. When `enable_azure_ad_token_refresh`
is set (or any AD-only auth), `api_key` is None and the client
constructor raised "The api_key client option must be set...", breaking
every Azure call with a v1 api_version.

The OpenAI SDK (>=2.20.0) accepts a callable for `api_key` and re-invokes
it on every request via `_refresh_api_key`, so we now forward
`azure_ad_token_provider` directly — preserving the per-request token
refresh behavior of the regular AzureOpenAI client and avoiding the
expiry hole that resolving the token once at client-creation time would
introduce. Static `azure_ad_token` strings fall through to `api_key`.

For the async path we wrap the sync provider returned by azure-identity
in an async function since AsyncOpenAI expects `Callable[[], Awaitable[str]]`.

Fixes #27945

https://claude.ai/code/session_01UnzrDSFUUgp5T2wRoPMxq5

* fix(azure): offload sync token provider to thread in v1 async wrapper

* fix(azure): include AD credential identity in v1 client cache key

---------

Co-authored-by: Claude <noreply@anthropic.com>
(cherry picked from commit 96a2e8b16d)

litellm/llms/azure/common_utils.py

@@ -1,3 +1,5 @@
+import asyncio
+import hashlib
 import json
 import os
 from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast
@@ -449,6 +451,25 @@ class BaseAzureLLM(BaseOpenAILLM):
         ] = None
         client_initialization_params: dict = locals()
         client_initialization_params["is_async"] = _is_async
+        _lp = litellm_params or {}
+        _ad_provider = _lp.get("azure_ad_token_provider")
+        _ad_token = _lp.get("azure_ad_token")
+        _client_secret = _lp.get("client_secret")
+        _azure_password = _lp.get("azure_password")
+        client_initialization_params["azure_ad_token"] = (
+            hashlib.sha256(_ad_token.encode()).hexdigest()
+            if isinstance(_ad_token, str)
+            else None
+        )
+        client_initialization_params["azure_ad_token_provider"] = (
+            f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}"
+            f"|tenant_id={_lp.get('tenant_id')}"
+            f"|client_id={_lp.get('client_id')}"
+            f"|client_secret={hashlib.sha256(_client_secret.encode()).hexdigest() if isinstance(_client_secret, str) else None}"
+            f"|azure_username={_lp.get('azure_username')}"
+            f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}"
+            f"|azure_scope={_lp.get('azure_scope')}"
+        )
         if client is None:
             cached_client = self.get_cached_openai_client(
                 client_initialization_params=client_initialization_params,
@@ -474,8 +495,29 @@ class BaseAzureLLM(BaseOpenAILLM):
             if self._is_azure_v1_api_version(api_version):
                 # Extract only params that OpenAI client accepts
                 # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview"
-                v1_params = {
-                    "api_key": azure_client_params.get("api_key"),
+                # The OpenAI client accepts a callable for `api_key` and re-invokes it
+                # on every request (via `_refresh_api_key`), so passing
+                # `azure_ad_token_provider` directly preserves Azure AD token refresh
+                # behavior that the regular AzureOpenAI client provides.
+                v1_api_key: Optional[Union[str, Callable[[], Any]]] = (
+                    azure_client_params.get("api_key")
+                    or azure_client_params.get("azure_ad_token_provider")
+                    or azure_client_params.get("azure_ad_token")
+                )
+                if _is_async is True and callable(v1_api_key):
+                    # AsyncOpenAI expects an async provider; wrap the sync provider
+                    # returned by azure-identity. Offload to a thread so a token
+                    # refresh (blocking HTTP call to AAD on cache miss) does not
+                    # stall the event loop.
+                    _sync_provider = v1_api_key
+
+                    async def _async_v1_api_key() -> str:
+                        return await asyncio.to_thread(_sync_provider)
+
+                    v1_api_key = _async_v1_api_key
+
+                v1_params: Dict[str, Any] = {
+                    "api_key": v1_api_key,
                     "base_url": f"{api_base}/openai/v1/",
                 }
                 if "timeout" in azure_client_params:

tests/test_litellm/llms/azure/test_azure_common_utils.py

@@ -1646,6 +1646,336 @@ def test_azure_v1_api_uses_openai_client(api_version):
         ), f"base_url should contain /openai/v1/, got {async_client.base_url}"
 
 
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_api_with_azure_ad_token_provider(api_version):
+    """
+    The v1 OpenAI client path must forward `azure_ad_token_provider` so Azure AD
+    auth works for `api_version` in {"v1", "latest", "preview"}.
+
+    Regression: https://github.com/BerriAI/litellm/issues/27945 — before the fix
+    the v1 branch only forwarded `api_key`, so AD-only configs raised
+    "The api_key client option must be set" on every request.
+
+    The OpenAI SDK accepts a callable for `api_key` and re-invokes it on every
+    request, so passing the provider directly preserves token refresh.
+    """
+    from openai import AsyncOpenAI, OpenAI
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+    token_value = "mock-azure-ad-token-from-provider"
+
+    def token_provider():
+        return token_value
+
+    init_return = {
+        "api_key": None,
+        "azure_endpoint": api_base,
+        "api_version": api_version,
+        "azure_ad_token": None,
+        "azure_ad_token_provider": token_provider,
+    }
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = init_return
+
+        client = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            _is_async=False,
+        )
+
+        assert isinstance(client, OpenAI)
+        # The SDK stores callables as `_api_key_provider` and refreshes
+        # `self.api_key` before each request.
+        assert client._api_key_provider is token_provider
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = init_return
+
+        async_client = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            _is_async=True,
+        )
+
+        assert isinstance(async_client, AsyncOpenAI)
+        # Async client requires an async provider; we wrap the sync provider
+        # so the SDK can `await` it.
+        assert async_client._api_key_provider is not None
+        assert async_client._api_key_provider is not token_provider
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_api_async_token_provider_resolves_to_current_token(api_version):
+    """
+    The async wrapper must call the underlying sync provider on each invocation
+    (not cache its first return value), so token rotation is honored.
+    """
+    import asyncio
+
+    from openai import AsyncOpenAI
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+    tokens = iter(["token-1", "token-2", "token-3"])
+
+    def rotating_provider():
+        return next(tokens)
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = {
+            "api_key": None,
+            "azure_endpoint": api_base,
+            "api_version": api_version,
+            "azure_ad_token": None,
+            "azure_ad_token_provider": rotating_provider,
+        }
+
+        async_client = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            _is_async=True,
+        )
+
+    assert isinstance(async_client, AsyncOpenAI)
+    loop = asyncio.new_event_loop()
+    try:
+        first = loop.run_until_complete(async_client._api_key_provider())
+        second = loop.run_until_complete(async_client._api_key_provider())
+    finally:
+        loop.close()
+
+    assert first == "token-1"
+    assert second == "token-2"
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_api_with_static_azure_ad_token(api_version):
+    """
+    When only `azure_ad_token` (a static string) is set, the v1 client should
+    receive it as `api_key`.
+    """
+    from openai import OpenAI
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+    token_value = "static-azure-ad-token"
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = {
+            "api_key": None,
+            "azure_endpoint": api_base,
+            "api_version": api_version,
+            "azure_ad_token": token_value,
+            "azure_ad_token_provider": None,
+        }
+
+        client = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            _is_async=False,
+        )
+
+        assert isinstance(client, OpenAI)
+        assert client.api_key == token_value
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_api_key_wins_over_ad_token(api_version):
+    """
+    Explicit `api_key` takes precedence over `azure_ad_token_provider` /
+    `azure_ad_token`, matching the priority documented in
+    `initialize_azure_sdk_client`.
+    """
+    from openai import OpenAI
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = {
+            "api_key": "explicit-key",
+            "azure_endpoint": api_base,
+            "api_version": api_version,
+            "azure_ad_token": "should-be-ignored",
+            "azure_ad_token_provider": lambda: "also-ignored",
+        }
+
+        client = base_llm.get_azure_openai_client(
+            api_key="explicit-key",
+            api_base=api_base,
+            api_version=api_version,
+            _is_async=False,
+        )
+
+        assert isinstance(client, OpenAI)
+        assert client.api_key == "explicit-key"
+        assert client._api_key_provider is None
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_client_cache_separates_distinct_ad_providers(api_version):
+    """
+    Two configs sharing api_base/api_version but with different AD token
+    providers must not share a cached OpenAI client, otherwise requests for
+    one config would be sent with another config's AD credentials.
+    """
+    from openai import AsyncOpenAI
+
+    litellm.in_memory_llm_clients_cache._cache = {}
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+
+    def provider_a():
+        return "token-a"
+
+    def provider_b():
+        return "token-b"
+
+    def _init_for(provider):
+        return {
+            "api_key": None,
+            "azure_endpoint": api_base,
+            "api_version": api_version,
+            "azure_ad_token": None,
+            "azure_ad_token_provider": provider,
+        }
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = _init_for(provider_a)
+        client_a = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params={"azure_ad_token_provider": provider_a},
+            _is_async=True,
+        )
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = _init_for(provider_b)
+        client_b = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params={"azure_ad_token_provider": provider_b},
+            _is_async=True,
+        )
+
+    assert isinstance(client_a, AsyncOpenAI)
+    assert isinstance(client_b, AsyncOpenAI)
+    assert client_a is not client_b
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_client_cache_separates_distinct_entra_credentials(api_version):
+    """
+    Configs that synthesize an AD provider from tenant_id/client_id/client_secret
+    must not share a cached client when those inputs differ.
+    """
+    from openai import AsyncOpenAI
+
+    litellm.in_memory_llm_clients_cache._cache = {}
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+
+    def synth_provider():
+        return "synthesized-token"
+
+    def _init_synth():
+        return {
+            "api_key": None,
+            "azure_endpoint": api_base,
+            "api_version": api_version,
+            "azure_ad_token": None,
+            "azure_ad_token_provider": synth_provider,
+        }
+
+    common = {
+        "api_key": None,
+        "api_base": api_base,
+        "api_version": api_version,
+        "_is_async": True,
+    }
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = _init_synth()
+        client_a = base_llm.get_azure_openai_client(
+            litellm_params={
+                "tenant_id": "tenant-a",
+                "client_id": "client-a",
+                "client_secret": "secret-a",
+            },
+            **common,
+        )
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = _init_synth()
+        client_b = base_llm.get_azure_openai_client(
+            litellm_params={
+                "tenant_id": "tenant-b",
+                "client_id": "client-b",
+                "client_secret": "secret-b",
+            },
+            **common,
+        )
+
+    assert isinstance(client_a, AsyncOpenAI)
+    assert isinstance(client_b, AsyncOpenAI)
+    assert client_a is not client_b
+
+
+@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"])
+def test_azure_v1_client_cache_reuses_for_identical_ad_config(api_version):
+    """
+    Identical AD configs should still share a cached client (regression guard
+    so the cache-key change doesn't accidentally disable caching).
+    """
+    from openai import AsyncOpenAI
+
+    litellm.in_memory_llm_clients_cache._cache = {}
+
+    base_llm = BaseAzureLLM()
+    api_base = "https://test.openai.azure.com"
+
+    def provider():
+        return "tok"
+
+    init_return = {
+        "api_key": None,
+        "azure_endpoint": api_base,
+        "api_version": api_version,
+        "azure_ad_token": None,
+        "azure_ad_token_provider": provider,
+    }
+
+    with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init:
+        mock_init.return_value = init_return
+        client_a = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params={"azure_ad_token_provider": provider},
+            _is_async=True,
+        )
+        client_b = base_llm.get_azure_openai_client(
+            api_key=None,
+            api_base=api_base,
+            api_version=api_version,
+            litellm_params={"azure_ad_token_provider": provider},
+            _is_async=True,
+        )
+
+    assert isinstance(client_a, AsyncOpenAI)
+    assert client_a is client_b
+
+
 def test_azure_traditional_api_uses_azure_openai_client():
     """
     Test that traditional Azure API versions still use AzureOpenAI client.