diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3da899a5610..3166e58106c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1149,7 +1149,7 @@ async def get_end_user_object( if end_user_id is None: return None - _key: Final = f"end_user_id:{end_user_id}" + _key: Final = _end_user_cache_key(end_user_id) # Check cache first cached_user_obj: Final = await user_api_key_cache.async_get_cache( @@ -1191,7 +1191,7 @@ async def get_end_user_object( # Save to cache await user_api_key_cache.async_set_cache( - key=f"end_user_id:{end_user_id}", + key=_key, value=_response, model_type=LiteLLM_EndUserTable, ) @@ -1206,6 +1206,39 @@ _END_USER_VALIDATION_NEGATIVE_TTL: Final = 60 _END_USER_VALIDATION_POSITIVE_TTL: Final = 300 +def _end_user_cache_key(end_user_id: str) -> str: + return f"end_user_id:{end_user_id}" + + +def _end_user_validation_cache_key(end_user_id: str) -> str: + return f"end_user_validation:{end_user_id}" + + +async def delete_cached_end_user_object( + end_user_id: str, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Every endpoint that mutates litellm_endusertable must call this: the auth path + reads the end-user object and the id-validation verdict cache-first, so without + invalidation a blocked, rebudgeted, or deleted customer keeps being served until + the TTL expires. Best-effort on both steps, and broadcast so workers that did not + handle the mutation drop their in-memory copy too. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + for cache_key in (_end_user_cache_key(end_user_id), _end_user_validation_cache_key(end_user_id)): + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: a cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached end-user entry %s; a stale end-user may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + async def resolve_and_validate_end_user_id( raw_end_user_id: str | None, prisma_client: PrismaClient | None, @@ -1241,7 +1274,7 @@ async def resolve_and_validate_end_user_id( if prisma_client is None: return raw_end_user_id - cache_key: Final = f"end_user_validation:{raw_end_user_id}" + cache_key: Final = _end_user_validation_cache_key(raw_end_user_id) cached: Final = await user_api_key_cache.async_get_cache(key=cache_key) if cached == "valid": return raw_end_user_id @@ -1904,7 +1937,19 @@ async def _delete_cache_key_object( hashed_token: str, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None, + broadcast_invalidation: bool = True, ): + """ + ``broadcast_invalidation`` is what makes a revoked or regenerated key stop + working fleet-wide right away; the local delete plus the Redis delete leave + every other worker serving its own in-memory copy until the TTL expires. + + Callers that evict a key each worker can already invalidate on its own (an + expired key, whose expiry every worker reads off the cached object) pass + False so a client replaying that key can't drive a broadcast per request. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + key: Final = hashed_token user_api_key_cache.delete_cache(key=key) @@ -1913,6 +1958,9 @@ async def _delete_cache_key_object( if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + if broadcast_invalidation: + await publish_auth_cache_invalidation(cache_key=key) + @log_db_metrics async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None): @@ -2105,6 +2153,8 @@ async def _delete_cache_access_object( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None = None, ): + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + key: Final = f"access_group_id:{access_group_id}" user_api_key_cache.delete_cache(key=key) @@ -2113,6 +2163,8 @@ async def _delete_cache_access_object( if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + await publish_auth_cache_invalidation(cache_key=key) + @log_db_metrics async def get_access_object( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9dc450befc2..9d74c1eedfb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1573,6 +1573,7 @@ async def _user_api_key_auth_builder( hashed_token=hash_token(api_key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + broadcast_invalidation=False, ) raise ProxyException( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 7fc8da42a3d..736bbb7ffdc 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,7 +1,15 @@ import asyncio import json +from collections.abc import Mapping from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Final +from typing import ( # noqa: TID251 # cast narrows a json str to the Literal + TYPE_CHECKING, + Final, + Literal, + TypeAlias, + cast, + get_args, +) from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.config_sync_pubsub import ( @@ -11,8 +19,13 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) if TYPE_CHECKING: + from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +InvalidationTarget: TypeAlias = Literal["user_api_key", "config_param"] + +DEFAULT_INVALIDATION_TARGET: Final[InvalidationTarget] = "user_api_key" +_INVALIDATION_TARGETS: Final[frozenset[str]] = frozenset(get_args(InvalidationTarget)) AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _POLL_TIMEOUT_SECONDS: Final = 1.0 @@ -29,32 +42,41 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: cache_key: str + cache: InvalidationTarget -def _cache_invalidation_message_json(cache_key: str) -> str: - return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) +def _cache_invalidation_message_json(cache_key: str, target: InvalidationTarget) -> str: + return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key, cache=target))) -def _cache_key_from_message_data(data: object) -> str | None: - if isinstance(data, bytes): - data = data.decode("utf-8", errors="replace") - if not isinstance(data, str): +def _invalidation_from_message_data(data: object) -> _CacheInvalidationMessage | None: + decoded: Final = data.decode("utf-8", errors="replace") if isinstance(data, bytes) else data + if not isinstance(decoded, str): return None try: - parsed: Final = json.loads(data) + parsed: Final = json.loads(decoded) except json.JSONDecodeError: return None if not isinstance(parsed, dict): return None cache_key: Final = parsed.get("cache_key") - return cache_key if isinstance(cache_key, str) else None + if not isinstance(cache_key, str): + return None + raw_target: Final = parsed.get("cache", DEFAULT_INVALIDATION_TARGET) + if raw_target not in _INVALIDATION_TARGETS: + return None + target: Final = cast(InvalidationTarget, raw_target) # cast-ok: membership checked against the Literal's args + return _CacheInvalidationMessage(cache_key=cache_key, cache=target) -async def publish_auth_cache_invalidation(cache_key: str) -> None: +async def publish_auth_cache_invalidation( + cache_key: str, + target: InvalidationTarget = DEFAULT_INVALIDATION_TARGET, +) -> None: """ Best-effort broadcast so every worker drops its local in-memory copy of a - mutated management object; without this, only the handling worker and Redis - are evicted and other workers keep serving the stale object until its TTL. + mutated object; without this, only the handling worker and Redis are + evicted and other workers keep serving the stale object until its TTL. """ redis_cache: Final = coordination_redis_cache() if redis_cache is None: @@ -67,21 +89,24 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: cache_key, ) return - await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + await client.publish( + auth_cache_invalidation_channel(redis_cache), + _cache_invalidation_message_json(cache_key=cache_key, target=target), + ) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) class AuthCacheInvalidationSubscriber: - __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ("_in_memory_caches", "_redis_cache", "_task") def __init__( self, redis_cache: "RedisCache", - user_api_key_cache: "UserApiKeyCache", + in_memory_caches: Mapping[InvalidationTarget, "InMemoryCache"], ) -> None: self._redis_cache = redis_cache - self._user_api_key_cache = user_api_key_cache + self._in_memory_caches = in_memory_caches self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -138,12 +163,12 @@ class AuthCacheInvalidationSubscriber: def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None - cache_key: Final = _cache_key_from_message_data(data) - if cache_key is None: + invalidation: Final = _invalidation_from_message_data(data) + if invalidation is None: return - in_memory_cache: Final = self._user_api_key_cache.in_memory_cache + in_memory_cache: Final = self._in_memory_caches.get(invalidation.cache) if in_memory_cache is not None: - in_memory_cache.delete_cache(cache_key) + in_memory_cache.delete_cache(invalidation.cache_key) @staticmethod async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index a51ff48aab6..9446d903441 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -10,6 +10,7 @@ All /customer management endpoints """ #### END-USER/CUSTOMER MANAGEMENT #### +from collections.abc import Iterable from datetime import datetime, timedelta from typing import Final @@ -21,6 +22,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import delete_cached_end_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.management_helpers.object_permission_utils import ( @@ -43,6 +45,13 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import ( router: Final = APIRouter() +async def _invalidate_cached_end_users(end_user_ids: Iterable[str]) -> None: + from litellm.proxy.proxy_server import user_api_key_cache + + for end_user_id in end_user_ids: + await delete_cached_end_user_object(end_user_id=end_user_id, user_api_key_cache=user_api_key_cache) + + def _to_customer_response(record: BaseModel) -> CustomerResponse: """Validate a raw end-user DB row into the typed customer response. @@ -96,6 +105,7 @@ async def block_user(data: BlockUsers): }, ) records.append(record) + await _invalidate_cached_end_users(data.user_ids) else: raise HTTPException( status_code=500, @@ -389,6 +399,8 @@ async def new_end_user( include={"litellm_budget_table": True, "object_permission": True}, ) + await _invalidate_cached_end_users((data.user_id,)) + return _to_customer_response(end_user_record) except Exception as e: verbose_proxy_logger.exception( @@ -632,6 +644,8 @@ async def update_end_user( raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + await _invalidate_cached_end_users((data.user_id,)) + return _to_customer_response(response) else: raise ValueError(f"user_id is required, passed user_id = {data.user_id}") @@ -705,6 +719,7 @@ async def delete_end_user( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) + await _invalidate_cached_end_users(data.user_ids) return DeleteCustomersResponse( deleted_customers=response, message="Successfully deleted customers with ids: " + str(data.user_ids), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d016eb57dd4..0b023df59ca 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6411,11 +6411,18 @@ class ProxyConfig: redis_cache: RedisCache | None, user_api_key_cache: UserApiKeyCache, ) -> None: + from litellm.proxy.utils import litellm_config_cache + if redis_cache is None or self.auth_cache_invalidation_subscriber is not None: return subscriber: Final = AuthCacheInvalidationSubscriber( redis_cache=redis_cache, - user_api_key_cache=user_api_key_cache, + in_memory_caches=MappingProxyType( + { + "user_api_key": user_api_key_cache.in_memory_cache, + "config_param": litellm_config_cache.in_memory_cache, + } + ), ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..7864f1977b7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -101,6 +101,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( @@ -2964,8 +2965,14 @@ async def evict_config_param(param_name: str) -> None: async def invalidate_config_param(param_name: str) -> None: - """Evict from both cache layers; call after every LiteLLM_Config write.""" + """Evict from both cache layers; call after every LiteLLM_Config write. + + ``evict_config_param`` only clears this worker's in-memory layer and the + shared Redis entry, so the broadcast is what stops the other workers from + serving their own in-memory copy of the pre-write value until it expires. + """ await evict_config_param(param_name) + await publish_auth_cache_invalidation(cache_key=_config_cache_key(param_name), target="config_param") await publish_config_param_change(param_name) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 298f8a31b64..f2ab1bc965b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5615,3 +5615,111 @@ def test_is_user_proxy_admin_rejects_view_only_admin(): assert _is_user_proxy_admin(user_obj=viewer) is False assert _is_user_proxy_admin(user_obj=admin) is True assert _is_user_proxy_admin(user_obj=None) is False + +class _RecordedInvalidations: + def __init__(self) -> None: + self.published: list = [] + + async def publish(self, cache_key: str, target: str = "user_api_key") -> None: + self.published.append((cache_key, target)) + + +@pytest.fixture +def recorded_invalidations(): + recorder = _RecordedInvalidations() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + recorder.publish, + ): + yield recorder + + +@pytest.mark.asyncio +async def test_delete_cache_key_object_broadcasts_revocation(recorded_invalidations): + """ + A revoked or regenerated key stays usable on every other worker until its + in-memory entry expires unless the eviction is broadcast. + """ + from litellm.proxy.auth.auth_checks import _delete_cache_key_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("hashed-token", {"key_name": "sk-1"}) + + await _delete_cache_key_object( + hashed_token="hashed-token", + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + + assert cache.in_memory_cache.get_cache("hashed-token") is None + assert recorded_invalidations.published == [("hashed-token", "user_api_key")] + + +@pytest.mark.asyncio +async def test_delete_cache_key_object_skips_broadcast_when_opted_out(recorded_invalidations): + """The expired-key auth path evicts on every worker on its own, so it must not + let a client replaying that key drive one broadcast per request.""" + from litellm.proxy.auth.auth_checks import _delete_cache_key_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + await _delete_cache_key_object( + hashed_token="hashed-token", + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=None, + broadcast_invalidation=False, + ) + + assert recorded_invalidations.published == [] + + +@pytest.mark.asyncio +async def test_delete_cache_access_object_broadcasts(recorded_invalidations): + from litellm.proxy.auth.auth_checks import _delete_cache_access_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("access_group_id:ag-1", {"models": []}) + + await _delete_cache_access_object(access_group_id="ag-1", user_api_key_cache=cache) + + assert cache.in_memory_cache.get_cache("access_group_id:ag-1") is None + assert recorded_invalidations.published == [("access_group_id:ag-1", "user_api_key")] + + +@pytest.mark.asyncio +async def test_delete_cached_end_user_object_evicts_and_broadcasts_both_entries(recorded_invalidations): + """ + The auth path caches the end-user row and, separately, the verdict on whether + the id resolves at all. Both survive a /customer mutation without this. + """ + from litellm.proxy.auth.auth_checks import delete_cached_end_user_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("end_user_id:eu-1", {"user_id": "eu-1", "blocked": False}) + cache.in_memory_cache.set_cache("end_user_validation:eu-1", "valid") + + await delete_cached_end_user_object(end_user_id="eu-1", user_api_key_cache=cache) + + assert cache.in_memory_cache.get_cache("end_user_id:eu-1") is None + assert cache.in_memory_cache.get_cache("end_user_validation:eu-1") is None + assert recorded_invalidations.published == [ + ("end_user_id:eu-1", "user_api_key"), + ("end_user_validation:eu-1", "user_api_key"), + ] + + +@pytest.mark.asyncio +async def test_delete_cached_end_user_object_still_broadcasts_when_eviction_fails(recorded_invalidations): + from litellm.proxy.auth.auth_checks import delete_cached_end_user_object + + cache = MagicMock() + cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + await delete_cached_end_user_object(end_user_id="eu-1", user_api_key_cache=cache) + + assert recorded_invalidations.published == [ + ("end_user_id:eu-1", "user_api_key"), + ("end_user_validation:eu-1", "user_api_key"), + ] 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 468e8aabae8..2911657fec1 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 @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest from redis.asyncio import Redis +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -69,8 +70,9 @@ class _FakeRedisCache: return self._client -def _invalidation_message(cache_key: str) -> dict: - return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()} +def _invalidation_message(cache_key: str, cache: Optional[str] = None) -> dict: + payload = {"cache_key": cache_key} if cache is None else {"cache_key": cache_key, "cache": cache} + return {"type": "message", "data": json.dumps(payload).encode()} @pytest.mark.asyncio @@ -82,7 +84,12 @@ async def test_publish_sends_cache_key_json_on_channel() -> None: ): await publish_auth_cache_invalidation(cache_key="project_id:p-1") - assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))] + assert client.published == [ + ( + AUTH_CACHE_INVALIDATION_CHANNEL, + json.dumps({"cache_key": "project_id:p-1", "cache": "user_api_key"}), + ) + ] @pytest.mark.asyncio @@ -129,7 +136,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) subscriber = AuthCacheInvalidationSubscriber( redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), - user_api_key_cache=cache, + in_memory_caches={"user_api_key": cache.in_memory_cache}, ) subscriber.start() try: @@ -151,11 +158,64 @@ async def test_subscriber_ignores_malformed_messages() -> None: subscriber = AuthCacheInvalidationSubscriber( redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), - user_api_key_cache=cache, + in_memory_caches={"user_api_key": cache.in_memory_cache}, ) subscriber._apply_message({"type": "message", "data": b"not json"}) subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) + subscriber._apply_message({"type": "message", "data": json.dumps({"cache_key": "project_id:p-1", "cache": "nope"}).encode()}) subscriber._apply_message("raw string") subscriber._apply_message(None) assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + + +@pytest.mark.asyncio +async def test_publish_targets_the_named_cache() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidation(cache_key="litellm_config:param:general_settings", target="config_param") + + assert json.loads(client.published[0][1]) == { + "cache_key": "litellm_config:param:general_settings", + "cache": "config_param", + } + + +@pytest.mark.asyncio +async def test_subscriber_routes_message_to_the_targeted_cache_only() -> None: + """ + A config-param broadcast must land on the config cache. Routing by target keeps + one worker's config write from evicting an unrelated, identically keyed auth entry. + """ + auth_cache = InMemoryCache() + config_cache = InMemoryCache() + shared_key = "litellm_config:param:general_settings" + auth_cache.set_cache(shared_key, {"from": "auth"}) + config_cache.set_cache(shared_key, {"from": "config"}) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + in_memory_caches={"user_api_key": auth_cache, "config_param": config_cache}, + ) + subscriber._apply_message(_invalidation_message(shared_key, cache="config_param")) + + assert config_cache.get_cache(shared_key) is None + assert auth_cache.get_cache(shared_key) is not None + + +@pytest.mark.asyncio +async def test_subscriber_defaults_untargeted_message_to_the_auth_cache() -> None: + """Messages published by a pre-upgrade worker carry no target during a rolling deploy.""" + auth_cache = InMemoryCache() + auth_cache.set_cache("project_id:p-1", {"models": []}) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + in_memory_caches={"user_api_key": auth_cache}, + ) + subscriber._apply_message(_invalidation_message("project_id:p-1")) + + assert auth_cache.get_cache("project_id:p-1") is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 0af5ad6cd9b..9b17d8ccc7d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -781,3 +781,82 @@ def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth): "deleted_customers": 2, "message": "Successfully deleted customers with ids: ['c1', 'c2']", } + +@pytest.fixture +def seeded_end_user_cache(): + """A cache already holding the auth-path entries for customer c1.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("end_user_id:c1", {"user_id": "c1", "blocked": False}) + cache.in_memory_cache.set_cache("end_user_validation:c1", "valid") + with patch("litellm.proxy.proxy_server.user_api_key_cache", cache): + yield cache + + +def _cached_end_user_entries(cache): + return ( + cache.in_memory_cache.get_cache("end_user_id:c1"), + cache.in_memory_cache.get_cache("end_user_validation:c1"), + ) + + +def test_block_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache): + """ + Without eviction the auth path keeps serving the pre-block end-user object, + so a blocked customer's requests still pass until the cache entry expires. + """ + mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock( + return_value=LiteLLM_EndUserTable(user_id="c1", blocked=True) + ) + + response = client.post("/customer/block", json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 200 + assert _cached_end_user_entries(seeded_end_user_cache) == (None, None) + + +def test_update_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache): + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=_row({"user_id": "c1", "blocked": False}) + ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + response = client.post( + "/customer/update", + json={"user_id": "c1", "alias": "Acme"}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200 + assert _cached_end_user_entries(seeded_end_user_cache) == (None, None) + + +def test_delete_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache): + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[LiteLLM_EndUserTable(user_id="c1", blocked=False)] + ) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1) + + response = client.post("/customer/delete", json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 200 + assert _cached_end_user_entries(seeded_end_user_cache) == (None, None) + + +def test_new_customer_evicts_cached_negative_validation(mock_prisma_client, mock_user_api_key_auth): + """ + The id-validation cache stores 'invalid' verdicts too, so creating a customer + for an id that was just rejected must drop that verdict. + """ + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("end_user_validation:c1", "invalid") + mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + with patch("litellm.proxy.proxy_server.user_api_key_cache", cache): + response = client.post("/customer/new", json={"user_id": "c1"}, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 200 + assert cache.in_memory_cache.get_cache("end_user_validation:c1") is None diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py index 761835078f4..e1583426ef1 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_config_param_cache.py @@ -265,3 +265,23 @@ async def test_prefetch_config_params_swallows_db_error_without_caching( prisma.db.litellm_config.find_many = AsyncMock(side_effect=RuntimeError("boom")) await prefetch_config_params(prisma, ["a", "b"]) assert _swap_config_cache._store == {} + +@pytest.mark.asyncio +async def test_invalidate_config_param_broadcasts_the_cache_key( + _swap_config_cache: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ + Evicting only this worker's copy leaves every other worker serving the + pre-write param value for the rest of the 60s TTL, so the write has to be + broadcast at the config cache. + """ + published: List[Any] = [] + + async def _record(cache_key: str, target: str = "user_api_key") -> None: + published.append((cache_key, target)) + + monkeypatch.setattr(utils_mod, "publish_auth_cache_invalidation", _record) + await invalidate_config_param("p5") + + assert published == [("litellm_config:param:p5", "config_param")]