feat: add ability to auth to azure with token - #21764
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR adds passwordless Azure AD (Entra ID) authentication for Redis in LiteLLM, mirroring the existing GCP IAM auth pattern. It introduces Key observations:
Confidence Score: 3/5
|
| Filename | Overview |
|---|---|
| litellm/_redis.py | Core implementation of Azure AD Redis auth — adds _build_azure_credential, create_azure_ad_redis_connect_func, and _generate_azure_ad_redis_token; correctly removes Azure kwargs from redis_kwargs; stores raw client credential value on function object (style concern); async cluster/pool/standard-async paths all fetch a one-shot token that will expire without a refresh path. |
| tests/test_litellm/test_utils.py | New Azure AD Redis tests added; most tests correctly use patch.dict("sys.modules", ...) for isolation, but test_redis_client_logic_azure_ad_auth calls _get_redis_client_logic with Azure AD enabled without mocking azure.identity, causing it to fail if azure-identity is not installed in CI. |
| docs/my-website/docs/caching/azure_redis_passwordless.md | New documentation page covering all three Azure auth modes (system-assigned managed identity, user-assigned managed identity, service principal); accurate and covers SSL requirement, prerequisites, and how the token rotation works. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[azure_redis_ad_token enabled] --> B[_get_redis_client_logic]
B --> C[_build_azure_credential]
C --> D{Which auth mode?}
D -->|Service Principal env vars| E[ClientSecretCredential]
D -->|User-Assigned Identity| F[ManagedIdentityCredential]
D -->|System-Assigned Identity| G[DefaultAzureCredential]
E & F & G --> H[Reusable credential object]
H --> I{Client type}
I -->|Sync redis.Redis| J[create_azure_ad_redis_connect_func]
J --> K[ad_connect closure - credential reused]
K --> L[Called per connection/reconnect]
L --> M[get_token on each connect]
M --> N[AUTH with fresh token - auto-renewed]
I -->|Async RedisCluster| O[_generate_azure_ad_redis_token]
I -->|Async Redis standard| O
I -->|Async ConnectionPool| O
O --> P[One-shot fetch at creation time]
P --> Q[Set as static password]
Q --> R[Expires in approx 1 hour - no refresh path]
Last reviewed commit: 32877e0
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
Automated patch bundle from next-100 unresolved backlog expansion.\nGenerated due limited direct branch-write access; please apply/cherry-pick minimal edits below.\n\n## PR #21764 — Unresolved thread summary
Minimal patch proposals
|
810ee68 to
9f2a828
Compare
| def test_generate_azure_ad_redis_token_import_error(): | ||
| """Test that _generate_azure_ad_redis_token raises ImportError when azure-identity is missing.""" | ||
| from unittest.mock import patch | ||
| from litellm._redis import _generate_azure_ad_redis_token | ||
|
|
||
| original_import = __builtins__["__import__"] | ||
|
|
||
| def mock_import(name, *args, **kwargs): | ||
| if name == "azure.identity": | ||
| raise ImportError("No module named 'azure.identity'") | ||
| return original_import(name, *args, **kwargs) | ||
|
|
||
| with patch("builtins.__import__", side_effect=mock_import): | ||
| with pytest.raises(ImportError) as exc_info: | ||
| _generate_azure_ad_redis_token() | ||
|
|
||
| assert "azure-identity is required" in str(exc_info.value) |
There was a problem hiding this comment.
Import-error test silently passes when azure-identity is already installed
The test patches builtins.__import__, but Python does not call __import__ for modules that are already present in sys.modules. If azure-identity is installed in the CI environment (or was loaded by an earlier test), azure.identity will already be cached in sys.modules, the mock is never invoked, no ImportError is raised, and pytest.raises(ImportError) fails the test — or, worse, the assertion silently passes because the exception was never raised and the pytest.raises block exited cleanly with a different path.
The reliable way to simulate a missing module is to set the key to None in sys.modules via patch.dict, which is how Python itself marks an import as explicitly unavailable:
def test_generate_azure_ad_redis_token_import_error():
from litellm._redis import _generate_azure_ad_redis_token
with patch.dict("sys.modules", {"azure.identity": None, "azure": None}):
with pytest.raises(ImportError) as exc_info:
_generate_azure_ad_redis_token()
assert "azure-identity is required" in str(exc_info.value)The same issue exists in test_generate_gcp_iam_access_token_import_error (line 2702), which patches builtins.__import__ in the same fragile way.
| try: | ||
| access_token = _generate_azure_ad_redis_token( | ||
| azure_client_id=_az_client_id, | ||
| azure_tenant_id=_az_tenant_id, | ||
| azure_client_secret=_az_client_secret, | ||
| ) | ||
| cluster_kwargs["password"] = access_token |
There was a problem hiding this comment.
Azure AD token generated once for async cluster — will expire
For async Redis cluster, the Azure AD token is fetched a single time at client-creation and stored as a static password. Azure AD tokens typically have a short TTL (around 1 hour). Once the token expires, every cluster operation will raise an authentication error until the entire RedisCluster object is reconstructed. There is no automatic re-fetch path here.
This is a meaningful divergence from the sync path, where redis_connect_func is invoked on every new connection and therefore obtains a fresh token automatically on reconnects.
Note that the GCP async-cluster path (line 615) has the exact same limitation. A minimal mitigation would be to document this constraint prominently and/or integrate a periodic token-refresh mechanism (e.g., a background task that recreates the cluster client before expiry) so long-running deployments don't silently break.
| _azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret( | ||
| "REDIS_AZURE_AD_TOKEN" | ||
| ) | ||
|
|
||
| if ( | ||
| _azure_redis_ad_token is not None | ||
| and str(_azure_redis_ad_token).lower() == "true" | ||
| and _gcp_service_account is not None | ||
| ): | ||
| verbose_logger.warning( | ||
| "Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. " | ||
| "Using GCP IAM. Remove one to avoid misconfiguration." | ||
| ) | ||
| # Clean up Azure-specific kwargs even though we're not using Azure AD | ||
| redis_kwargs.pop("azure_redis_ad_token", None) | ||
| redis_kwargs.pop("azure_client_id", None) | ||
| redis_kwargs.pop("azure_tenant_id", None) | ||
| redis_kwargs.pop("azure_client_secret", None) |
There was a problem hiding this comment.
Azure custom kwargs leak into redis.Redis() when azure_redis_ad_token is not "true"
The cleanup of azure_redis_ad_token, azure_client_id, azure_tenant_id, and azure_client_secret from redis_kwargs only happens inside the two branches that are gated on str(_azure_redis_ad_token).lower() == "true". If someone sets REDIS_AZURE_AD_TOKEN to any value other than "true" (e.g., "false"), _azure_redis_ad_token is non-None, neither branch executes, and all four Azure-specific keys remain in redis_kwargs. They then get forwarded to redis.Redis(**redis_kwargs), which raises a TypeError because Redis doesn't accept these custom parameters.
A fix would be to always remove the Azure-specific keys before returning from _get_redis_client_logic, regardless of which branch is taken — similar to popping the GCP keys at the end of the GCP block.
20de80b to
8c1273f
Compare
| self.send_command("AUTH", access_token, check_health=False) | ||
| auth_response = self.read_response() | ||
|
|
||
| if str_if_bytes(auth_response) != "OK": | ||
| raise AuthenticationError("Azure AD authentication failed for Redis") | ||
|
|
||
| return ad_connect | ||
|
|
||
|
|
||
| def get_redis_url_from_environment(): | ||
| if "REDIS_URL" in os.environ: | ||
| return os.environ["REDIS_URL"] |
There was a problem hiding this comment.
Inconsistent empty-username handling between sync and async paths
In the sync ad_connect closure, username defaults to "" when REDIS_USERNAME is not set, and AUTH "" <token> is always sent. In the async cluster path (around line 512–514), the username is only set if non-empty:
# async path (cluster)
_username = os.environ.get("REDIS_USERNAME", "")
if _username:
cluster_kwargs["username"] = _usernameAzure Cache for Redis with AAD authentication requires the username to be the principal's Object ID. Sending an empty string (AUTH "" <token>) is not valid for most ACL-configured Azure Redis instances and will result in a hard AuthenticationError that the AuthenticationWrongNumberOfArgsError fallback will not catch (since the argument count is correct; it's the value that's wrong).
The sync path should mirror the async path and only include the username when it's non-empty:
username = os.environ.get("REDIS_USERNAME", "")
if username:
auth_args = (username, access_token)
else:
auth_args = (access_token,)
self.send_command("AUTH", *auth_args, check_health=False)| _client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET") | ||
|
|
||
| if _client_id and _tenant_id and _client_secret: | ||
| credential = ClientSecretCredential( | ||
| client_id=_client_id, | ||
| tenant_id=_tenant_id, | ||
| client_secret=_client_secret, | ||
| ) | ||
| elif _client_id: | ||
| credential = ManagedIdentityCredential(client_id=_client_id) | ||
| else: | ||
| credential = DefaultAzureCredential() | ||
|
|
||
| token = credential.get_token(AZURE_REDIS_SCOPE) | ||
| return token.token | ||
|
|
||
|
|
||
| def create_azure_ad_redis_connect_func( | ||
| azure_client_id: Optional[str] = None, | ||
| azure_tenant_id: Optional[str] = None, | ||
| azure_client_secret: Optional[str] = None, | ||
| ) -> Callable: | ||
| """ | ||
| Creates a custom Redis connection function for Azure AD authentication. | ||
|
|
||
| Used for sync Redis clients. Generates a fresh Azure AD token on each | ||
| connection/reconnection, ensuring token refresh is handled automatically. | ||
|
|
||
| Args: | ||
| azure_client_id: Optional Azure client ID | ||
| azure_tenant_id: Optional Azure tenant ID | ||
| azure_client_secret: Optional Azure client secret |
There was a problem hiding this comment.
New credential object created on every Redis connection
_generate_azure_ad_redis_token instantiates a brand-new ClientSecretCredential, ManagedIdentityCredential, or DefaultAzureCredential on every call. Because this function is invoked inside ad_connect (which runs on every Redis connection and reconnection in the sync path), a fresh credential object and its internal HTTP client/session pool is created repeatedly throughout the application's lifetime.
Azure SDK credentials are designed to be long-lived — they cache tokens internally and handle expiry/refresh transparently. Recreating them per-connection means:
- Token caching is bypassed (a new token request is sent to Azure AD on each reconnect)
- Each
ClientSecretCredential/ManagedIdentityCredentialallocates its own HTTP connection pool that is never released, leading to resource exhaustion under load
The credential should be created once inside create_azure_ad_redis_connect_func (captured by the closure) and reused across connections. This is analogous to how service_account is captured in the outer scope of create_gcp_iam_redis_connect_func. The _generate_azure_ad_redis_token helper can still be called with the pre-built credential, or get_token can be called directly on it inside ad_connect — either way the SDK handles caching and silent renewal of the underlying token.
| redis_connect_func, "_azure_redis_ad_token" | ||
| ): | ||
| _az_client_id = getattr(redis_connect_func, "_azure_client_id", None) | ||
| _az_tenant_id = getattr(redis_connect_func, "_azure_tenant_id", None) | ||
| _az_client_secret = getattr( | ||
| redis_connect_func, "_azure_client_secret", None | ||
| ) | ||
|
|
||
| verbose_logger.debug("Generating Azure AD token for async Redis cluster") | ||
| try: | ||
| access_token = _generate_azure_ad_redis_token( | ||
| azure_client_id=_az_client_id, | ||
| azure_tenant_id=_az_tenant_id, | ||
| azure_client_secret=_az_client_secret, | ||
| ) | ||
| cluster_kwargs["password"] = access_token | ||
| # Set username if available | ||
| _username = os.environ.get("REDIS_USERNAME", "") | ||
| if _username: |
There was a problem hiding this comment.
Token expiry not documented for connection pool path
Like the async cluster path (which has an explicit NOTE warning about token TTL), the connection pool path fetches a one-shot Azure AD token at pool-creation time and sets it as a static password. Azure AD tokens typically expire in ~1 hour. Once the token expires, every connection checkout from the pool will fail authentication, and the pool itself has no refresh mechanism.
This is the same fundamental limitation called out in the cluster block, but there is no equivalent warning here. Consider adding a comment to make the constraint visible to future maintainers:
# Handle Azure AD / GCP IAM auth for connection pool — async pools don't
# support redis_connect_func, so resolve the token now and set as password.
# NOTE: The token is fetched once at pool creation. Azure AD tokens typically
# expire in ~1 hour; once expired, all connections from this pool will fail
# with an auth error. The pool (and the LiteLLM Redis client that owns it)
# must be recreated to obtain a fresh token.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)| redis_kwargs["redis_connect_func"]._azure_tenant_id = _azure_tenant_id | ||
| redis_kwargs["redis_connect_func"]._azure_client_secret = _azure_client_secret | ||
|
|
There was a problem hiding this comment.
Azure credential value stored as plaintext attribute on function object
The _azure_client_secret value is stored as a plain attribute on the redis_connect_func closure object. Unlike the GCP path, which only attaches the service account name (not a credential value), this exposes a raw auth value that could be inadvertently captured in debug logs, stack traces, or any introspection of redis_kwargs (e.g. repr(redis_connect_func)).
The credential object built inside _build_azure_credential already encapsulates this value internally and handles all token fetching. A cleaner approach for the async cluster path would be to attach the pre-built credential object to redis_connect_func rather than the raw credential components, so the raw value never needs to be re-extracted outside the closure.
| def test_redis_client_logic_azure_ad_auth(): | ||
| """Test that _get_redis_client_logic sets up Azure AD auth when REDIS_AZURE_AD_TOKEN=true.""" | ||
| from litellm._redis import _get_redis_client_logic | ||
|
|
||
| redis_kwargs = _get_redis_client_logic( | ||
| host="myredis.redis.cache.windows.net", | ||
| port="6380", | ||
| azure_redis_ad_token="true", | ||
| ssl=True, | ||
| ) | ||
|
|
||
| # Should have redis_connect_func set | ||
| assert "redis_connect_func" in redis_kwargs | ||
| assert hasattr(redis_kwargs["redis_connect_func"], "_azure_redis_ad_token") | ||
| assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True | ||
|
|
||
| # Azure-specific kwargs should be removed | ||
| assert "azure_redis_ad_token" not in redis_kwargs | ||
| assert "azure_client_id" not in redis_kwargs |
There was a problem hiding this comment.
Test will fail in CI without azure-identity installed
This test calls _get_redis_client_logic with the Azure AD flag enabled, which immediately invokes create_azure_ad_redis_connect_func → _build_azure_credential → from azure.identity import .... Because the import is eager (not deferred inside the connect closure), an ImportError is raised at test setup time if the optional azure-identity package is not present.
All other Azure tests in this PR correctly use patch.dict("sys.modules", ...) for isolation, but this test does not. With no credentials provided, DefaultAzureCredential() is instantiated during _get_redis_client_logic, requiring the actual package to be installed.
The fix mirrors the pattern used by test_generate_azure_ad_redis_token: mock azure.identity via patch.dict before calling _get_redis_client_logic, so the test is self-contained and does not depend on the optional package being present in CI.
|
closing in favour of #27556 |
Address review issues on PR #21764 cherry-pick: - Add AzureADCredentialProvider that wraps the live azure-identity credential, so async cluster, async standard, and connection-pool paths fetch tokens via the SDK's internal cache + silent refresh on every connection — instead of baking a single point-in-time token as the password (which would expire ~1h after pool creation and break all subsequent reconnects). - Stop attaching raw client_id / tenant_id / client_secret to the redis_connect_func object. The credential closure already holds them; exposing them as function attributes risked leaks via inspection or logging. Async paths now read the already-built credential object via `_azure_credential` instead. - Connection-pool path now picks up REDIS_USERNAME for ACL-configured Azure Redis instances, matching the cluster + async paths. - Mock azure.identity via sys.modules in test_redis_client_logic_ azure_ad_auth so the test no longer requires azure-identity to be installed in the CI environment. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
4 participants
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
📖 Documentation
✅ Test
Changes
azure-identitylibrary to extract fresh OAuth access tokens mapped specifically forhttps://redis.azure.com/.defaultscopes.azure_redis_ad_token: true) mapped overazure_client_idspecifications.E2E Verification Details
azure-identitylibrary to return our local Redis password as a fake AD token.AUTH <token>command.x-litellm-cache-keyheader.