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"],