Merge pull request #37740 from BerriAI/litellm_redis_url_pool_credential_provider

fix(redis): apply Azure AD and GCP IAM auth to every async client path

Author: mateo-berri

litellm/_redis.py

@@ -17,6 +17,7 @@
 
 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)

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.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