fix(key management): populate failed_tokens for admin partial-delete path

The admin bulk-delete path (prisma_client.delete_data) was never comparing
the DB return value against the requested tokens, so failed_tokens was
always empty for admins even when the DB silently skipped some tokens.

Adds the same mismatch check as the non-admin path, and a new test that
exercises admin bulk-delete returning fewer tokens than requested.

Fixes gap identified by Greptile review on #21609.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros Pro 2026-02-19 19:28:03 -03:00
parent 079bc364bb
commit b9f36645f3
2 changed files with 51 additions and 0 deletions

View File

@ -2998,6 +2998,10 @@ async def delete_verification_tokens(
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
failed_tokens = [
token for token in tokens if token not in deleted_tokens
]
else:
deletion_tasks = [
prisma_client.delete_data(tokens=[key.token])

View File

@ -208,3 +208,50 @@ async def test_delete_tokens_non_admin_token_not_in_db_returns_failed_tokens(
"token-2 was not found in the DB and must appear in failed_tokens"
)
assert "hashed-token-1" in result["deleted_keys"]
# ---------------------------------------------------------------------------
# Test 4 admin, DB bulk-delete returns fewer tokens → failed_tokens populated
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_delete_tokens_admin_partial_db_failure_returns_failed_tokens(
monkeypatch,
):
"""
PROXY_ADMIN requests deletion of two tokens; the DB bulk-delete only
removes one (e.g. the other was concurrently deleted). The unremoved
token must appear in `failed_tokens` previously it would be silently
swallowed since the admin path never compared returned vs. requested counts.
"""
key1 = _make_token("hashed-token-1")
key2 = _make_token("hashed-token-2")
# DB reports only token-1 as deleted
mock_prisma = _mock_prisma(
keys=[key1, key2],
deleted_tokens=["hashed-token-1"],
)
mock_cache = MagicMock()
mock_cache.delete_cache = MagicMock()
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)
result, _ = await delete_verification_tokens(
tokens=["hashed-token-1", "hashed-token-2"],
user_api_key_cache=mock_cache,
user_api_key_dict=_admin_user(),
)
assert "failed_tokens" in result
assert "hashed-token-2" in result["failed_tokens"], (
"token-2 was not deleted by the DB and must appear in failed_tokens for admins too"
)
assert "hashed-token-1" in result["deleted_keys"]