From 8e89aa3c4dc2bea9ae7f0166a45d213faccefa23 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 3 Nov 2026 15:25:08 -0800 Subject: [PATCH] fix(proxy): evict already blocked keys again when POST /key/revoke sees them Cache eviction is best-effort, so a report could save the block while another worker kept a cached copy of the key. A later report skipped already blocked keys and never retried the eviction. Every matched key is now evicted on every report, and only newly blocked keys get the database update, audit row and warning --- .../key_revocation_endpoints.py | 13 +++++++------ .../test_key_revocation_endpoints.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_revocation_endpoints.py b/litellm/proxy/management_endpoints/key_revocation_endpoints.py index 237f68adba0..2c1be464f22 100644 --- a/litellm/proxy/management_endpoints/key_revocation_endpoints.py +++ b/litellm/proxy/management_endpoints/key_revocation_endpoints.py @@ -81,16 +81,17 @@ async def _block_live_keys( ) -> None: from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - newly_blocked: Final = tuple( - (row.token, row) for row in await tokens.find_by_tokens(hashed_tokens) if row.token and row.blocked is not True - ) - if not newly_blocked: - return + matched: Final = tuple((row.token, row) for row in await tokens.find_by_tokens(hashed_tokens) if row.token) + newly_blocked: Final = tuple((token, row) for token, row in matched if row.blocked is not True) blocked_tokens: Final = tuple(token for token, _ in newly_blocked) await tokens.block_tokens(blocked_tokens) await delete_cache_key_objects( - hashed_tokens=blocked_tokens, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj + hashed_tokens=tuple(token for token, _ in matched), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) + if not blocked_tokens: + return if is_audit_logging_enabled(): await audit(tuple(_audit_row(row, token) for token, row in newly_blocked)) verbose_proxy_logger.warning( diff --git a/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py index accd09d72d8..42b878263a9 100644 --- a/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py @@ -233,6 +233,23 @@ def test_an_already_blocked_key_gets_no_second_update_or_audit_row(monkeypatch, assert audit_rows == [] +def test_reporting_an_already_blocked_key_again_evicts_a_stale_cached_copy(monkeypatch, audit_rows, caplog): + table: Final = _install(monkeypatch, {BLOCKED_HASH: _row(BLOCKED_HASH, blocked=True)}) + cache: Final = proxy_server.user_api_key_cache + cache.set_cache(key=BLOCKED_HASH, value=UserAPIKeyAuth(token=BLOCKED_HASH, blocked=False).model_dump()) + assert cache.get_cache(key=BLOCKED_HASH) is not None + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + assert _revoke(BLOCKED) == ACCEPTED + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert cache.get_cache(key=BLOCKED_HASH) is None + assert (table.updates, audit_rows) == ([], []) + assert not any("Blocked" in record.getMessage() for record in caplog.records), caplog.records + + @pytest.mark.parametrize( "authorization", [f"Bearer {MASTER_KEY}", "Bearer sk-not-a-real-key", "Basic Zm9vOmJhcg==", "garbage"],