fix(azure): exchange oidc azure_ad_token lazily through the threaded token provider

An oidc/ azure_ad_token was exchanged for an Entra access token while building the
SDK client, so the first async chat request ran the federated token mint and the
token endpoint POST on the event loop. The exchange now runs as the client's token
provider, which async clients already call in a worker thread, and the exchange
cache keeps later calls cheap while letting a cached client pick up a fresh token

litellm/llms/azure/common_utils.py

@@ -4,7 +4,7 @@ import json
 import os
 import re
 from collections.abc import Awaitable, Callable, Mapping
-from functools import lru_cache
+from functools import lru_cache, partial
 from types import MappingProxyType
 from typing import Final, Literal, NamedTuple, cast
 
@@ -703,12 +703,14 @@ class BaseAzureLLM(BaseOpenAILLM):
 
         if azure_ad_token is not None and azure_ad_token.startswith("oidc/"):
             verbose_logger.debug("Using Azure OIDC Token for Azure Auth")
-            azure_ad_token = get_azure_ad_token_from_oidc(
+            azure_ad_token_provider = partial(
+                get_azure_ad_token_from_oidc,
                 azure_ad_token=azure_ad_token,
                 azure_client_id=client_id,
                 azure_tenant_id=tenant_id,
                 scope=scope,
             )
+            azure_ad_token = None
         elif not api_key and azure_ad_token_provider is None and litellm.enable_azure_ad_token_refresh is True:
             verbose_logger.debug(
                 "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth"

tests/unit/llms/azure/test_azure_common_utils.py

@@ -223,6 +223,7 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch):
         is_async=False,
     )
 
+    assert result["azure_ad_token_provider"]() == "mock-oidc-token"
     setup_mocks["oidc_token"].assert_called_once_with(
         azure_ad_token="oidc/test-token",
         azure_client_id=None,
@@ -230,8 +231,6 @@ def test_initialize_with_oidc_token(setup_mocks, monkeypatch):
         scope="https://cognitiveservices.azure.com/.default",
     )
 
-    # Verify expected result
-    assert result["azure_ad_token"] == "mock-oidc-token"
 
 
 def test_initialize_with_oidc_token_and_client_params(setup_mocks):
@@ -251,6 +250,7 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks):
     )
 
     # Verify that get_azure_ad_token_from_oidc was called with the correct parameters
+    assert result["azure_ad_token_provider"]() == "mock-oidc-token"
     setup_mocks["oidc_token"].assert_called_once_with(
         azure_ad_token="oidc/test-token",
         azure_client_id="test-client-id",
@@ -258,8 +258,6 @@ def test_initialize_with_oidc_token_and_client_params(setup_mocks):
         scope="test-azure-scope",
     )
 
-    # Verify expected result
-    assert result["azure_ad_token"] == "mock-oidc-token"
 
 
 def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch):
@@ -280,6 +278,7 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch):
     )
 
     # Verify that get_azure_ad_token_from_oidc was called with environment variables
+    assert result["azure_ad_token_provider"]() == "mock-oidc-token"
     setup_mocks["oidc_token"].assert_called_once_with(
         azure_ad_token="oidc/test-token",
         azure_client_id="env-client-id",
@@ -287,8 +286,6 @@ def test_initialize_with_oidc_token_fallback_to_env(setup_mocks, monkeypatch):
         scope="https://cognitiveservices.azure.com/.default",
     )
 
-    # Verify expected result
-    assert result["azure_ad_token"] == "mock-oidc-token"
 
 
 def test_initialize_with_ad_token_provider(setup_mocks, monkeypatch):
@@ -2217,6 +2214,69 @@ async def test_non_v1_azure_token_refresh_does_not_block_the_event_loop():
         await client.close()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("api_version", ["2024-02-01", "v1"])
+async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_request(monkeypatch, api_version):
+    import threading
+
+    import httpx
+
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    exchange_threads: Final[list[int]] = []
+
+    def entra_token_endpoint(request: httpx.Request) -> httpx.Response:
+        exchange_threads.append(threading.get_ident())
+        assert b"client_assertion=federated-oidc-token" in request.content
+        return httpx.Response(200, json={"access_token": f"exchanged-{api_version}", "expires_in": 3600})
+
+    def echo_authorization(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json={
+                "id": "chatcmpl-1",
+                "object": "chat.completion",
+                "created": 0,
+                "model": "gpt",
+                "choices": [
+                    {
+                        "index": 0,
+                        "finish_reason": "stop",
+                        "message": {"role": "assistant", "content": request.headers["authorization"]},
+                    }
+                ],
+            },
+        )
+
+    monkeypatch.setenv("FEDERATED_OIDC_TOKEN", "federated-oidc-token")
+    monkeypatch.setattr(
+        litellm,
+        "module_level_client",
+        HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(entra_token_endpoint))),
+    )
+    monkeypatch.setattr(
+        litellm, "aclient_session", httpx.AsyncClient(transport=httpx.MockTransport(echo_authorization))
+    )
+    loop_thread: Final = threading.get_ident()
+    client: Final = BaseAzureLLM().get_azure_openai_client(
+        api_key=None,
+        api_base=f"https://oidc-exchange-{api_version}.openai.azure.com",
+        api_version=api_version,
+        litellm_params={
+            "azure_ad_token": "oidc/env/FEDERATED_OIDC_TOKEN",
+            "client_id": f"oidc-client-{api_version}",
+            "tenant_id": "oidc-tenant",
+        },
+        _is_async=True,
+    )
+
+    response: Final = await client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}])
+
+    assert response.choices[0].message.content == f"Bearer exchanged-{api_version}"
+    assert len(exchange_threads) == 1
+    assert loop_thread not in exchange_threads
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("api_version", ["2024-02-01", "v1"])
 async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token(monkeypatch, api_version):