From d54071b04629342fb688bfe4fb49b851ec29dba2 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:53:20 +0000 Subject: [PATCH] fix(proxy): broadcast key eviction from /key/delete too Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 8 +++- .../test_key_management_endpoints.py | 46 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a5078c50fc0..53a99c917d2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4061,7 +4061,7 @@ async def delete_verification_tokens( - List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted, this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj failed_tokens: list = [] try: @@ -4125,7 +4125,11 @@ async def delete_verification_tokens( user_api_key_cache.delete_cache(key) # remove hash token from cache hashed_token = hash_token(cast(str, key)) - user_api_key_cache.delete_cache(hashed_token) + await _delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) return { "deleted_keys": deleted_tokens, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e8709f3af34..940a333779c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15442,3 +15442,49 @@ async def test_migrate_encryption_endpoint_rejects_proxy_admin_viewer(): assert exc_info.value.status_code == 403 mock_migrate.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_verification_tokens_broadcasts_key_eviction(monkeypatch): + """ + /key/delete evicting only the handling worker's copy leaves the deleted key + authenticating requests on every other worker until its cache entry expires. + """ + published: list = [] + + async def _record(cache_key: str, target: str = "user_api_key") -> None: + published.append((cache_key, target)) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-token-1")] + ) + mock_prisma_client.delete_data = AsyncMock(return_value=["hashed-token-1"]) + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", + lambda token: token, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.hash_token", + lambda token: token, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + _record, + ) + + await delete_verification_tokens( + tokens=["hashed-token-1"], + user_api_key_cache=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert published == [("hashed-token-1", "user_api_key")]