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 35

fix(redis): regenerate GCP IAM token per connection for async cluster clients

Browse source 
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.
This commit is contained in:
Ishaan Jaffer 2026-03-23 09:24:05 -07:00
parent 9e98a7d7b3
commit 67b752ca8f
2 changed files with 185 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/_redis.py
View file

@ -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}"

197
tests/test_litellm/test_redis.py
View file

@ -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)"