scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

refactor(proxy): move login throttle sentinels into constants

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-17 16:42:04 +00:00
parent 9a365d2021
commit b227a8c4c9
4 changed files with 36 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/constants.py
View file

@ -1644,6 +1644,11 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = int(
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int(
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
)
LOGIN_THROTTLE_CACHE_KEY_PREFIX: Final = "login_fail"
LOGIN_THROTTLE_UNKNOWN_SOURCE: Final = "unknown"
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS: Final = 20_000
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS: Final = 10_000
LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"

37
litellm/proxy/auth/login_throttle.py
View file

@ -17,7 +17,6 @@ import time
from collections.abc import Mapping
from dataclasses import dataclass
from functools import cache
from types import MappingProxyType
from typing import Final, Literal, NamedTuple, NoReturn, Protocol, TypeAlias
from fastapi import Request, status
@ -27,6 +26,14 @@ from redis.exceptions import RedisError
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.constants import (
EMPTY_MAPPING,
LOGIN_THROTTLE_CACHE_KEY_PREFIX,
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS,
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS,
LOGIN_THROTTLE_NOT_BLOCKED,
LOGIN_THROTTLE_UNKNOWN_SOURCE,
)
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.network import TrustedProxyConfig, normalize_cidr_ranges, resolve_client_ip
from litellm.secret_managers.main import get_secret_bool
@ -45,12 +52,6 @@ WINDOW_KEY: Final = "failed_login_window_seconds"
BLOCK_KEY: Final = "failed_login_block_seconds"
TRUSTED_PROXY_RANGES_KEY: Final = "trusted_proxy_ranges"
_CACHE_KEY_PREFIX: Final = "login_fail"
_UNKNOWN_SOURCE: Final = "unknown"
_MAX_TRACKED_COUNTERS: Final = 20_000
_MAX_TRACKED_BLOCKS: Final = 10_000
_NO_SETTINGS: Final[Mapping[str, object]] = MappingProxyType({})
_NOT_BLOCKED: Final = (0, 0)
_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError)
_LOCAL_BLOCK_EXPIRY: Final = TypeAdapter[float | None](float | None)
_SOURCE_LIMIT_OVERRIDES: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object])
@ -94,9 +95,11 @@ _RECORD_FAILURE_LUA: Final = (
)
_COUNTERS: Final = InMemoryCache(
max_size_in_memory=_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS
)
_BLOCKS: Final = InMemoryCache(
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_BLOCKS, default_ttl=DEFAULT_FAILED_LOGIN_BLOCK_SECONDS
)
_BLOCKS: Final = InMemoryCache(max_size_in_memory=_MAX_TRACKED_BLOCKS, default_ttl=DEFAULT_FAILED_LOGIN_BLOCK_SECONDS)
@cache
@ -249,13 +252,13 @@ class LoginThrottle:
general_settings: Mapping[str, object] | None,
redis_cache: RedisCache | None,
) -> LoginThrottle:
settings: Final = general_settings if general_settings is not None else _NO_SETTINGS
settings: Final = general_settings if general_settings is not None else EMPTY_MAPPING
proxies: Final = declared_proxy_ranges(settings)
resolved, _ = resolve_client_ip(
request, TrustedProxyConfig(use_forwarded_for=bool(proxies), trusted_proxy_cidrs=proxies or ())
)
return cls(
client_ip=resolved or _UNKNOWN_SOURCE,
client_ip=resolved or LOGIN_THROTTLE_UNKNOWN_SOURCE,
source_limit=_source_limit(settings, resolved) if proxies is not None and resolved is not None else None,
user_limit=_int_setting(settings, USER_LIMIT_KEY, DEFAULT_MAX_FAILED_LOGIN_ATTEMPTS_PER_USER),
window_seconds=_int_setting(settings, WINDOW_KEY, DEFAULT_FAILED_LOGIN_WINDOW_SECONDS),
@ -270,10 +273,10 @@ class LoginThrottle:
group: Final = source_group(self.client_ip)
user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest()
return _Keys(
pair_counter=f"{_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}",
pair_block=f"{_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}",
source_counter=f"{_CACHE_KEY_PREFIX}:{{{group}}}:source",
source_block=f"{_CACHE_KEY_PREFIX}:{{{group}}}:block:source",
pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}",
pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}",
source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:source",
source_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:source",
)
async def attempt(self, username: str) -> LoginAttempt:
@ -305,14 +308,14 @@ class LoginThrottle:
async def _shared_block_ttls(self, keys: _Keys) -> _BlockTtls:
if self.redis_cache is None:
return _NOT_BLOCKED
return LOGIN_THROTTLE_NOT_BLOCKED
try:
return _LUA_BLOCK_TTLS.validate_python(
await self.redis_cache.async_register_script(_BLOCK_TTLS_LUA)(keys, ())
)
except _REDIS_FAILURES as err:
self._warn_redis(err)
return _NOT_BLOCKED
return LOGIN_THROTTLE_NOT_BLOCKED
def _local_block_ttls(self, keys: _Keys) -> _BlockTtls:
return self._local_block_ttl(keys.pair_block), self._local_block_ttl(keys.source_block)

20
tests/test_litellm/proxy/auth/test_login_utils.py
View file

@ -1361,7 +1361,7 @@ class _DownRedis(_FakeRedis):
@pytest.mark.asyncio
async def test_redis_is_the_only_counter_while_it_answers(monkeypatch):
"""Every worker must spend the same budget, see the same block, and a success must clear the pair for all."""
from litellm.proxy.auth.login_throttle import _CACHE_KEY_PREFIX
from litellm.constants import LOGIN_THROTTLE_CACHE_KEY_PREFIX
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
@ -1371,7 +1371,7 @@ async def test_redis_is_the_only_counter_while_it_answers(monkeypatch):
second_worker = _throttle(user_limit=2, stores=_stores(), redis_cache=redis)
assert [await _fail(first_worker, username="user@corp.com") for _ in range(3)] == ["401"] * 3
assert not [k for k in first_worker.counters.cache_dict if str(k).startswith(_CACHE_KEY_PREFIX)], (
assert not [k for k in first_worker.counters.cache_dict if str(k).startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX)], (
"with Redis answering, no worker may keep a counter of its own"
)
assert not first_worker.blocks.cache_dict
@ -1437,7 +1437,8 @@ async def test_a_failed_redis_delete_still_clears_this_workers_counter(monkeypat
async def test_counters_do_not_share_the_key_authentication_cache(monkeypatch):
"""Regression: throttle entries must not evict cached credentials from user_api_key_cache."""
from litellm.proxy import proxy_server as ps
from litellm.proxy.auth.login_throttle import _CACHE_KEY_PREFIX, LoginThrottle
from litellm.constants import LOGIN_THROTTLE_CACHE_KEY_PREFIX
from litellm.proxy.auth.login_throttle import LoginThrottle
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
@ -1453,7 +1454,7 @@ async def test_counters_do_not_share_the_key_authentication_cache(monkeypatch):
assert await _fail(throttle, username=f"made-up-{i}@example.com") == "401"
added = set(ps.user_api_key_cache.in_memory_cache.cache_dict) - auth_cache_keys_before
assert not [k for k in added if str(k).startswith(_CACHE_KEY_PREFIX)]
assert not [k for k in added if str(k).startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX)]
def test_settings_that_arrive_as_environment_strings_are_honored():
@ -1557,17 +1558,12 @@ async def test_a_username_spray_cannot_evict_an_active_block(monkeypatch):
"""Counters and blocks live in separate bounded stores, so a flood of made-up pairs fills the counter
store while the blocks it already earned stay in force."""
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.login_throttle import (
_BLOCKS,
_COUNTERS,
_MAX_TRACKED_BLOCKS,
_MAX_TRACKED_COUNTERS,
LoginThrottle,
)
from litellm.constants import LOGIN_THROTTLE_MAX_TRACKED_BLOCKS, LOGIN_THROTTLE_MAX_TRACKED_COUNTERS
from litellm.proxy.auth.login_throttle import _BLOCKS, _COUNTERS, LoginThrottle
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
assert _MAX_TRACKED_COUNTERS >= 10_000 and _MAX_TRACKED_BLOCKS >= 10_000
assert LOGIN_THROTTLE_MAX_TRACKED_COUNTERS >= 10_000 and LOGIN_THROTTLE_MAX_TRACKED_BLOCKS >= 10_000
assert _COUNTERS is not _BLOCKS
counters, blocks = InMemoryCache(max_size_in_memory=50), InMemoryCache(max_size_in_memory=50)
throttle = LoginThrottle(

5
tests/test_litellm/proxy/proxy_server/conftest.py
View file

@ -521,13 +521,14 @@ def reset_login_throttle(monkeypatch):
window, so without this a failed sign-in test could block unrelated tests later.
Only the throttle's own keys are removed, so other cache entries remain untouched.
"""
from litellm.constants import LOGIN_THROTTLE_CACHE_KEY_PREFIX
from litellm.proxy import proxy_server as ps
from litellm.proxy.auth.login_throttle import _BLOCKS, _CACHE_KEY_PREFIX, _COUNTERS
from litellm.proxy.auth.login_throttle import _BLOCKS, _COUNTERS
def _drop_throttle_keys() -> None:
for store in (_COUNTERS, _BLOCKS):
for key in tuple(store.cache_dict) + tuple(store.ttl_dict):
if key.startswith(_CACHE_KEY_PREFIX):
if key.startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX):
store.delete_cache(key)
monkeypatch.setattr(ps, "redis_usage_cache", None)