diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 95b127b5b2a..61d7078ae4c 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -86,6 +86,10 @@ class UserApiKeyCache(DualCache): default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl ) + def update_in_memory_max_size(self, max_size: int | None) -> None: + super().update_in_memory_max_size(max_size) + self.key_object_cache.update_in_memory_max_size(max_size) + def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None ) -> None: diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index f24175a1922..09523cd1901 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -414,6 +414,36 @@ class TestUserKeyObjectPartition: assert cache.in_memory_cache_for(HASHED_TOKEN) is cache.key_object_cache.in_memory_cache assert cache.in_memory_cache_for(end_user_cache_key("u1")) is cache.in_memory_cache + def test_update_in_memory_max_size_applies_to_key_object_partition(self): + cache = UserApiKeyCache( + in_memory_cache=InMemoryCache(max_size_in_memory=2), + key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2), + ) + cache.update_in_memory_max_size(3) + + tokens = tuple(hashlib.sha256(f"sk-key-{i}".encode()).hexdigest() for i in range(3)) + for token in tokens: + cache.set_cache(token, _make_key_obj(token), model_type=UserAPIKeyAuth, ttl=100) + for i in range(3): + cache.set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=100) + + first_key = cache.get_cache(tokens[0], model_type=UserAPIKeyAuth) + assert first_key is not None, "key partition still evicts at its old capacity" + assert first_key.token == tokens[0] + assert cache.get_cache(end_user_cache_key("u0")) == {"user_id": "u0"} + + def test_update_in_memory_max_size_none_resets_key_object_partition_to_default(self): + cache = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=1)) + cache.update_in_memory_max_size(None) + + tokens = tuple(hashlib.sha256(f"sk-key-{i}".encode()).hexdigest() for i in range(2)) + for token in tokens: + cache.set_cache(token, _make_key_obj(token), model_type=UserAPIKeyAuth, ttl=100) + + first_key = cache.get_cache(tokens[0], model_type=UserAPIKeyAuth) + assert first_key is not None + assert first_key.token == tokens[0] + class TestManagementObjectTTL: """