diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index aef2fb3930..4688445fee 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1027,10 +1027,15 @@ async def regenerate_key_fn( new_token = f"sk-{secrets.token_urlsafe(16)}" new_token_hash = hash_token(new_token) + new_token_key_name = f"sk-...{new_token[-4:]}" # update new token in DB updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, data={"token": new_token_hash} + where={"token": hashed_api_key}, + data={ + "token": new_token_hash, + "key_name": new_token_key_name, + }, ) updated_token_dict = {} if updated_token is not None: diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index e41a67c08b..49a4c95c75 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -3035,4 +3035,6 @@ async def test_regenerate_api_key(prisma_client): new_key.models == models ), f"Expected models {models} but got {new_key.models}" + assert new_key.key_name == f"sk-...{new_key.key[-4:]}" + pass