fix(redis): type ElastiCache IAM configuration

Generated with AI

Co-Authored-By: Claude Code

litellm/_redis.py

@@ -181,25 +181,24 @@ def _is_true(value: object | None) -> bool:
     return value is True or (isinstance(value, str) and value.lower() == "true")
 
 
-def _build_elasticache_iam_provider(redis_kwargs: dict) -> ElastiCacheIAMCredentialProvider | None:
-    if not _is_true(redis_kwargs.get("aws_iam_auth")):
-        return None
-
-    required_settings: Final = {
-        "aws_iam_user_name": redis_kwargs.get("aws_iam_user_name"),
-        "aws_iam_cache_name": redis_kwargs.get("aws_iam_cache_name"),
-        "aws_iam_region": redis_kwargs.get("aws_iam_region")
-        or get_secret_str("AWS_REGION")
-        or get_secret_str("AWS_DEFAULT_REGION"),
-    }
-    missing_settings: Final = tuple(name for name, value in required_settings.items() if not value)
+def _build_elasticache_iam_provider(
+    user_name: object | None,
+    cache_name: object | None,
+    region: object | None,
+) -> ElastiCacheIAMCredentialProvider:
+    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(required_settings["aws_iam_user_name"]),
-        cache_name=str(required_settings["aws_iam_cache_name"]),
-        region=str(required_settings["aws_iam_region"]),
+        user_name=str(user_name),
+        cache_name=str(cache_name),
+        region=str(region),
     )
 
 
@@ -507,10 +506,14 @@ def _get_redis_client_logic(**env_overrides):
                 "for Redis. Using Azure AD. Remove one to avoid misconfiguration."
             )
         elif _aws_iam_enabled:
-            aws_provider: Final = _build_elasticache_iam_provider(redis_kwargs)
-            if aws_provider is not None:
-                verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
-                redis_kwargs["credential_provider"] = aws_provider
+            verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
+            redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(
+                user_name=redis_kwargs.get("aws_iam_user_name"),
+                cache_name=redis_kwargs.get("aws_iam_cache_name"),
+                region=redis_kwargs.get("aws_iam_region")
+                or get_secret_str("AWS_REGION")
+                or get_secret_str("AWS_DEFAULT_REGION"),
+            )
 
     redis_kwargs.pop("gcp_service_account", None)
     redis_kwargs.pop("gcp_ssl_ca_certs", None)

litellm/_redis_credential_provider.py

@@ -1,11 +1,19 @@
+from __future__ import annotations
+
 import asyncio
 import threading
 import time
-from typing import Any, Final, Protocol
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Final, Protocol
 from urllib.parse import quote
 
 from redis.credentials import CredentialProvider
 
+if TYPE_CHECKING:
+    from botocore.credentials import Credentials
+else:
+    Credentials = Any  # rebind-ok: runtime alias for the type-checking-only botocore import
+
 # Azure AD scope for Redis Cache for Azure.
 AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
 
@@ -109,38 +117,24 @@ _ELASTICACHE_SERVICE_NAME: Final = "elasticache"
 _ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
 
 
-class _FrozenBotocoreCredentials(Protocol):
-    access_key: str
-    secret_key: str
-    token: str | None
-
-
-class _BotocoreCredentials(Protocol):
-    def get_frozen_credentials(self) -> _FrozenBotocoreCredentials: ...
-
-
-class _BotocoreCredentialsResolver(Protocol):
-    def __call__(self) -> _BotocoreCredentials | None: ...
-
-
 class ElastiCacheIAMCredentialProvider(CredentialProvider):
     def __init__(
         self,
         user_name: str,
         cache_name: str,
         region: str,
-        credentials_resolver: _BotocoreCredentialsResolver | None = None,
+        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
         self._region = region
         self._credentials_resolver = credentials_resolver or self._resolve_credentials
-        self._credentials: _BotocoreCredentials | None = None
+        self._credentials: Credentials | None = None
         self._token_lifetime_seconds = token_lifetime_seconds
 
     @staticmethod
-    def _resolve_credentials() -> Any:
+    def _resolve_credentials() -> Credentials | None:
         try:
             import botocore.session
         except ImportError as e:
@@ -187,6 +181,14 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider):
         return await asyncio.to_thread(self._get_credentials)
 
 
+class _AzureToken(Protocol):
+    token: str
+
+
+class _AzureCredential(Protocol):
+    def get_token(self, scope: str) -> _AzureToken: ...
+
+
 class AzureADCredentialProvider(CredentialProvider):
     """
     redis.credentials.CredentialProvider implementation that supplies Azure AD
@@ -198,7 +200,7 @@ class AzureADCredentialProvider(CredentialProvider):
     fail authentication after the initial token expired (~1 hour TTL).
     """
 
-    def __init__(self, credential: Any, username: str | None = None) -> None:
+    def __init__(self, credential: _AzureCredential, username: str | None = None) -> None:
         self._credential = credential
         self._username = username
 

ruff-strict-budget.json

@@ -24,7 +24,7 @@
     "limit": 133
   },
   "ANN401": {
-    "limit": 1153
+    "limit": 1152
   },
   "ASYNC230": {
     "limit": 11

tests/test_litellm/test_redis_credential_provider.py

@@ -4,10 +4,7 @@ from urllib.parse import parse_qs, urlsplit
 
 import pytest
 
-from litellm._redis_credential_provider import (
-    ElastiCacheIAMCredentialProvider,
-    _BotocoreCredentials,
-)
+from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider
 
 
 class _FakeCredentials:
@@ -20,16 +17,6 @@ class _FakeCredentials:
         return self
 
 
-class _FakeResolver:
-    def __init__(self, credentials: _BotocoreCredentials | None) -> None:
-        self.credentials = credentials
-        self.calls = 0
-
-    def __call__(self):
-        self.calls += 1
-        return self.credentials
-
-
 class _RotatingFakeCredentials:
     def __init__(self) -> None:
         self.calls = 0
@@ -46,6 +33,16 @@ class _RotatingFakeCredentials:
         )
 
 
+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(

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -24559,6 +24559,26 @@ 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 User Name
+             * @description AWS ElastiCache IAM user name
+             */
+            aws_iam_user_name?: string | null;
             /**
              * Host
              * @description Redis hostname