diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index be761e1258b..26ad89a70cc 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -376,7 +376,9 @@ class DualCache(BaseCache): ) # async_batch_set_cache - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs): + async def async_set_cache_pipeline( + self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs + ): """ Batch write values to the cache """ diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index fb2ca6372c0..dbe11882b3c 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -196,9 +196,7 @@ class AuthCacheInvalidationSubscriber: for additional_cache in self._additional_in_memory_caches: additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) return - in_memory_cache: Final = self._user_api_key_cache.in_memory_cache - if in_memory_cache is not None: - in_memory_cache.delete_cache(parsed.cache_key) + self._user_api_key_cache.in_memory_cache_for(parsed.cache_key).delete_cache(parsed.cache_key) for additional_cache in self._additional_in_memory_caches: additional_cache.delete_cache(parsed.cache_key) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 554a6ae8d1a..1d4024bb84e 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -147,8 +147,11 @@ async def memory_usage_in_mem_cache( llm_router.cache.in_memory_cache.ttl_dict ) - num_items_in_user_api_key_cache: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len( - user_api_key_cache.in_memory_cache.ttl_dict + num_items_in_user_api_key_cache: Final = ( + len(user_api_key_cache.in_memory_cache.cache_dict) + + len(user_api_key_cache.in_memory_cache.ttl_dict) + + len(user_api_key_cache.key_object_cache.in_memory_cache.cache_dict) + + len(user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict) ) num_items_in_proxy_logging_obj_cache: Final = len( @@ -189,6 +192,8 @@ async def memory_usage_in_mem_cache_items( return { "user_api_key_cache": user_api_key_cache.in_memory_cache.cache_dict, "user_api_key_ttl": user_api_key_cache.in_memory_cache.ttl_dict, + "user_key_object_cache": user_api_key_cache.key_object_cache.in_memory_cache.cache_dict, + "user_key_object_ttl": user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict, "llm_router_cache": llm_router_in_memory_cache_dict, "llm_router_ttl": llm_router_in_memory_ttl_dict, "proxy_logging_obj_cache": proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict, @@ -294,7 +299,9 @@ async def get_memory_summary( try: # User API key cache - user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len( + user_api_key_cache.key_object_cache.in_memory_cache.cache_dict + ) total_cache_items += user_cache_items caches["user_api_keys"] = { "count": user_cache_items, @@ -429,10 +436,16 @@ def _get_cache_memory_stats( cache_stats: Final[dict[str, object]] = {} try: # User API key cache - user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) - user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + key_object_in_memory_cache: Final = user_api_key_cache.key_object_cache.in_memory_cache + user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) + sys.getsizeof( + key_object_in_memory_cache.cache_dict + ) + user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + sys.getsizeof( + key_object_in_memory_cache.ttl_dict + ) cache_stats["user_api_key_cache"] = { - "num_items": len(user_api_key_cache.in_memory_cache.cache_dict), + "num_items": len(user_api_key_cache.in_memory_cache.cache_dict) + + len(key_object_in_memory_cache.cache_dict), "cache_dict_size_bytes": user_cache_size, "ttl_dict_size_bytes": user_ttl_size, "total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2), diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 76982d30306..cb72088ee4a 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -1,11 +1,15 @@ from __future__ import annotations +import re +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -14,6 +18,13 @@ if TYPE_CHECKING: T = TypeVar("T", bound=BaseModel) +_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") + + +def is_user_key_cache_key(key: str) -> bool: + """Only user-key objects are cached under a bare ``hash_token`` digest; every other object uses a prefixed key.""" + return _HASHED_TOKEN_CACHE_KEY.fullmatch(key) is not None + class UserApiKeyCache(DualCache): """ @@ -36,10 +47,50 @@ class UserApiKeyCache(DualCache): ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). + User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, + ``key_object_cache``, so churn in the other management objects cannot evict them. Both + partitions share the same Redis backend and TTL settings. + ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. """ + def __init__( + self, + in_memory_cache: InMemoryCache | None = None, + redis_cache: RedisCache | None = None, + default_in_memory_ttl: float | None = None, + default_redis_ttl: float | None = None, + key_object_in_memory_cache: InMemoryCache | None = None, + ) -> None: + super().__init__( + in_memory_cache=in_memory_cache, + redis_cache=redis_cache, + default_in_memory_ttl=default_in_memory_ttl, + default_redis_ttl=default_redis_ttl, + ) + self.key_object_cache: Final = DualCache( + in_memory_cache=key_object_in_memory_cache or InMemoryCache(), + redis_cache=redis_cache, + default_in_memory_ttl=default_in_memory_ttl, + default_redis_ttl=default_redis_ttl, + ) + + def in_memory_cache_for(self, key: str) -> InMemoryCache: + return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache + + def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: + super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) + self.key_object_cache.update_cache_ttl( + default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl + ) + + def attach_redis_cache( + self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None + ) -> None: + super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) + self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) + @overload def get_cache( self, @@ -71,7 +122,11 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - cached: Final = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + cached: Final = ( + self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + if is_user_key_cache_key(key) + else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + ) if model_type is None: return cached if cached is None: @@ -117,8 +172,14 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) - cached: Final = await super().async_get_cache( - key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + cached: Final = ( + await self.key_object_cache.async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if is_user_key_cache_key(key) + else await super().async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) ) if model_type is None: return cached @@ -137,20 +198,49 @@ class UserApiKeyCache(DualCache): def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if key is not None and is_user_key_cache_key(key): + return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + if key is not None and is_user_key_cache_key(key): + return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None: + def delete_cache(self, key: str) -> None: + if is_user_key_cache_key(key): + self.key_object_cache.delete_cache(key) + return + super().delete_cache(key) + + async def async_delete_cache(self, key: str) -> None: + if is_user_key_cache_key(key): + await self.key_object_cache.async_delete_cache(key) + return + await super().async_delete_cache(key) + + def flush_cache(self) -> None: + super().flush_cache() + self.key_object_cache.in_memory_cache.flush_cache() + + async def async_set_cache_pipeline( + self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object + ) -> None: """ Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. """ - normalized: Final = [(key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list] - return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs) + normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list) + key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0])) + other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0])) + if key_object_entries: + await self.key_object_cache.async_set_cache_pipeline( + cache_list=key_object_entries, local_only=local_only, **kwargs + ) + if other_entries: + await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) #: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 749a940de0e..db4467dda5b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3729,7 +3729,7 @@ async def delete_key_fn( ) verbose_proxy_logger.debug( - "/keys/delete - cache after delete: %s", user_api_key_cache.in_memory_cache.cache_dict + "/keys/delete - cache after delete: %s", user_api_key_cache.key_object_cache.in_memory_cache.cache_dict ) asyncio.create_task( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 5a99139a67f..c53706ac938 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7150,7 +7150,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals key = "sk-alice-key" cache = UserApiKeyCache() - cache.in_memory_cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"}) + cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"}) proxy_globals.user_api_key_cache = cache proxy_globals.prisma_client = object() diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 7d5fc1a3544..4e2059ac30b 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -1,4 +1,5 @@ import asyncio +import hashlib import json from typing import Iterable, List, Optional, Tuple from unittest.mock import patch @@ -7,6 +8,7 @@ import pytest from redis.asyncio import Redis from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -145,6 +147,35 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] +@pytest.mark.asyncio +async def test_subscriber_deletes_key_object_partition_entry_on_message() -> None: + """ + LIT-7563 moved user-key objects into their own in-memory partition; a key + invalidation broadcast must still evict the hashed-token entry there, or a + deleted key keeps authenticating on other workers until its TTL expires. + """ + hashed_token = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() + cache = UserApiKeyCache() + cache.set_cache(hashed_token, UserAPIKeyAuth(token=hashed_token), model_type=UserAPIKeyAuth) + assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message(hashed_token)]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + for _ in range(200): + if cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None + + @pytest.mark.asyncio async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> 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 33e0d8bf38e..2d5d76ed542 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 @@ -1,3 +1,4 @@ +import hashlib import json from typing import Any @@ -10,10 +11,14 @@ from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + end_user_cache_key, get_management_object_ttl, + is_user_key_cache_key, ) from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum +HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() + class CapturingInMemoryCache(InMemoryCache): """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" @@ -204,9 +209,7 @@ class TestUserApiKeyCache: # Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values # for dict-based models (deserialize returns None). - await cache.in_memory_cache.async_set_cache( - key="k", value="invalid-payload-not-a-dict" - ) + await cache.in_memory_cache.async_set_cache(key="k", value="invalid-payload-not-a-dict") value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) assert value is None @@ -224,6 +227,141 @@ class TestUserApiKeyCache: fake.set_cache("k2", {"ok": NotSerializable()}) +class TestUserKeyObjectPartition: + """ + Regression for LIT-7563: user-key objects share one 200-entry ``InMemoryCache`` with + every other management object, so end-user / team / tag churn evicts hot keys and + forces a ``LiteLLM_VerificationToken`` lookup on the next request. + """ + + @pytest.mark.parametrize( + ("key", "expected"), + [ + (HASHED_TOKEN, True), + (HASHED_TOKEN.upper(), False), + (f"team_id:{HASHED_TOKEN}", False), + (end_user_cache_key("u1"), False), + ("sk-lit7563-hot-key", False), + ], + ) + def test_is_user_key_cache_key(self, key: str, expected: bool): + assert is_user_key_cache_key(key) is expected + + @pytest.mark.asyncio + async def test_management_object_churn_does_not_evict_key_object(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100) + for i in range(2): + await cache.async_set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200) + + key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict + + def test_sync_write_and_read_route_to_key_object_partition(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100) + for i in range(2): + cache.set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200) + + key_obj = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + + @pytest.mark.asyncio + async def test_redis_hit_backfills_key_object_partition_with_configured_ttl(self): + redis = FakeRedisCache() + writer = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30) + await writer.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + key_partition = CapturingInMemoryCache() + reader = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30, key_object_in_memory_cache=key_partition) + key_obj = await reader.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert key_partition.last_ttl == 30 + assert HASHED_TOKEN not in reader.in_memory_cache.cache_dict + + @pytest.mark.asyncio + async def test_update_cache_ttl_applies_to_key_object_partition(self): + key_partition = CapturingInMemoryCache() + cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition) + cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7) + + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + assert key_partition.last_ttl == 7 + + @pytest.mark.asyncio + async def test_attach_redis_cache_applies_to_key_object_partition(self): + redis = FakeRedisCache() + cache = UserApiKeyCache() + cache.attach_redis_cache(redis) + + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + other_worker = UserApiKeyCache(redis_cache=redis) + key_obj = await other_worker.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + + @pytest.mark.asyncio + async def test_delete_removes_key_object_from_partition_and_redis(self): + redis = FakeRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is not None + + cache.delete_cache(HASHED_TOKEN) + + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert await redis.async_get_cache(HASHED_TOKEN) is None + + @pytest.mark.asyncio + async def test_async_delete_removes_key_object_from_partition_and_redis(self): + redis = FakeRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + + await cache.async_delete_cache(HASHED_TOKEN) + + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert await redis.async_get_cache(HASHED_TOKEN) is None + + @pytest.mark.asyncio + async def test_pipeline_write_routes_each_entry_to_its_partition(self): + cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) + await cache.async_set_cache_pipeline( + [(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN))] + + [(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}) for i in range(2)], + ttl=100, + ) + + key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + assert key_obj is not None + assert key_obj.token == HASHED_TOKEN + assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict + assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + + def test_flush_clears_key_object_partition(self): + cache = UserApiKeyCache() + cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + cache.set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) + + cache.flush_cache() + + assert cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert cache.get_cache(end_user_cache_key("u1")) is None + + def test_in_memory_cache_for_routes_by_key(self): + cache = UserApiKeyCache() + 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 + + class TestManagementObjectTTL: """ Regression for LIT-3338: ``general_settings.user_api_key_cache_ttl`` (which the @@ -238,19 +376,13 @@ class TestManagementObjectTTL: def test_falls_back_to_constant_when_no_default_configured(self): cache = UserApiKeyCache() assert cache.default_in_memory_ttl is None - assert ( - get_management_object_ttl(cache) - == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(cache) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL def test_resolves_on_a_plain_dual_cache(self): # Many call sites are typed UserApiKeyCache but exercised in tests with a # bare DualCache; the resolver must work on the base type, not just the subclass. assert get_management_object_ttl(DualCache(default_in_memory_ttl=300)) == 300 - assert ( - get_management_object_ttl(DualCache()) - == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(DualCache()) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL @pytest.mark.asyncio async def test_management_write_uses_configured_ttl_over_constant(self): @@ -260,9 +392,7 @@ class TestManagementObjectTTL: redis_cache=FakeRedisCache(), default_in_memory_ttl=300, ) - assert get_management_object_ttl(cache) != ( - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL - ) + assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) await cache.async_set_cache( "team_id:abc",