From c67ca115ae50004248ffddc33e0b7a3956425b95 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 6 May 2025 20:45:15 -0700 Subject: [PATCH] Return 404 when `delete_verification_tokens` (`POST /key/delete`) fails to find keys (#10604) (#10605) * Return 404 when delete_verification_tokens fails to find keys * Add test_delete_nonexistent_key_returns_404 Co-authored-by: Marc Abramowitz --- .../key_management_endpoints.py | 6 ++++ .../test_key_generate_prisma.py | 35 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 89d0284750d..4ff0aa9c1b0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1486,6 +1486,12 @@ async def delete_verification_tokens( where={"token": {"in": tokens}} ) + if len(_keys_being_deleted) == 0: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": "No keys found"}, + ) + # Assuming 'db' is your Prisma Client instance # check if admin making request - don't filter by user-id if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 73128628aef..7aa0d3a8938 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -4056,3 +4056,38 @@ async def test_reset_budget_job(prisma_client, entity_type): assert entity_after is not None assert entity_after.spend == 0.0 + +def test_delete_nonexistent_key_returns_404(prisma_client): + # Try to delete a key that does not exist, expect a 404 error + import random, string + from litellm.proxy._types import KeyRequest, UserAPIKeyAuth, LitellmUserRoles, ProxyException + from litellm.proxy.management_endpoints.key_management_endpoints import delete_key_fn + from starlette.datastructures import URL + from fastapi import Request + + print("prisma client=", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + # Generate a random key that does not exist + random_key = "sk-" + ''.join(random.choices(string.ascii_letters + string.digits, k=24)) + delete_key_request = KeyRequest(keys=[random_key]) + bearer_token = "Bearer sk-1234" + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + # use admin to auth in + result = await litellm.proxy.proxy_server.user_api_key_auth(request=request, api_key=bearer_token) + result.user_role = LitellmUserRoles.PROXY_ADMIN + try: + await delete_key_fn(data=delete_key_request, user_api_key_dict=result) + pytest.fail("Expected ProxyException 404 for non-existent key, but delete_key_fn did not raise.") + except ProxyException as e: + print("Caught ProxyException:", e) + assert str(e.code) == "404" + assert "No keys found" in str(e.message) or "No matching keys or aliases found to delete" in str(e.message) + import asyncio + asyncio.run(test()) + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}")