diff --git a/litellm/_redis.py b/litellm/_redis.py index 2bf32d71b21..69cd536b833 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"] @@ -266,7 +289,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) @@ -428,23 +451,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 15907190998..4ca386b6d83 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -1,16 +1,20 @@ +import json +import os +from unittest.mock import MagicMock, call, patch + +import pytest +import redis +import redis.asyncio as async_redis + from litellm._redis import ( - get_redis_url_from_environment, + GCPIAMCredentialProvider, _get_redis_cluster_kwargs, get_redis_async_client, get_redis_client, get_redis_connection_pool, + get_redis_url_from_environment, ) -import json -import os -import pytest -from unittest.mock import MagicMock, patch -import redis -import redis.asyncio as async_redis + def test_get_redis_url_from_environment_single_url(monkeypatch): """Test when REDIS_URL is directly provided""" @@ -23,6 +27,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 @@ -39,6 +44,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 @@ -55,6 +61,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 @@ -69,6 +76,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 @@ -85,6 +93,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 @@ -100,6 +109,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 @@ -110,9 +120,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""" @@ -124,57 +138,66 @@ 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" + @patch("litellm._redis.init_redis_cluster") def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch): @@ -194,6 +217,7 @@ def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch): "startup_nodes" in call_kwargs ), "startup_nodes must be forwarded to init_redis_cluster" + @patch("litellm._redis.async_redis.RedisCluster") def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch): """ @@ -207,12 +231,18 @@ def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch): mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" - assert len(call_kwargs["startup_nodes"]) == 1, "should forward exactly 1 cluster node" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to async RedisCluster" + assert ( + len(call_kwargs["startup_nodes"]) == 1 + ), "should forward exactly 1 cluster node" @patch("litellm._redis.async_redis.RedisCluster") -def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, monkeypatch): +def test_async_client_prefers_cluster_over_url_via_env_var( + mock_cluster_cls, monkeypatch +): """ Test get_redis_async_client returns async RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -227,10 +257,15 @@ def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, mon mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to async RedisCluster" + @patch("litellm._redis.init_redis_cluster") -def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, monkeypatch): +def test_sync_client_prefers_cluster_over_url_via_env_var( + mock_init_cluster, monkeypatch +): """ Test get_redis_client returns RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -246,11 +281,16 @@ def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, mon mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to init_redis_cluster" + assert ( + "startup_nodes" in call_kwargs + ), "startup_nodes must be forwarded to init_redis_cluster" assert len(call_kwargs["startup_nodes"]) == 1 + @patch("litellm._redis.init_redis_cluster") -def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_cluster, monkeypatch): +def test_sync_client_preserves_password_for_cluster_when_url_also_set( + mock_init_cluster, monkeypatch +): """ Test _get_redis_client_logic does not strip password from redis_kwargs when startup_nodes is present even if REDIS_URL is also set. @@ -264,7 +304,9 @@ def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_ mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert "password" in call_kwargs, "password must not be stripped when routing to cluster" + assert ( + "password" in call_kwargs + ), "password must not be stripped when routing to cluster" assert call_kwargs["password"] == "secret" @@ -287,3 +329,70 @@ def test_sync_client_url_used_when_no_cluster(mock_from_url, monkeypatch): get_redis_client() mock_from_url.assert_called_once() + + +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)"