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:
parent
079bc364bb
commit
b9f36645f3
@ -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])
|
||||
|
||||
@ -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"]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user