diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2884c8d374e..7be1a5a5e0c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1654,7 +1654,7 @@ async def _get_and_validate_existing_key( LiteLLM_VerificationToken: The existing key row Raises: - HTTPException: If key is not found + ProxyException: 404 if key is not found """ if prisma_client is None: raise HTTPException( @@ -1662,12 +1662,7 @@ async def _get_and_validate_existing_key( detail={"error": "Database not connected"}, ) - from litellm.proxy.proxy_server import hash_token - - if token.startswith("sk-"): - hashed_token = hash_token(token=token) - else: - hashed_token = token + hashed_token = _hash_token_if_needed(token=token) existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( where={"token": hashed_token} @@ -2116,28 +2111,11 @@ async def update_key_fn( key = data_json.pop("key") # get the row from db - if prisma_client is None: - raise Exception("Not connected to DB!") - - from litellm.proxy.proxy_server import hash_token - - if data.key.startswith("sk-"): - hashed_token = hash_token(token=data.key) - else: - hashed_token = data.key - - existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} + existing_key_row = await _get_and_validate_existing_key( + token=data.key, + prisma_client=prisma_client, ) - if existing_key_row is None: - raise ProxyException( - message=f"Key not found. Passed key={data.key}", - type=ProxyErrorTypes.not_found_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) - await _validate_update_key_data( data=data, existing_key_row=existing_key_row, 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 3bff35fbd55..48a1f7936c5 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 @@ -1678,6 +1678,56 @@ async def test_unblock_key_nonexistent_key_returns_404(monkeypatch): mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called() +@pytest.mark.asyncio +async def test_update_key_nonexistent_key_returns_404(monkeypatch): + """ + Test that update_key_fn returns 404 (not misleading 401) when the body + key doesn't exist in the database, even when the caller is authenticated + as a proxy admin via the Authorization header. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + # find_unique returns None → key does not exist + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user" + ) + + data = UpdateKeyRequest(key="sk-does-not-exist-key") + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=mock_request, + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert exc_info.value.code == "404" + assert "not found" in str(exc_info.value.message).lower() + assert "Authentication Error" not in str(exc_info.value.message) + + @pytest.mark.asyncio async def test_block_key_existing_key_succeeds(monkeypatch): """