From 193089d94ccc47fd6441fce03e5fa9775ffb05f9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 09:24:05 -0700 Subject: [PATCH] fix(redis): regenerate GCP IAM token per connection for async cluster clients Async RedisCluster was generating the IAM token once at startup and storing it as a static password. After the 1-hour GCP token TTL, any new connection (including to newly-discovered cluster nodes) would fail to authenticate. Fix: introduce GCPIAMCredentialProvider that implements redis-py's CredentialProvider protocol. It calls _generate_gcp_iam_access_token() on every new connection, matching what the sync redis_connect_func already does. async_redis.RedisCluster accepts a credential_provider kwarg which is invoked per-connection. --- litellm/_redis.py | 49 ++++++---- tests/test_litellm/test_redis.py | 155 ++++++++++++++++++++++++------- 2 files changed, 156 insertions(+), 48 deletions(-) diff --git a/litellm/_redis.py b/litellm/_redis.py index b754c1f4330..472b3c2c80d 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -7,12 +7,13 @@ # # Thank you users! We ❤️ you! - Krrish & Ishaan +import asyncio import inspect import json # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os -from typing import Callable, List, Optional, Union +from typing import Callable, List, Optional, Tuple, Union import redis # type: ignore import redis.asyncio as async_redis # type: ignore @@ -178,6 +179,28 @@ def create_gcp_iam_redis_connect_func( return iam_connect +class GCPIAMCredentialProvider: + """ + redis.credentials.CredentialProvider implementation that generates a fresh GCP IAM + token on every new connection. This fixes the 1-hour token expiry issue for async + Redis cluster clients, which previously generated the token once at startup and + cached it as a static password. + """ + + def __init__(self, gcp_service_account: str) -> None: + self._gcp_service_account = gcp_service_account + + def get_credentials(self) -> Tuple[str]: + token = _generate_gcp_iam_access_token(self._gcp_service_account) + return (token,) + + async def get_credentials_async(self) -> Tuple[str]: + token = await asyncio.to_thread( + _generate_gcp_iam_access_token, self._gcp_service_account + ) + return (token,) + + def get_redis_url_from_environment(): if "REDIS_URL" in os.environ: return os.environ["REDIS_URL"] @@ -262,7 +285,7 @@ def _get_redis_client_logic(**env_overrides): service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs ) # Store GCP service account in redis_connect_func for async cluster access - redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account + redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account # type: ignore[attr-defined] # Remove GCP-specific kwargs that shouldn't be passed to Redis client redis_kwargs.pop("gcp_service_account", None) @@ -434,23 +457,15 @@ def get_redis_async_client( f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" ) - # If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password + # If GCP IAM is configured (indicated by redis_connect_func), attach a + # credential_provider that regenerates the token on every new connection. + # This mirrors the sync behaviour where redis_connect_func is called per + # connection, and avoids the 1-hour token expiry bug where the old code + # generated the token once at startup and set it as a static password. if redis_connect_func and gcp_service_account: - verbose_logger.debug( - "DEBUG: Generating IAM token for service account (value not logged for security reasons)" + cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider( + gcp_service_account ) - try: - # Generate IAM access token using the helper function - access_token = _generate_gcp_iam_access_token(gcp_service_account) - cluster_kwargs["password"] = access_token - verbose_logger.debug( - "DEBUG: Successfully generated GCP IAM access token for async Redis cluster" - ) - except Exception as e: - verbose_logger.error(f"Failed to generate GCP IAM access token: {e}") - from redis.exceptions import AuthenticationError - - raise AuthenticationError("Failed to generate GCP IAM access token") else: verbose_logger.debug( f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 4709faea4bc..da0d350da2f 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,9 +1,17 @@ -from litellm._redis import get_redis_url_from_environment, _get_redis_cluster_kwargs, get_redis_async_client import os +from unittest.mock import MagicMock, call, patch + import pytest -from unittest.mock import MagicMock, patch import redis.asyncio as async_redis +from litellm._redis import ( + GCPIAMCredentialProvider, + _get_redis_cluster_kwargs, + get_redis_async_client, + get_redis_url_from_environment, +) + + def test_get_redis_url_from_environment_single_url(monkeypatch): """Test when REDIS_URL is directly provided""" # Set the environment variable @@ -15,6 +23,7 @@ def test_get_redis_url_from_environment_single_url(monkeypatch): # Assert that the returned URL matches the expected value assert redis_url == "redis://redis-server:6379/0" + def test_get_redis_url_from_environment_host_port(monkeypatch): """Test when REDIS_HOST and REDIS_PORT are provided""" # Set the environment variables @@ -31,6 +40,7 @@ def test_get_redis_url_from_environment_host_port(monkeypatch): # Assert that the returned URL matches the expected value assert redis_url == "redis://redis-server:6379" + def test_get_redis_url_from_environment_with_ssl(monkeypatch): """Test when SSL is enabled""" # Set the environment variables @@ -47,6 +57,7 @@ def test_get_redis_url_from_environment_with_ssl(monkeypatch): # Assert that the returned URL uses rediss:// protocol assert redis_url == "rediss://redis-server:6379" + def test_get_redis_url_from_environment_with_username_password(monkeypatch): """Test when username and password are provided""" # Set the environment variables @@ -61,6 +72,7 @@ def test_get_redis_url_from_environment_with_username_password(monkeypatch): # Assert that the returned URL includes username:password@ assert redis_url == "redis://user:password@redis-server:6379" + def test_get_redis_url_from_environment_with_password_only(monkeypatch): """Test when only password is provided""" # Set the environment variables @@ -77,6 +89,7 @@ def test_get_redis_url_from_environment_with_password_only(monkeypatch): # Assert that the returned URL includes :password@ assert redis_url == "redis://password@redis-server:6379" + def test_get_redis_url_from_environment_with_all_options(monkeypatch): """Test when all options are provided""" # Set the environment variables @@ -92,6 +105,7 @@ def test_get_redis_url_from_environment_with_all_options(monkeypatch): # Assert that the returned URL includes all components assert redis_url == "rediss://user:password@redis-server:6379" + def test_get_redis_url_from_environment_missing_host_port(monkeypatch): """Test error when required variables are missing""" # Make sure these environment variables don't exist @@ -102,9 +116,13 @@ def test_get_redis_url_from_environment_missing_host_port(monkeypatch): # Call the function and expect a ValueError with pytest.raises(ValueError) as excinfo: get_redis_url_from_environment() - + # Check the error message - assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) + assert ( + "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" + in str(excinfo.value) + ) + def test_get_redis_url_from_environment_missing_port(monkeypatch): """Test error when only REDIS_HOST is provided but REDIS_PORT is missing""" @@ -116,54 +134,129 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): # Call the function and expect a ValueError with pytest.raises(ValueError) as excinfo: get_redis_url_from_environment() - + # Check the error message - assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) + assert ( + "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" + in str(excinfo.value) + ) + def test_max_connections_in_cluster_kwargs(): """Test that max_connections is included in Redis cluster kwargs""" kwargs = _get_redis_cluster_kwargs() - assert "max_connections" in kwargs, "max_connections should be in available Redis cluster kwargs" + assert ( + "max_connections" in kwargs + ), "max_connections should be in available Redis cluster kwargs" + def test_get_redis_async_client_with_connection_pool(): """Test that connection_pool parameter is properly passed to Redis client""" # Create a mock connection pool mock_pool = MagicMock(spec=async_redis.BlockingConnectionPool) - + # Mock the Redis client creation - with patch('litellm._redis.async_redis.Redis') as mock_redis, \ - patch('litellm._redis._get_redis_client_logic') as mock_logic: - + with patch("litellm._redis.async_redis.Redis") as mock_redis, patch( + "litellm._redis._get_redis_client_logic" + ) as mock_logic: + # Configure mock to return basic redis kwargs - mock_logic.return_value = { - "host": "localhost", - "port": 6379, - "db": 0 - } - + mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} + # Call get_redis_async_client with connection_pool get_redis_async_client(connection_pool=mock_pool) - + # Verify Redis was called with connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert "connection_pool" in call_kwargs, "connection_pool should be passed to Redis client" - assert call_kwargs["connection_pool"] == mock_pool, "connection_pool should match the provided pool" + assert ( + "connection_pool" in call_kwargs + ), "connection_pool should be passed to Redis client" + assert ( + call_kwargs["connection_pool"] == mock_pool + ), "connection_pool should match the provided pool" + def test_get_redis_async_client_without_connection_pool(): """Test that Redis client works without connection_pool parameter""" - with patch('litellm._redis.async_redis.Redis') as mock_redis, \ - patch('litellm._redis._get_redis_client_logic') as mock_logic: - + with patch("litellm._redis.async_redis.Redis") as mock_redis, patch( + "litellm._redis._get_redis_client_logic" + ) as mock_logic: + # Configure mock to return basic redis kwargs - mock_logic.return_value = { - "host": "localhost", - "port": 6379, - "db": 0 - } - + mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} + # Call get_redis_async_client without connection_pool get_redis_async_client() - + # Verify Redis was called without connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert "connection_pool" not in call_kwargs, "connection_pool should not be in kwargs when not provided" + assert ( + "connection_pool" not in call_kwargs + ), "connection_pool should not be in kwargs when not provided" + + +def test_gcp_iam_credential_provider_get_credentials(): + """GCPIAMCredentialProvider.get_credentials() returns a fresh token tuple on every call.""" + service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" + + with patch( + "litellm._redis._generate_gcp_iam_access_token", return_value="tok-1" + ) as mock_gen: + provider = GCPIAMCredentialProvider(service_account) + creds = provider.get_credentials() + + assert creds == ("tok-1",) + mock_gen.assert_called_once_with(service_account) + + +def test_gcp_iam_credential_provider_regenerates_token_on_each_call(): + """Each call to get_credentials() generates a new token (no caching).""" + service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" + tokens = ["tok-1", "tok-2", "tok-3"] + + with patch( + "litellm._redis._generate_gcp_iam_access_token", side_effect=tokens + ) as mock_gen: + provider = GCPIAMCredentialProvider(service_account) + results = [provider.get_credentials() for _ in range(3)] + + assert results == [("tok-1",), ("tok-2",), ("tok-3",)] + assert mock_gen.call_count == 3 + + +def test_get_redis_async_client_gcp_cluster_uses_credential_provider(): + """ + When startup_nodes + gcp_service_account are provided, the async cluster client + must be constructed with a GCPIAMCredentialProvider — not a static password. + This ensures that the 1-hour IAM token expiry does not cause auth failures. + """ + startup_nodes = [{"host": "redis-node-1", "port": 6379}] + + mock_connect_func = MagicMock() + mock_connect_func._gcp_service_account = ( + "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com" + ) + + redis_kwargs = { + "startup_nodes": startup_nodes, + "redis_connect_func": mock_connect_func, + } + + with patch("litellm._redis.async_redis.RedisCluster") as mock_cluster, patch( + "litellm._redis._get_redis_client_logic", return_value=redis_kwargs + ): + get_redis_async_client() + + assert mock_cluster.called + cluster_call_kwargs = mock_cluster.call_args[1] + + # Must use credential_provider, not a static password + assert ( + "credential_provider" in cluster_call_kwargs + ), "async GCP cluster must use credential_provider for per-connection token refresh" + assert isinstance( + cluster_call_kwargs["credential_provider"], GCPIAMCredentialProvider + ) + assert ( + "password" not in cluster_call_kwargs + ), "async GCP cluster must not use a static password (expires after 1h)"