From eec8e1a746154e9a9465f0abeb9d166d208b602f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 26 Jun 2025 15:50:08 -0700
Subject: [PATCH] fix(key_management_endpoints.py): add new param `new_key` for
 setting the regenerated key value

user request
---
 litellm/proxy/_types.py                         |  3 +--
 .../key_management_endpoints.py                 | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 63937297aa9..db0959a1d43 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -351,13 +351,11 @@ class LiteLLMRoutes(enum.Enum):
         "/v1beta/models/{model_name}:countTokens",
         "/v1beta/models/{model_name}:generateContent",
         "/v1beta/models/{model_name}:streamGenerateContent",
-
         "/models/{model_name}:countTokens",
         "/models/{model_name}:generateContent",
         "/models/{model_name}:streamGenerateContent",
     ]
 
-
     apply_guardrail_routes = [
         "/guardrails/apply_guardrail",
     ]
@@ -789,6 +787,7 @@ class UpdateKeyRequest(KeyRequestBase):
 class RegenerateKeyRequest(GenerateKeyRequest):
     # This needs to be different from UpdateKeyRequest, because "key" is optional for this
     key: Optional[str] = None
+    new_key: Optional[str] = None
     duration: Optional[str] = None
     spend: Optional[float] = None
     metadata: Optional[dict] = None
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 0b1c7f523c0..ab78bd1d977 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -696,6 +696,7 @@ async def prepare_key_update_data(
 ):
     data_json: dict = data.model_dump(exclude_unset=True)
     data_json.pop("key", None)
+    data_json.pop("new_key", None)
     non_default_values = {}
     for k, v in data_json.items():
         if k in LiteLLM_ManagementEndpoint_MetadataFields:
@@ -1860,6 +1861,9 @@ async def regenerate_key_fn(
     Parameters:
     - key: str (path parameter) - The key to regenerate
     - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
+        - key: Optional[str] - The key to regenerate.
+        - new_master_key: Optional[str] - The new master key to use, if key is the master key.
+        - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used.
         - key_alias: Optional[str] - User-friendly key alias
         - user_id: Optional[str] - User ID associated with key
         - team_id: Optional[str] - Team ID associated with key
@@ -1983,7 +1987,18 @@ async def regenerate_key_fn(
 
         verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
 
-        new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
+        if data and data.new_key is not None:
+            new_token = data.new_key
+            if not data.new_key.startswith("sk-"):
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={
+                        "error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key."
+                    },
+                )
+        else:
+            new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
+
         new_token_hash = hash_token(new_token)
         new_token_key_name = f"sk-...{new_token[-4:]}"