From 7103bfd003096acec3ff88bb5495513da25e549c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:52:25 -0700 Subject: [PATCH] fix(azure): await async azure_ad_token_provider results in the threaded wrapper --- litellm/llms/azure/common_utils.py | 14 +++---- litellm/llms/custom_httpx/llm_http_handler.py | 20 ++++++--- .../llms/azure/test_azure_common_utils.py | 42 +++++++++++++++++++ 3 files changed, 63 insertions(+), 13 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index fdf19eb9cf8..cce2265d280 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -46,7 +46,6 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object): _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams) -_AZURE_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter(Callable[[], str]) class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object): @@ -136,10 +135,11 @@ def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: def _async_azure_ad_token_provider( - sync_provider: Callable[[], str], + provider: Callable[[], str | Awaitable[str]], ) -> Callable[[], Awaitable[str]]: async def async_provider() -> str: - return await asyncio.to_thread(sync_provider) + token: Final = await asyncio.to_thread(provider) + return token if isinstance(token, str) else await token return async_provider @@ -607,7 +607,7 @@ class BaseAzureLLM(BaseOpenAILLM): or v1_client_params.get("azure_ad_token") ) v1_api_key: Final[str | Callable[[], str | Awaitable[str]] | None] = ( - _async_azure_ad_token_provider(_AZURE_TOKEN_PROVIDER_ADAPTER.validate_python(v1_token_provider)) + _async_azure_ad_token_provider(v1_token_provider) if _is_async is True and callable(v1_token_provider) else v1_token_provider ) @@ -639,11 +639,9 @@ class BaseAzureLLM(BaseOpenAILLM): async_azure_client_params_with_threaded_provider: Final[_AsyncAzureClientParams] = ( { **async_azure_client_params, - "azure_ad_token_provider": _async_azure_ad_token_provider( - _AZURE_TOKEN_PROVIDER_ADAPTER.validate_python(async_token_provider_candidate) - ), + "azure_ad_token_provider": _async_azure_ad_token_provider(async_token_provider_candidate), } - if callable(async_token_provider_candidate) + if async_token_provider_candidate is not None else async_azure_client_params ) openai_client = AsyncAzureOpenAI(**async_azure_client_params_with_threaded_provider) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 66ed80b7afb..fc35531f73d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -3148,7 +3148,9 @@ class BaseLLMHTTPHandler: headers: Final = await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=extra_headers or {}, model="None", litellm_params=litellm_params + headers=extra_headers or {}, + model="None", + litellm_params=litellm_params, ) if extra_headers: @@ -3403,7 +3405,9 @@ class BaseLLMHTTPHandler: headers: Final = await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=extra_headers or {}, model="None", litellm_params=litellm_params + headers=extra_headers or {}, + model="None", + litellm_params=litellm_params, ) if extra_headers: @@ -3568,7 +3572,9 @@ class BaseLLMHTTPHandler: headers: Final = await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=extra_headers or {}, model="None", litellm_params=litellm_params + headers=extra_headers or {}, + model="None", + litellm_params=litellm_params, ) if extra_headers: @@ -4626,7 +4632,9 @@ class BaseLLMHTTPHandler: headers: Final = await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=extra_headers or {}, model="None", litellm_params=litellm_params + headers=extra_headers or {}, + model="None", + litellm_params=litellm_params, ) if extra_headers: @@ -4802,7 +4810,9 @@ class BaseLLMHTTPHandler: validated_headers: Final = await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=extra_headers or {}, model=model, litellm_params=litellm_params + headers=extra_headers or {}, + model=model, + litellm_params=litellm_params, ) if extra_headers: diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 4a4201c5fe9..c60e3e8a9d8 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -2217,6 +2217,48 @@ 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_async_azure_ad_token_provider_reaches_azure_as_bearer_token(monkeypatch, api_version): + import httpx + + 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"]}, + } + ], + }, + ) + + async def async_token_provider() -> str: + return "async-token" + + monkeypatch.setattr( + litellm, "aclient_session", httpx.AsyncClient(transport=httpx.MockTransport(echo_authorization)) + ) + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key=None, + api_base=f"https://async-token-provider-{api_version}.openai.azure.com", + api_version=api_version, + litellm_params={"azure_ad_token_provider": async_token_provider}, + _is_async=True, + ) + + response: Final = await client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "Bearer async-token" + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client.