From 456b97e635fe87fbb63adbb18eb4837634d2f93f Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 16 Jul 2026 14:15:24 +0000
Subject: [PATCH] feat(secret_managers): support Azure Workload Identity for
 Key Vault auth

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
---
 litellm/proxy/proxy_server.py                 | 26 ++++++++-
 .../proxy/proxy_server/test_lifecycle.py      | 54 +++++++++++++++++++
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 6661474d215..e51e1f0b16b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -126,6 +126,7 @@ from litellm.utils import (
 
 if TYPE_CHECKING:
     from aiohttp import ClientSession
+    from azure.core.credentials import TokenCredential
     from opentelemetry.trace import Span as _Span
 
     from litellm.integrations.opentelemetry import OpenTelemetry
@@ -2040,12 +2041,33 @@ def _resolve_pydantic_type(typ) -> List:
     return typs
 
 
+def _get_azure_key_vault_credential() -> "TokenCredential":
+    """
+    Build the credential used to authenticate to Azure Key Vault.
+
+    When ``AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY`` is truthy, use
+    ``WorkloadIdentityCredential`` (federated token auth for AKS workload identity,
+    no client secret needed); it reads ``AZURE_TENANT_ID``, ``AZURE_CLIENT_ID`` and
+    ``AZURE_FEDERATED_TOKEN_FILE`` from the environment. Otherwise fall back to
+    ``DefaultAzureCredential``.
+    """
+    from litellm.secret_managers.main import str_to_bool
+
+    if str_to_bool(os.getenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY")) is True:
+        from azure.identity import WorkloadIdentityCredential
+
+        return WorkloadIdentityCredential()
+
+    from azure.identity import DefaultAzureCredential
+
+    return DefaultAzureCredential()
+
+
 def load_from_azure_key_vault(use_azure_key_vault: bool = False):
     if use_azure_key_vault is False:
         return
 
     try:
-        from azure.identity import DefaultAzureCredential
         from azure.keyvault.secrets import SecretClient
 
         # Set your Azure Key Vault URI
@@ -2054,7 +2076,7 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False):
         if KVUri is None:
             raise Exception("Error when loading keys from Azure Key Vault: AZURE_KEY_VAULT_URI is not set.")
 
-        credential = DefaultAzureCredential()
+        credential = _get_azure_key_vault_credential()
 
         # Create the SecretClient using the credential
         client = SecretClient(vault_url=KVUri, credential=credential)
diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
index a3f5049ef1d..e75a6c435d5 100644
--- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
+++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
@@ -417,6 +417,60 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch)
     assert result is None
 
 
+def _install_fake_azure_identity(monkeypatch):
+    """Install a fake ``azure.identity`` so credential selection can be exercised
+    without the real Azure SDK installed. Returns the two sentinel classes."""
+    import sys
+    import types
+
+    class FakeDefaultAzureCredential:
+        pass
+
+    class FakeWorkloadIdentityCredential:
+        pass
+
+    azure_module = types.ModuleType("azure")
+    identity_module = types.ModuleType("azure.identity")
+    identity_module.DefaultAzureCredential = FakeDefaultAzureCredential
+    identity_module.WorkloadIdentityCredential = FakeWorkloadIdentityCredential
+    azure_module.identity = identity_module
+
+    monkeypatch.setitem(sys.modules, "azure", azure_module)
+    monkeypatch.setitem(sys.modules, "azure.identity", identity_module)
+    return FakeDefaultAzureCredential, FakeWorkloadIdentityCredential
+
+
+def test_get_azure_key_vault_credential_defaults_to_default_credential(monkeypatch):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.delenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", raising=False)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, default_cls)
+    assert not isinstance(credential, workload_cls)
+
+
+@pytest.mark.parametrize("truthy", ["true", "True", "TRUE"])
+def test_get_azure_key_vault_credential_uses_workload_identity_when_enabled(monkeypatch, truthy):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.setenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", truthy)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, workload_cls)
+    assert not isinstance(credential, default_cls)
+
+
+@pytest.mark.parametrize("falsy", ["false", "False", "0", "unset-like"])
+def test_get_azure_key_vault_credential_uses_default_when_disabled(monkeypatch, falsy):
+    default_cls, workload_cls = _install_fake_azure_identity(monkeypatch)
+    monkeypatch.setenv("AZURE_KEY_VAULT_USE_WORKLOAD_IDENTITY", falsy)
+
+    credential = ps._get_azure_key_vault_credential()
+
+    assert isinstance(credential, default_cls)
+
+
 # ---------------------------------------------------------------------------
 # cost_tracking
 # ---------------------------------------------------------------------------