From f54e4e664bd4c0bdf2677a162da616db67eb8ea7 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Sat, 11 Apr 2026 22:36:40 -0400 Subject: [PATCH] fix(proxy): use _hash_token_if_needed for cache invalidation in bulk update and key rotation (#25552) Two code paths in key_management_endpoints.py call hash_token() unconditionally when invalidating the user_api_key_cache after a key update. When the caller passes a pre-hashed token ID (not an sk- prefixed key), hash_token() double-hashes it, producing a cache key that does not match the actual cached entry. Cache invalidation silently fails. This is compounded by update_cache() which writes the stale cached key object back with a fresh 60s TTL after every successful request, preventing natural TTL expiry. The stale entry (with outdated fields like max_budget=None) persists indefinitely under load. PR #24969 fixed this in update_key_fn but missed two other call sites: - _process_single_key_update (bulk update path) - _execute_virtual_key_regeneration (key rotation path) Fix: replace hash_token() with _hash_token_if_needed() in both locations, matching the pattern already used elsewhere in the file. Co-authored-by: Claude Opus 4.6 (1M context) --- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 164 ++++++++++++++++++ 2 files changed, 166 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e8a691ce9..24396fbd18 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1849,7 +1849,7 @@ async def _process_single_key_update( # Delete cache await _delete_cache_key_object( - hashed_token=hash_token(key_update_item.key), + hashed_token=_hash_token_if_needed(key_update_item.key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -3721,7 +3721,7 @@ async def _execute_virtual_key_regeneration( if hashed_api_key or key: await _delete_cache_key_object( - hashed_token=hash_token(key), + hashed_token=_hash_token_if_needed(key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) 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 096e0b2bc4..a16bc078cf 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 @@ -8746,3 +8746,167 @@ def test_validate_public_image_url_accepts_http_and_noop_empty(): _validate_public_image_url(None, "logo_url") _validate_public_image_url("", "logo_url") _validate_public_image_url(" ", "logo_url") + + +@pytest.mark.asyncio +async def test_process_single_key_update_cache_invalidation_with_token_hash(): + """ + _process_single_key_update must pass the token hash as-is (not + double-hashed) to _delete_cache_key_object when the key is already a + pre-hashed token ID rather than an sk- prefixed key. + + Without this, cache invalidation silently fails: the wrong cache entry + is deleted while the stale entry (with outdated fields) persists and + gets refreshed indefinitely by update_cache on every successful request. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _process_single_key_update, + ) + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequestItem, + ) + + token_hash = "abc123def456" + + existing_key = LiteLLM_VerificationToken( + token=token_hash, + user_id="user-1", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=existing_key + ) + mock_updated = MagicMock() + mock_updated.model_dump.return_value = {"max_budget": 100.0} + mock_prisma_client.update_data = AsyncMock(return_value={"data": mock_updated}) + + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_llm_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data", + return_value={"max_budget": 100.0}, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + return_value=None, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) as mock_delete_cache, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ): + key_update_item = BulkUpdateKeyRequestItem( + key=token_hash, + max_budget=100.0, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + llm_router=mock_llm_router, + ) + + mock_delete_cache.assert_called_once() + call_kwargs = mock_delete_cache.call_args.kwargs + # The token hash should be passed as-is, NOT double-hashed + assert call_kwargs["hashed_token"] == token_hash + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_hash(): + """ + _execute_virtual_key_regeneration must pass the token hash as-is (not + double-hashed) to _delete_cache_key_object when the key is a + pre-hashed token ID. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + token_hash = "abc123def456" + + existing_key = LiteLLM_VerificationToken( + token=token_hash, + user_id="user-1", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + mock_prisma_client = AsyncMock() + # _execute_virtual_key_regeneration calls dict(updated_token) which + # needs the return value to be iterable as key-value pairs. + class DictLikeResult: + def __init__(self, data): + self._data = data + def __iter__(self): + return iter(self._data.items()) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=DictLikeResult({"token": "new-hashed-token", "key_name": "sk-...ab12", "user_id": "user-1"}) + ) + mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock( + return_value=None + ) + mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data) + + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) as mock_delete_cache, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data", + new_callable=AsyncMock, + return_value={}, + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key=token_hash, + key=token_hash, + data=None, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + ) + + mock_delete_cache.assert_called_once() + call_kwargs = mock_delete_cache.call_args.kwargs + # The token hash should be passed as-is, NOT double-hashed + assert call_kwargs["hashed_token"] == token_hash