From ca87d4e6e382fe55cfb10f0a17f15fa0475f3d59 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 10:59:33 -0700 Subject: [PATCH] refactor(redis-store): login throttle store takes the source group and username, not keys --- litellm/proxy/auth/login_throttle.py | 17 +++++++------- litellm/redis/stores/login_throttle.py | 32 ++++++++++++++++++-------- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/auth/login_throttle.py b/litellm/proxy/auth/login_throttle.py index 205d7f48601..c17d92aedc8 100644 --- a/litellm/proxy/auth/login_throttle.py +++ b/litellm/proxy/auth/login_throttle.py @@ -298,7 +298,7 @@ class LoginThrottle: """Refuses a blocked key before any credential is looked at; otherwise hands back the attempt to settle.""" if not self.enabled: return LoginAttempt(throttle=self, username=username) - block: Final = await self._active_block(self._keys(username)) + block: Final = await self._active_block(username) if block is None: return LoginAttempt(throttle=self, username=username) verbose_proxy_logger.warning( @@ -310,9 +310,9 @@ class LoginThrottle: ) raise self.refused(block.retry_after) - async def _active_block(self, keys: LoginKeys) -> Block | None: - local: Final = self._local_block_ttls(keys) - shared: Final = await self._shared_block_ttls(keys) + async def _active_block(self, username: str) -> Block | None: + local: Final = self._local_block_ttls(self._keys(username)) + shared: Final = await self._shared_block_ttls(username) user_ttl: Final = max(local[0], shared[0]) source_ttl: Final = max(local[1], shared[1]) if self.source_limit is not None and source_ttl > 0: @@ -322,11 +322,11 @@ class LoginThrottle: return None @with_service_target(_LOGIN_THROTTLE_TARGET) - async def _shared_block_ttls(self, keys: LoginKeys) -> BlockTtls: + async def _shared_block_ttls(self, username: str) -> BlockTtls: store: Final = self._store() if store is None: return LOGIN_THROTTLE_NOT_BLOCKED - blocks: Final = await store.blocks(keys) + blocks: Final = await store.blocks(source_group(self.client_ip), username) return LOGIN_THROTTLE_NOT_BLOCKED if isinstance(blocks, RedisFailure) else blocks def _local_block_ttls(self, keys: LoginKeys) -> BlockTtls: @@ -347,7 +347,8 @@ class LoginThrottle: None if store is None else await store.record_failure( - keys, + source_group(self.client_ip), + username, user_limit=self.user_limit, source_limit=source_limit, window=timedelta(seconds=self.window_seconds), @@ -376,7 +377,7 @@ class LoginThrottle: keys: Final = self._keys(username) store: Final = self._store() if store is not None: - _ = await store.clear(keys) + _ = await store.clear(source_group(self.client_ip), username) self.counters.delete_cache(keys.pair_counter) @staticmethod diff --git a/litellm/redis/stores/login_throttle.py b/litellm/redis/stores/login_throttle.py index 56394f8fa95..12b72ed4201 100644 --- a/litellm/redis/stores/login_throttle.py +++ b/litellm/redis/stores/login_throttle.py @@ -50,30 +50,44 @@ def login_keys(source_group: str, username: str) -> LoginKeys: class LoginThrottleStore(Protocol): - async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure: ... + async def blocks(self, source_group: str, username: str) -> BlockTtls | RedisFailure: ... async def record_failure( - self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta + self, + source_group: str, + username: str, + *, + user_limit: int, + source_limit: int, + window: timedelta, + block: timedelta, ) -> BlockTtls | RedisFailure: ... - async def clear(self, keys: LoginKeys) -> None | RedisFailure: ... + async def clear(self, source_group: str, username: str) -> None | RedisFailure: ... @dataclass(frozen=True, slots=True) class TransportLoginThrottleStore: transport: RedisTransport - async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure: - return _ttls(await self.transport.eval(_BLOCK_TTLS_LUA, keys, ())) + async def blocks(self, source_group: str, username: str) -> BlockTtls | RedisFailure: + return _ttls(await self.transport.eval(_BLOCK_TTLS_LUA, login_keys(source_group, username), ())) async def record_failure( - self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta + self, + source_group: str, + username: str, + *, + user_limit: int, + source_limit: int, + window: timedelta, + block: timedelta, ) -> BlockTtls | RedisFailure: args: Final = (user_limit, source_limit, _whole_seconds(window), _whole_seconds(block)) - return _ttls(await self.transport.eval(_RECORD_FAILURE_LUA, keys, args)) + return _ttls(await self.transport.eval(_RECORD_FAILURE_LUA, login_keys(source_group, username), args)) - async def clear(self, keys: LoginKeys) -> None | RedisFailure: - deleted: Final = await self.transport.delete(keys.pair_counter) + async def clear(self, source_group: str, username: str) -> None | RedisFailure: + deleted: Final = await self.transport.delete(login_keys(source_group, username).pair_counter) return deleted if isinstance(deleted, RedisFailure) else None