scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_internal_user_spend_log_detail_route

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-09-10 15:13:59 -07:00
commit 9bdaa54643
11 changed files with 864 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
litellm/_redis.py
View file

@ -24,6 +24,7 @@ from redis.credentials import CredentialProvider
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
ElastiCacheIAMCredentialProvider,
GCPIAMCredentialProvider,
_generate_gcp_iam_access_token,
)
@ -38,6 +39,14 @@ from ._logging import verbose_logger
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
_AWS_IAM_KWARG_NAMES: Final = (
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
)
def _unwrapped_init_args(cls: type) -> frozenset[str]:
"""Every parameter on a single class's own ``__init__``, decorator-unwrapped.
@ -75,6 +84,7 @@ def _get_redis_kwargs():
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
*_AWS_IAM_KWARG_NAMES,
}
available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
@ -270,6 +280,42 @@ def _redis_kwargs_from_environment():
return return_dict
def _coerces_to_true(value: object | None) -> bool:
return _str_to_bool(value) if isinstance(value, str) else bool(value)
def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool:
if redis_kwargs.get("startup_nodes") is not None:
return _coerces_to_true(redis_kwargs.get("ssl"))
url: Final = redis_kwargs.get("url")
if isinstance(url, str):
return urlsplit(url).scheme.lower() == "rediss"
return _coerces_to_true(redis_kwargs.get("ssl"))
def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider:
user_name: Final = redis_kwargs.get("aws_iam_user_name")
cache_name: Final = redis_kwargs.get("aws_iam_cache_name")
region: Final = (
redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION")
)
required_settings: Final = (
("aws_iam_user_name", user_name),
("aws_iam_cache_name", cache_name),
("aws_iam_region", region),
)
missing_settings: Final = tuple(name for name, value in required_settings if not value)
if missing_settings:
raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings))
return ElastiCacheIAMCredentialProvider(
user_name=str(user_name),
cache_name=str(cache_name),
region=str(region),
is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")),
)
def create_gcp_iam_redis_connect_func(
service_account: str,
ssl_ca_certs: str | None = None,
@ -540,6 +586,7 @@ def _get_redis_client_logic(**env_overrides):
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
_aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth"))
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
@ -567,6 +614,22 @@ def _get_redis_client_logic(**env_overrides):
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
if _aws_iam_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured "
"for Redis. Using GCP IAM. Remove one to avoid misconfiguration."
)
elif _aws_iam_enabled and _azure_ad_enabled:
verbose_logger.warning(
"Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured "
"for Redis. Using Azure AD. Remove one to avoid misconfiguration."
)
elif _aws_iam_enabled:
if not _uses_tls(redis_kwargs):
raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS")
verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs)
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
@ -575,6 +638,8 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("azure_client_id", None)
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
for aws_iam_key in _AWS_IAM_KWARG_NAMES:
redis_kwargs.pop(aws_iam_key, None)
if redis_kwargs.get("credential_provider") is not None:
redis_kwargs.pop("redis_connect_func", None)

85
litellm/_redis_credential_provider.py
View file

@ -1,10 +1,17 @@
from __future__ import annotations
import asyncio
import threading
import time
from typing import Final, Protocol
from collections.abc import Callable
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import urlencode
from redis.credentials import CredentialProvider
if TYPE_CHECKING:
from botocore.credentials import Credentials
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
@ -117,6 +124,82 @@ class GCPIAMCredentialProvider(CredentialProvider):
return (token,)
_ELASTICACHE_SERVICE_NAME: Final = "elasticache"
_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache"
class ElastiCacheIAMCredentialProvider(CredentialProvider):
def __init__(
self,
user_name: str,
cache_name: str,
region: str,
is_serverless: bool = False,
credentials_resolver: Callable[[], Credentials | None] | None = None,
token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS,
) -> None:
self._user_name = user_name
self._cache_name = cache_name.lower()
self._region = region
self._is_serverless = is_serverless
self._credentials_resolver = credentials_resolver or self._resolve_credentials
self._credentials: Credentials | None = None
self._token_lifetime_seconds = token_lifetime_seconds
@staticmethod
def _resolve_credentials() -> Credentials | None:
try:
import botocore.session
except ImportError as e:
raise ImportError(
"botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
) from e
return botocore.session.get_session().get_credentials()
def _get_credentials(self) -> tuple[str, str]:
credentials: Final = self._credentials if self._credentials is not None else self._credentials_resolver()
if credentials is None:
raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication")
self._credentials = credentials
frozen_credentials: Final = credentials.get_frozen_credentials()
try:
from botocore.auth import SigV4QueryAuth
from botocore.awsrequest import AWSRequest
except ImportError as e:
raise ImportError(
"botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
) from e
query: Final = urlencode(
(
("Action", "connect"),
("User", self._user_name),
*((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()),
)
)
request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}")
SigV4QueryAuth(
frozen_credentials,
_ELASTICACHE_SERVICE_NAME,
self._region,
expires=self._token_lifetime_seconds,
).add_auth(request)
signed_url: Final = request.url
if signed_url is None:
raise RuntimeError("Unable to generate AWS ElastiCache IAM credentials")
return self._user_name, signed_url.removeprefix("https://")
def get_credentials(self) -> tuple[str, str]:
return self._get_credentials()
async def get_credentials_async(self) -> tuple[str, str]:
return await asyncio.to_thread(self._get_credentials)
class AzureADCredentialProvider(CredentialProvider):
"""
redis.credentials.CredentialProvider implementation that supplies Azure AD

7
litellm/proxy/_types.py
View file

@ -2432,6 +2432,13 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase):
)
sentinel_password: str | None = Field(None, description="password for the sentinel nodes")
service_name: str | None = Field(None, description="sentinel service name")
aws_iam_auth: bool | str | None = Field(None, description="enable AWS ElastiCache IAM authentication")
aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name")
aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name")
aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication")
aws_iam_serverless: bool | str | None = Field(
None, description="the ElastiCache cache is serverless rather than a self-designed cluster"
)
def has_connection_target(self) -> bool:
return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes))

2
litellm/proxy/management_endpoints/coordination_redis_endpoints.py
View file

@ -140,7 +140,7 @@ def _merge_over_saved(
def _validated_params(settings: Mapping[str, object]) -> CoordinationRedisParams:
"""Validate settings the way startup does: resolve env refs, then require a connection target."""
try:
params: Final = CoordinationRedisParams(**_resolve_env_refs(settings))
params: Final = CoordinationRedisParams.model_validate(_resolve_env_refs(settings))
except ValidationError as e:
invalid_fields: Final = sorted({str(error["loc"][0]) for error in e.errors() if error["loc"]})
raise HTTPException(

8
litellm/proxy/proxy_server.py
View file

@ -4957,7 +4957,9 @@ class ProxyConfig:
if not isinstance(raw_params, dict):
raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params")
coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(raw_params))
coordination_params: Final = CoordinationRedisParams.model_validate(
_resolve_coordination_redis_env_refs(raw_params)
)
if not coordination_params.has_connection_target():
raise ValueError(
"general_settings.coordination_redis needs a connection target: "
@ -9285,7 +9287,9 @@ class ProxyStartupEvent:
if persisted is None:
return None
coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(persisted))
coordination_params: Final = CoordinationRedisParams.model_validate(
_resolve_coordination_redis_env_refs(persisted)
)
if not coordination_params.has_connection_target():
verbose_proxy_logger.warning(
"coordination_redis saved in the database names no connection target; ignoring it."

45
litellm/types/management_endpoints/cache_settings_endpoints.py
View file

@ -237,4 +237,49 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [
ui_field_name="SSL Check Hostname",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_auth",
field_type="Boolean",
field_value=None,
field_description="Enable AWS ElastiCache IAM authentication",
field_default=False,
ui_field_name="AWS IAM Authentication",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_user_name",
field_type="String",
field_value=None,
field_description="AWS ElastiCache IAM user name",
field_default=None,
ui_field_name="AWS IAM User Name",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_cache_name",
field_type="String",
field_value=None,
field_description="AWS ElastiCache cache name",
field_default=None,
ui_field_name="AWS IAM Cache Name",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_region",
field_type="String",
field_value=None,
field_description="AWS region for ElastiCache IAM authentication",
field_default=None,
ui_field_name="AWS IAM Region",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_value=None,
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
redis_type=None,
),
]

37
litellm/types/management_endpoints/coordination_redis_endpoints.py
View file

@ -102,4 +102,41 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]]
ui_field_name="Service Name",
section="sentinel",
),
CoordinationRedisSettingsField(
field_name="aws_iam_auth",
field_type="Boolean",
field_description="Enable AWS ElastiCache IAM authentication",
field_default=False,
ui_field_name="AWS IAM Authentication",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_user_name",
field_type="String",
field_description="AWS ElastiCache IAM user name",
ui_field_name="AWS IAM User Name",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_cache_name",
field_type="String",
field_description="AWS ElastiCache cache name",
ui_field_name="AWS IAM Cache Name",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_region",
field_type="String",
field_description="AWS region for ElastiCache IAM authentication",
ui_field_name="AWS IAM Region",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
section="connection",
),
]

2
tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py
View file

@ -481,7 +481,7 @@ async def test_archive_source_registers_and_is_served_verbatim_in_marketplace():
assert response.action == "created"
assert response.plugin.source == _ARCHIVE_SOURCE
marketplace = json.loads((await get_marketplace()).body)
marketplace = json.loads((await get_marketplace(request=MagicMock())).body)
assert marketplace["plugins"] == [{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "version": "1.0.0"}]

357
tests/test_litellm/test_redis.py
View file

@ -10,13 +10,16 @@ from redis.credentials import CredentialProvider
import litellm
from litellm._redis import (
_AWS_IAM_KWARG_NAMES,
_async_auth_kwargs,
_coerce_redis_kwargs_types,
_get_redis_client_logic,
_get_redis_cluster_kwargs,
_get_redis_env_kwarg_mapping,
_get_redis_kwargs,
_get_redis_url_kwargs,
_pretty_print_redis_config,
_uses_tls,
get_redis_async_client,
get_redis_client,
get_redis_connection_pool,
@ -24,12 +27,14 @@ from litellm._redis import (
)
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
ElastiCacheIAMCredentialProvider,
GCPIAMCredentialProvider,
_token_cache,
)
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
from litellm.proxy._types import CoordinationRedisParams
class _StubCredentialProvider(CredentialProvider):
@ -78,6 +83,8 @@ def clean_redis_environment(monkeypatch):
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
"AWS_REGION",
"AWS_DEFAULT_REGION",
*_get_redis_env_kwarg_mapping(),
):
monkeypatch.delenv(var, raising=False)
@ -110,6 +117,29 @@ def test_credential_provider_is_not_environment_derived():
assert "credential_provider" not in mapping.values()
_AWS_IAM_SETTINGS = {
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
}
def test_aws_iam_settings_are_environment_derived():
allowed = _get_redis_kwargs()
mapping = _get_redis_env_kwarg_mapping()
assert _AWS_IAM_SETTINGS <= allowed
assert set(_AWS_IAM_KWARG_NAMES) == _AWS_IAM_SETTINGS
assert {f for f in CoordinationRedisParams.model_fields if f.startswith("aws_iam_")} == _AWS_IAM_SETTINGS
assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth"
assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name"
assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name"
assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region"
assert mapping["REDIS_AWS_IAM_SERVERLESS"] == "aws_iam_serverless"
def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment):
provider = _StubCredentialProvider()
@ -300,6 +330,333 @@ def test_gcp_kwargs_never_survive_client_logic(clean_redis_environment, override
assert "gcp_ssl_ca_certs" not in redis_kwargs
def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch):
monkeypatch.setenv("REDIS_AWS_IAM_AUTH", "true")
monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user")
monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com")
monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1")
monkeypatch.setenv("REDIS_AWS_IAM_SERVERLESS", "1")
monkeypatch.setenv("REDIS_SSL", "1")
redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is True
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize(
"transport",
[
pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "off"}, id="host_ssl_off_string"),
pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}]},
id="cluster_without_ssl",
),
pytest.param(
{
"url": "rediss://cache.example.com:6379",
"startup_nodes": [{"host": "cache.example.com", "port": 6379}],
},
id="cluster_without_ssl_ignores_url_scheme",
),
pytest.param(
{
"sentinel_nodes": [("sentinel.example.com", 26379)],
"service_name": "cache",
},
id="sentinel_without_ssl",
),
],
)
def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, transport):
with pytest.raises(ValueError, match="requires TLS"):
_get_redis_client_logic(
**transport,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
@pytest.mark.parametrize(
"transport",
[
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "True"}, id="host_ssl_true_capitalized"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}], "ssl": "1"},
id="cluster_ssl_one_string",
),
pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"),
pytest.param(
{
"startup_nodes": [{"host": "cache.example.com", "port": 6379}],
"ssl": True,
},
id="cluster",
),
pytest.param(
{
"sentinel_nodes": [("sentinel.example.com", 26379)],
"service_name": "cache",
"ssl": True,
},
id="sentinel",
),
],
)
def test_aws_iam_auth_accepts_tls_connections(clean_redis_environment, transport):
redis_kwargs = _get_redis_client_logic(
**transport,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
@pytest.mark.parametrize("ssl", ["true", "True", "TRUE", "1", "yes", "YES", "false", "0", "no", "off", "", "maybe"])
def test_tls_detection_agrees_with_the_ssl_kwarg_coercion(ssl):
assert _uses_tls({"ssl": ssl}) is _coerce_redis_kwargs_types({"ssl": ssl})["ssl"]
def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
url="rediss://url-user:url-pass@cache.example.com:6380",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
username="static-user",
password="static-password",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
assert redis_kwargs["url"] == "rediss://cache.example.com:6380"
assert "username" not in redis_kwargs
assert "password" not in redis_kwargs
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"])
def test_aws_iam_missing_setting_fails_closed(clean_redis_environment, missing):
settings = {
"aws_iam_auth": True,
"aws_iam_user_name": "iam-user",
"aws_iam_cache_name": "cache.example.com",
"aws_iam_region": "us-east-1",
"ssl": True,
}
settings[missing] = None
with pytest.raises(ValueError, match=missing):
_get_redis_client_logic(host="cache.example.com", port=6379, **settings)
@pytest.mark.parametrize("region_var", ["AWS_REGION", "AWS_DEFAULT_REGION"])
def test_aws_iam_region_falls_back_to_environment(clean_redis_environment, monkeypatch, region_var):
monkeypatch.setenv(region_var, "sa-east-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._region == "sa-east-1"
def test_aws_iam_region_prefers_aws_region_over_default_region(clean_redis_environment, monkeypatch):
monkeypatch.setenv("AWS_REGION", "sa-east-1")
monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
)
assert redis_kwargs["credential_provider"]._region == "sa-east-1"
def test_aws_iam_region_prefers_explicit_over_environment(clean_redis_environment, monkeypatch):
monkeypatch.setenv("AWS_REGION", "sa-east-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="explicit-region",
)
assert redis_kwargs["credential_provider"]._region == "explicit-region"
def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user-value",
aws_iam_cache_name="iam-cache-value",
aws_iam_region="iam-region-value",
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._user_name == "iam-user-value"
assert provider._cache_name == "iam-cache-value"
assert provider._region == "iam-region-value"
@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no", "off"])
def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
aws_iam_auth=aws_iam_auth,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("aws_iam_auth", [True, "true", "True", "TRUE", "1", "yes"])
def test_aws_iam_auth_enabled_by_any_truthy_flag(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=aws_iam_auth,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
@pytest.mark.parametrize(
"aws_iam_serverless, expected",
[
pytest.param(None, False, id="unset"),
pytest.param(False, False, id="bool_false"),
pytest.param("false", False, id="string_false"),
pytest.param("0", False, id="string_zero"),
pytest.param(True, True, id="bool_true"),
pytest.param("true", True, id="string_true"),
pytest.param("1", True, id="string_one"),
],
)
def test_aws_iam_serverless_flag_reaches_the_provider(clean_redis_environment, aws_iam_serverless, expected):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
aws_iam_serverless=aws_iam_serverless,
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is expected
assert "aws_iam_serverless" not in redis_kwargs
def test_explicit_provider_wins_over_aws_iam(clean_redis_environment):
provider = _StubCredentialProvider()
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
credential_provider=provider,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert redis_kwargs["credential_provider"] is provider
assert "aws_iam_auth" not in redis_kwargs
def test_gcp_wins_over_aws_iam(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
gcp_service_account="sa@example.com",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert redis_kwargs["redis_connect_func"]._gcp_service_account == "sa@example.com"
def test_azure_wins_over_aws_iam(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
azure_redis_ad_token="true",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True
def test_async_cluster_installs_aws_iam_provider(clean_redis_environment):
startup_nodes = [{"host": "cluster-node", "port": 6379}]
client = get_redis_async_client(
startup_nodes=startup_nodes,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert isinstance(client.connection_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment):
provider = _StubCredentialProvider()

236
tests/test_litellm/test_redis_credential_provider.py Normal file
View file

@ -0,0 +1,236 @@
import asyncio
import builtins
import sys
from types import SimpleNamespace
from urllib.parse import parse_qs, urlsplit
import pytest
from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider
class _FakeCredentials:
def __init__(self, access_key: str) -> None:
self.access_key = access_key
self.secret_key = "synthetic-secret"
self.token = "synthetic-session-token"
def get_frozen_credentials(self):
return self
class _RotatingFakeCredentials:
def __init__(self) -> None:
self.calls = 0
def __bool__(self) -> bool:
return False
def get_frozen_credentials(self):
self.calls += 1
return SimpleNamespace(
access_key=f"AKIA-SYNTHETIC-{self.calls}",
secret_key="synthetic-secret",
token="synthetic-session-token",
)
class _FakeResolver:
def __init__(self, credentials: _FakeCredentials | _RotatingFakeCredentials | None) -> None:
self.credentials = credentials
self.calls = 0
def __call__(self):
self.calls += 1
return self.credentials
def test_elasticache_provider_signs_expected_query():
resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC"))
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
user_name, token = provider.get_credentials()
parsed = urlsplit("https://" + token)
query = parse_qs(parsed.query)
assert user_name == "iam-user"
assert parsed.netloc == "cache.example.com"
assert query["Action"] == ["connect"]
assert query["User"] == ["iam-user"]
assert query["X-Amz-Expires"] == ["900"]
assert "elasticache" in query["X-Amz-Credential"][0]
assert query["X-Amz-Credential"][0].split("/")[2] == "us-east-1"
assert not token.startswith("https://")
def test_elasticache_provider_resolves_credentials_once_but_refreshes_signature():
rotating_credentials = _RotatingFakeCredentials()
resolver = _FakeResolver(rotating_credentials)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
first = provider.get_credentials()
second = provider.get_credentials()
async_result = asyncio.run(provider.get_credentials_async())
assert first[0] == second[0] == async_result[0] == "iam-user"
assert first[1] != second[1]
assert async_result[1] != second[1]
assert resolver.calls == 1
assert rotating_credentials.calls == 3
def test_elasticache_provider_uses_botocore_session_credentials(monkeypatch):
credentials = _FakeCredentials("AKIA-SYNTHETIC")
monkeypatch.setattr("botocore.session.get_session", lambda: SimpleNamespace(get_credentials=lambda: credentials))
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
)
user_name, token = provider.get_credentials()
assert user_name == "iam-user"
assert "AKIA-SYNTHETIC" in token
def test_elasticache_provider_reports_missing_botocore(monkeypatch):
original_import = builtins.__import__
def import_without_botocore(name, *args, **kwargs):
if name == "botocore.session":
raise ImportError("synthetic missing dependency")
return original_import(name, *args, **kwargs)
monkeypatch.delitem(sys.modules, "botocore.session", raising=False)
monkeypatch.setattr(builtins, "__import__", import_without_botocore)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
)
with pytest.raises(ImportError, match="pip install boto3"):
provider.get_credentials()
def test_elasticache_provider_reports_missing_credentials():
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=_FakeResolver(None),
)
with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"):
provider.get_credentials()
def test_elasticache_provider_reports_missing_signing_dependency(monkeypatch):
original_import = builtins.__import__
def import_without_botocore_auth(name, *args, **kwargs):
if name == "botocore.auth":
raise ImportError("synthetic missing dependency")
return original_import(name, *args, **kwargs)
monkeypatch.delitem(sys.modules, "botocore.auth", raising=False)
monkeypatch.setattr(builtins, "__import__", import_without_botocore_auth)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
with pytest.raises(ImportError, match="pip install boto3"):
provider.get_credentials()
def test_elasticache_provider_recovers_after_a_failed_resolution():
resolver = _FakeResolver(None)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"):
provider.get_credentials()
resolver.credentials = _FakeCredentials("AKIA-SYNTHETIC")
user_name, token = provider.get_credentials()
assert user_name == "iam-user"
assert token
assert resolver.calls == 2
@pytest.mark.parametrize(
"provider_kwargs, expected_operation_params",
[
pytest.param({}, frozenset({"Action", "User"}), id="default_is_self_designed"),
pytest.param({"is_serverless": False}, frozenset({"Action", "User"}), id="self_designed"),
pytest.param({"is_serverless": True}, frozenset({"Action", "User", "ResourceType"}), id="serverless"),
],
)
def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_operation_params):
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
**provider_kwargs,
)
_, token = provider.get_credentials()
query_string = urlsplit("https://" + token).query
param_names = tuple(pair.split("=", 1)[0] for pair in query_string.split("&"))
first_auth_param = next(i for i, name in enumerate(param_names) if name.startswith("X-Amz-"))
query = parse_qs(query_string)
assert frozenset(param_names[:first_auth_param]) == expected_operation_params
assert all(name.startswith("X-Amz-") for name in param_names[first_auth_param:])
assert query.get("ResourceType") == (["ServerlessCache"] if "ResourceType" in expected_operation_params else None)
assert query["X-Amz-Signature"]
def test_elasticache_provider_lowercases_the_cache_name():
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="Mixed-Case-Cache",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
_, token = provider.get_credentials()
assert urlsplit("https://" + token).netloc == "mixed-case-cache"
def test_elasticache_provider_encodes_reserved_characters_in_the_user_name():
user_name = "iam user/with+reserved&chars"
provider = ElastiCacheIAMCredentialProvider(
user_name=user_name,
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
returned_user_name, token = provider.get_credentials()
query = parse_qs(urlsplit("https://" + token).query)
assert returned_user_name == user_name
assert query["User"] == [user_name]
assert query["Action"] == ["connect"]

25
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -26371,6 +26371,31 @@ export interface components {
* independently of the response-cache backend in `litellm_settings.cache_params`.
*/
CoordinationRedisParams: {
/**
* Aws Iam Auth
* @description enable AWS ElastiCache IAM authentication
*/
aws_iam_auth?: boolean | string | null;
/**
* Aws Iam Cache Name
* @description AWS ElastiCache cache name
*/
aws_iam_cache_name?: string | null;
/**
* Aws Iam Region
* @description AWS region for ElastiCache IAM authentication
*/
aws_iam_region?: string | null;
/**
* Aws Iam Serverless
* @description the ElastiCache cache is serverless rather than a self-designed cluster
*/
aws_iam_serverless?: boolean | string | null;
/**
* Aws Iam User Name
* @description AWS ElastiCache IAM user name
*/
aws_iam_user_name?: string | null;
/**
* Host
* @description Redis hostname