scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #37740 from BerriAI/litellm_redis_url_pool_credential_provider

Browse source 
fix(redis): apply Azure AD and GCP IAM auth to every async client path
This commit is contained in:
Mateo Wang 2026-08-20 17:24:10 -07:00 • committed by GitHub
commit 987478abe4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 201 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

91
litellm/_redis.py
View file

@ -17,6 +17,7 @@ from typing import Final
import redis
import redis.asyncio as async_redis
from redis.credentials import CredentialProvider
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
@ -134,6 +135,7 @@ def _get_redis_cluster_kwargs(client=None):
"ssl_check_hostname",
"ssl_ca_certs",
"redis_connect_func", # Needed for sync clusters and IAM detection
"credential_provider",
"gcp_service_account",
"gcp_ssl_ca_certs",
"azure_redis_ad_token",
@ -549,14 +551,22 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
return sentinel.master_for(service_name, **connection_kwargs)
def _sentinel_auth_kwargs(connection_kwargs: dict, sentinel_password: str | None) -> dict:
"""The Sentinel monitors are separate servers that authenticate with their own password, so the
data node's credential provider never belongs on them: leaving it there makes redis-py send the
data node's token to a monitor, which fails whether the monitor is unauthenticated or has its
own password."""
kept: Final = ((k, v) for k, v in connection_kwargs.items() if k != "credential_provider")
return dict(kept, password=sentinel_password)
def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = dict(connection_kwargs)
sentinel_kwargs["password"] = sentinel_password
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
@ -574,6 +584,36 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
return sentinel.master_for(service_name, **connection_kwargs)
def _async_credential_provider(redis_connect_func: object | None) -> CredentialProvider | None:
"""The Azure AD and GCP IAM connect funcs run their AUTH exchange with the blocking client
API, so on an async connection their ``send_command``/``read_response`` calls return
coroutines nobody awaits and every connect fails. Async paths authenticate through a
``CredentialProvider`` instead, which redis-py consults per connection so the token stays
fresh. Any other ``redis_connect_func`` is left where it is, since redis-py awaits it
itself when it is a coroutine function."""
gcp_service_account: Final = getattr(redis_connect_func, "_gcp_service_account", None)
if gcp_service_account is not None:
return GCPIAMCredentialProvider(gcp_service_account)
azure_credential: Final = getattr(redis_connect_func, "_azure_credential", None)
if azure_credential is not None:
return AzureADCredentialProvider(azure_credential, username=os.environ.get("REDIS_USERNAME") or None)
return None
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
which supersedes any static username or password redis-py would otherwise reject it with."""
credential_provider: Final = _async_credential_provider(redis_kwargs.get("redis_connect_func"))
if credential_provider is None:
return redis_kwargs
superseded: Final = frozenset({"redis_connect_func", "username", "password"})
kept: Final = ((k, v) for k, v in redis_kwargs.items() if k not in superseded)
return dict(kept, credential_provider=credential_provider) # mutable-ok: the branches below mutate these kwargs
def get_redis_client(**env_overrides):
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
@ -600,7 +640,7 @@ def get_redis_async_client(
connection_pool: async_redis.BlockingConnectionPool | None = None,
**env_overrides,
) -> async_redis.Redis | async_redis.RedisCluster:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
if "startup_nodes" in redis_kwargs:
from redis.cluster import ClusterNode
@ -611,28 +651,12 @@ def get_redis_async_client(
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
# Handle GCP IAM authentication for async clusters
redis_connect_func = cluster_kwargs.pop("redis_connect_func", None)
# Use a CredentialProvider so the IAM token is regenerated on every new
# connection — mirrors the sync path where redis_connect_func is invoked
# per connection. Without this, the token would expire after ~1 hour.
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
# Handle Azure AD authentication for async clusters via CredentialProvider
# so the credential's internal cache + silent refresh runs per connection
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: Final[list[ClusterNode]] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
cluster_kwargs.pop("startup_nodes", None)
cluster_kwargs.pop("redis_connect_func", None)
# Default to a periodic health check + TCP keepalive so a connection silently dropped
# by a cluster restart (e.g. ElastiCache Serverless maintenance) is revalidated and
@ -667,19 +691,6 @@ def get_redis_async_client(
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
return _init_async_redis_sentinel(redis_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
# Redis client. The async client doesn't support redis_connect_func, but it
# does honour credential_provider — which is called per connection, so the
# underlying SDK can refresh tokens silently before they expire.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
if connection_pool is not None:
@ -693,7 +704,7 @@ def get_redis_async_client(
def get_redis_connection_pool(
**env_overrides,
) -> async_redis.BlockingConnectionPool | None:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "startup_nodes" in redis_kwargs:
@ -714,18 +725,6 @@ def get_redis_connection_pool(
)
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
# connections re-fetch tokens via the SDK's internal cache + silent refresh
# rather than reusing a single token captured at pool creation.
redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
if redis_kwargs.pop("ssl", None):
redis_kwargs["connection_class"] = async_redis.SSLConnection
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)

156
tests/test_litellm/test_redis.py
View file

@ -1,4 +1,5 @@
import json
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
@ -14,6 +15,7 @@ from litellm._redis import (
)
from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
GCPIAMCredentialProvider,
_token_cache,
)
@ -910,3 +912,157 @@ def test_url_allowlist_always_carries_socket_timeouts():
allowed = _get_redis_url_kwargs()
assert "socket_timeout" in allowed
assert "socket_connect_timeout" in allowed
AZURE_AD_CONNECT_FUNC = {"_azure_credential": object()}
GCP_IAM_CONNECT_FUNC = {"_gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}
@pytest.mark.parametrize(
"markers, provider_cls",
[
(AZURE_AD_CONNECT_FUNC, AzureADCredentialProvider),
(GCP_IAM_CONNECT_FUNC, GCPIAMCredentialProvider),
],
ids=["azure_ad", "gcp_iam"],
)
def test_async_url_client_authenticates_through_credential_provider(markers, provider_cls):
"""A REDIS_URL config with Azure AD or GCP IAM must still reach the server with a credential.
The url branch forwards redis_connect_func straight to the async connection, which runs
its AUTH exchange with the blocking client API and dies, so the branch has to hand the
connection a CredentialProvider instead.
"""
redis_kwargs = {
"url": "rediss://redis-host:6380",
"redis_connect_func": SimpleNamespace(**markers),
}
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
client = get_redis_async_client()
connection_kwargs = client.connection_pool.connection_kwargs
assert isinstance(connection_kwargs.get("credential_provider"), provider_cls)
assert "redis_connect_func" not in connection_kwargs
@pytest.mark.parametrize(
"markers, provider_cls",
[
(AZURE_AD_CONNECT_FUNC, AzureADCredentialProvider),
(GCP_IAM_CONNECT_FUNC, GCPIAMCredentialProvider),
],
ids=["azure_ad", "gcp_iam"],
)
def test_async_url_connection_pool_authenticates_through_credential_provider(markers, provider_cls):
"""Same for the pool-based path: every connection the pool hands out needs the provider."""
redis_kwargs = {
"url": "rediss://redis-host:6380",
"redis_connect_func": SimpleNamespace(**markers),
}
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
pool = get_redis_connection_pool()
assert isinstance(pool.connection_kwargs.get("credential_provider"), provider_cls)
assert "redis_connect_func" not in pool.connection_kwargs
def test_async_url_client_drops_username_alongside_credential_provider():
"""redis-py refuses a connection given both a username and a credential_provider, and
AzureADCredentialProvider already carries REDIS_USERNAME, so the username must be dropped.
"""
redis_kwargs = {
"url": "rediss://redis-host:6380",
"username": "redis-user",
"redis_connect_func": SimpleNamespace(**AZURE_AD_CONNECT_FUNC),
}
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
client = get_redis_async_client()
pool = client.connection_pool
assert "username" not in pool.connection_kwargs
pool.connection_class(**pool.connection_kwargs)
@pytest.mark.parametrize("build_pool", [False, True], ids=["client", "pool"])
def test_async_url_keeps_a_coroutine_connect_func(build_pool):
"""redis-py awaits a coroutine redis_connect_func on an async connection, so one we cannot
turn into a credential provider has to be left where it is rather than dropped.
"""
async def connect(connection):
return None
redis_kwargs = {
"url": "rediss://redis-host:6380",
"redis_connect_func": connect,
}
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
pool = get_redis_connection_pool() if build_pool else get_redis_async_client().connection_pool
assert pool.connection_kwargs["redis_connect_func"] is connect
assert "credential_provider" not in pool.connection_kwargs
def test_async_cluster_drops_a_connect_func_it_cannot_pass_on():
"""redis-py's async RedisCluster has no redis_connect_func parameter, so a connect func that
is not translated into a credential provider has to be dropped rather than forwarded.
"""
async def connect(connection):
return None
redis_kwargs = {
"startup_nodes": [{"host": "cluster-node", "port": 6379}],
"redis_connect_func": connect,
}
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
client = get_redis_async_client()
assert isinstance(client, async_redis.RedisCluster)
@pytest.mark.parametrize(
"markers, provider_cls",
[
(AZURE_AD_CONNECT_FUNC, AzureADCredentialProvider),
(GCP_IAM_CONNECT_FUNC, GCPIAMCredentialProvider),
],
ids=["azure_ad", "gcp_iam"],
)
@pytest.mark.parametrize(
"sentinel_password",
[None, "sentinel-secret"],
ids=["unauthenticated_monitors", "password_protected_monitors"],
)
def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers, provider_cls, sentinel_password):
"""The Sentinel monitors are separate servers with their own password, so the data node's token
never belongs on them: redis-py refuses it next to a Sentinel password, and sends it to an
unauthenticated monitor as an AUTH the monitor rejects.
"""
redis_kwargs = {
"sentinel_nodes": [("sentinel-1", 26379)],
"sentinel_password": sentinel_password,
"service_name": "mymaster",
"redis_connect_func": SimpleNamespace(**markers),
}
with patch("litellm._redis.async_redis.Sentinel") as mock_sentinel_cls:
with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs):
get_redis_async_client()
sentinel_kwargs = mock_sentinel_cls.call_args[1]["sentinel_kwargs"]
assert sentinel_kwargs["password"] == sentinel_password
assert "credential_provider" not in sentinel_kwargs
monitor_connection = async_redis.Connection(host="sentinel-1", port=26379, **sentinel_kwargs)
assert monitor_connection.credential_provider is None
assert bool(monitor_connection.username or monitor_connection.password) is bool(sentinel_password)
master_kwargs = mock_sentinel_cls.return_value.master_for.call_args[1]
assert isinstance(master_kwargs["credential_provider"], provider_cls)
assert "password" not in master_kwargs