From c9a7e960074351bd1c889bf1bbe1e67eec20ddfb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:01:40 -0700 Subject: [PATCH] fix(key_management): evict every requested key on delete, including keys whose row is already gone A batch delete that names a live key and a key whose row was removed outside the API evicted only the live key, so the gone key kept authenticating from every worker's cache. Every requested key hash is evicted now, while raw non-key cache keys are still left alone --- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 46 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2eafeac8118..3021b21de51 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -79,7 +79,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, TEAM_OR_ORG_ADMIN, is_team_admin @@ -5569,7 +5569,7 @@ async def delete_verification_tokens( raise e await delete_cache_key_objects( - hashed_tokens=tuple(key.token for key in authorized_keys if key.token is not None), + hashed_tokens=tuple(token for token in hashed_tokens if is_user_key_cache_key(token)), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index a4368e547dc..5627ad2b492 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -5901,6 +5901,52 @@ async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeyp assert denied.value.code == "401" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deleter", + [ + UserAPIKeyAuth(user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value), + UserAPIKeyAuth(user_id="user-123", api_key="sk-owner", user_role=LitellmUserRoles.INTERNAL_USER.value), + ], + ids=["proxy_admin", "key_owner"], +) +async def test_delete_verification_tokens_evicts_a_requested_key_whose_row_is_already_gone( + monkeypatch: pytest.MonkeyPatch, deleter: UserAPIKeyAuth +) -> None: + """A key whose row was removed outside the API stays cached on every worker, so deleting it in a + batch beside a live key must still stop it everywhere instead of only evicting the rows found""" + live_key: Final = "sk-" + uuid4().hex + gone_key: Final = "sk-" + uuid4().hex + live_token: Final = hash_token(live_key) + gone_token: Final = hash_token(gone_key) + deleting_worker_cache: Final = UserApiKeyCache() + peer_worker_cache: Final = UserApiKeyCache() + for cache in (deleting_worker_cache, peer_worker_cache): + for token in (live_token, gone_token): + cache.set_cache(key=token, value=UserAPIKeyAuth(token=token, user_id="user-123")) + coordination_redis: Final = _SharedCoordinationRedis() + peer_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=coordination_redis, user_api_key_cache=peer_worker_cache + ) + peer_subscriber.start() + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_key_row(_owned_key_row(live_token))) + + try: + await delete_verification_tokens( + tokens=[live_key, gone_key], user_api_key_cache=deleting_worker_cache, user_api_key_dict=deleter + ) + await coordination_redis.until_every_subscriber_applied_its_messages() + finally: + await peer_subscriber.stop() + + assert [ + cache.get_cache(key=token) + for cache in (deleting_worker_cache, peer_worker_cache) + for token in (live_token, gone_token) + ] == [None, None, None, None] + + @pytest.mark.asyncio async def test_delete_verification_tokens_leaves_unowned_cache_entries_alone_on_every_worker(monkeypatch: pytest.MonkeyPatch) -> None: """/key/delete accepts any string, and a name that is not a key row passes through unhashed. A key