diff --git a/litellm/_redis.py b/litellm/_redis.py index 0acc01fa14f..f3f3c4424de 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -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) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 896ca2de399..3aa4bc58f13 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -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