From 84888aa2ebb5ed89b8ea027e14ed161e455e3082 Mon Sep 17 00:00:00 2001 From: micli Date: Mon, 2 Mar 2026 09:45:19 +0800 Subject: [PATCH] fix(azure): support system-assigned managed identity token provider --- litellm/llms/azure/azure.py | 14 ++ .../get_azure_ad_token_provider.py | 6 +- .../test_azure_image_generation_init.py | 203 ++++++++++++++++++ .../test_get_azure_ad_token_provider.py | 31 +++ 4 files changed, 253 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 61cfd54b565..874d50d80d4 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1251,6 +1251,20 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_base=api_base, is_async=False, ) + + # If MI/DefaultAzureCredential resolved a token provider inside + # initialize_azure_sdk_client, use it to set the Authorization header + # for the raw httpx request (SDK clients handle this internally, but + # image generation uses raw httpx). + if api_key is None and "Authorization" not in headers: + resolved_provider = azure_client_params.get("azure_ad_token_provider") + resolved_token = azure_client_params.get("azure_ad_token") + if resolved_provider is not None and resolved_token is None: + resolved_token = resolved_provider() + if resolved_token: + headers.pop("api-key", None) + headers["Authorization"] = f"Bearer {resolved_token}" + if aimg_generation is True: return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 184d959b964..a54f0ff7d5f 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -92,7 +92,11 @@ def get_azure_ad_token_provider( tenant_id=os.environ["AZURE_TENANT_ID"], ) elif cred == AzureCredentialType.ManagedIdentityCredential: - credential = ManagedIdentityCredential(client_id=os.environ["AZURE_CLIENT_ID"]) + managed_identity_client_id = os.environ.get("AZURE_CLIENT_ID") + if managed_identity_client_id: + credential = ManagedIdentityCredential(client_id=managed_identity_client_id) + else: + credential = ManagedIdentityCredential() elif cred == AzureCredentialType.CertificateCredential: if os.getenv("AZURE_CERTIFICATE_PASSWORD"): credential = CertificateCredential( 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 44bcc9f954a..560ba535615 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 @@ -448,3 +448,206 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name(): f"Request body 'model' field should be base_model '{base_model}', " f"but got: {request_data.get('model')}" ) + +class TestAzureImageGenerationManagedIdentityAuth: + """ + Tests for Azure image generation with Managed Identity (MI) / DefaultAzureCredential. + + The image generation code path uses raw httpx requests (not the Azure SDK client), + so the MI token provider resolved by initialize_azure_sdk_client must be explicitly + used to set the Authorization header. + """ + + def test_should_set_authorization_header_from_mi_token_provider(self): + """ + Test that when Managed Identity resolves a token provider via + initialize_azure_sdk_client, the Authorization header is set on the + raw httpx request used for image generation. + """ + mock_token = "mock-azure-ad-token-from-mi-12345" + mock_token_provider = Mock(return_value=mock_token) + + azure_chat = AzureChatCompletion() + + mock_client_params = { + "api_key": None, + "azure_endpoint": "https://my-aoai.openai.azure.com/", + "api_version": "2024-06-01", + "azure_ad_token": None, + "azure_ad_token_provider": mock_token_provider, + } + + with patch.object( + azure_chat, "initialize_azure_sdk_client", return_value=mock_client_params + ), patch.object( + azure_chat, "aimage_generation", new_callable=AsyncMock + ) as mock_aimg: + mock_aimg.return_value = MagicMock() + + azure_chat.image_generation( + prompt="test prompt", + timeout=30.0, + optional_params={}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base="https://my-aoai.openai.azure.com/", + api_version="2024-06-01", + azure_ad_token=None, + azure_ad_token_provider=None, + aimg_generation=True, + litellm_params={}, + ) + + call_kwargs = mock_aimg.call_args + headers_passed = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {}) + assert "Authorization" in headers_passed, ( + "Authorization header should be set when MI token provider is resolved" + ) + assert headers_passed["Authorization"] == f"Bearer {mock_token}" + assert "api-key" not in headers_passed, ( + "api-key header should be removed when using MI token" + ) + + def test_should_use_explicit_token_provider_over_resolved_one(self): + """ + Test that if azure_ad_token_provider is already passed (not None), + its token takes priority over the one resolved by initialize_azure_sdk_client. + """ + explicit_token = "explicit-token-passed-by-caller" + explicit_provider = Mock(return_value=explicit_token) + + azure_chat = AzureChatCompletion() + + mock_client_params = { + "api_key": None, + "azure_endpoint": "https://my-aoai.openai.azure.com/", + "api_version": "2024-06-01", + "azure_ad_token": None, + "azure_ad_token_provider": Mock(return_value="resolved-by-init-sdk-client"), + } + + with patch.object( + azure_chat, "initialize_azure_sdk_client", return_value=mock_client_params + ), patch.object( + azure_chat, "aimage_generation", new_callable=AsyncMock + ) as mock_aimg: + mock_aimg.return_value = MagicMock() + + azure_chat.image_generation( + prompt="test prompt", + timeout=30.0, + optional_params={}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base="https://my-aoai.openai.azure.com/", + api_version="2024-06-01", + azure_ad_token=None, + azure_ad_token_provider=explicit_provider, + aimg_generation=True, + litellm_params={}, + ) + + call_kwargs = mock_aimg.call_args + headers_passed = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {}) + assert headers_passed["Authorization"] == f"Bearer {explicit_token}" + explicit_provider.assert_called_once() + + def test_should_not_set_authorization_when_api_key_present(self): + """ + Test that when api_key is provided, MI token resolution is skipped. + """ + azure_chat = AzureChatCompletion() + + mock_client_params = { + "api_key": "my-api-key", + "azure_endpoint": "https://my-aoai.openai.azure.com/", + "api_version": "2024-06-01", + "azure_ad_token": None, + "azure_ad_token_provider": None, + } + + with patch.object( + azure_chat, "initialize_azure_sdk_client", return_value=mock_client_params + ), patch.object( + azure_chat, "aimage_generation", new_callable=AsyncMock + ) as mock_aimg: + mock_aimg.return_value = MagicMock() + + azure_chat.image_generation( + prompt="test prompt", + timeout=30.0, + optional_params={}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json", "api-key": "my-api-key"}, + model="gpt-image-1", + api_key="my-api-key", + api_base="https://my-aoai.openai.azure.com/", + api_version="2024-06-01", + azure_ad_token=None, + azure_ad_token_provider=None, + aimg_generation=True, + litellm_params={}, + ) + + call_kwargs = mock_aimg.call_args + headers_passed = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {}) + assert "Authorization" not in headers_passed + assert headers_passed.get("api-key") == "my-api-key" + + def test_should_set_authorization_in_sync_path(self): + """ + Test the sync path: MI token should be set in headers for sync httpx request. + """ + mock_token = "sync-mi-token-67890" + mock_token_provider = Mock(return_value=mock_token) + + azure_chat = AzureChatCompletion() + + mock_client_params = { + "api_key": None, + "azure_endpoint": "https://my-aoai.openai.azure.com/", + "api_version": "2024-06-01", + "azure_ad_token": None, + "azure_ad_token_provider": mock_token_provider, + } + + mock_httpx_response = MagicMock() + mock_httpx_response.json.return_value = { + "created": 1234567890, + "data": [{"url": "https://example.com/image.png", "revised_prompt": "test"}], + } + + with patch.object( + azure_chat, "initialize_azure_sdk_client", return_value=mock_client_params + ), patch.object( + azure_chat, "create_azure_base_url", return_value="https://my-aoai.openai.azure.com/openai/deployments/gpt-image-1/images/generations?api-version=2024-06-01" + ), patch.object( + azure_chat, "make_sync_azure_httpx_request", return_value=mock_httpx_response + ) as mock_sync_request: + + azure_chat.image_generation( + prompt="test prompt", + timeout=30.0, + optional_params={}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base="https://my-aoai.openai.azure.com/", + api_version="2024-06-01", + azure_ad_token=None, + azure_ad_token_provider=None, + aimg_generation=False, + litellm_params={}, + ) + + call_kwargs = mock_sync_request.call_args + headers_passed = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers", {}) + assert "Authorization" in headers_passed, ( + "Sync path: Authorization header should be set from MI token provider" + ) + assert headers_passed["Authorization"] == f"Bearer {mock_token}" \ No newline at end of file diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py index f02f59cccc0..c8a8cbd696a 100644 --- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py @@ -95,6 +95,37 @@ class TestGetAzureAdTokenProvider: token = result() assert token == "mock-managed-identity-token" + @patch.dict( + os.environ, + { + "AZURE_SCOPE": "https://cognitiveservices.azure.com/.default", + "AZURE_CREDENTIAL": "ManagedIdentityCredential", + }, + clear=True, + ) + @patch("azure.identity.get_bearer_token_provider") + @patch("azure.identity.ManagedIdentityCredential") + def test_get_azure_ad_token_provider_managed_identity_credential_without_client_id( + self, mock_managed_identity_credential, mock_get_bearer_token_provider + ): + """Test get_azure_ad_token_provider with system-assigned ManagedIdentityCredential.""" + mock_credential_instance = MagicMock() + mock_managed_identity_credential.return_value = mock_credential_instance + + mock_token_provider = MagicMock(return_value="mock-system-managed-identity-token") + mock_get_bearer_token_provider.return_value = mock_token_provider + + result = get_azure_ad_token_provider() + + assert callable(result) + mock_managed_identity_credential.assert_called_once_with() + mock_get_bearer_token_provider.assert_called_once_with( + mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ) + + token = result() + assert token == "mock-system-managed-identity-token" + @patch.dict( os.environ, {