fix: add key_rotation_email to generate_key_helper_fn, KeyResponse type, and SQL migration

litellm-proxy-extras/litellm_proxy_extras/migrations/20260331000000_add_key_rotation_email/migration.sql

@@ -0,0 +1,7 @@
+-- Add key_rotation_email to LiteLLM_VerificationToken and LiteLLM_DeletedVerificationToken
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken"
+  ADD COLUMN IF NOT EXISTS "key_rotation_email" TEXT;
+
+ALTER TABLE "LiteLLM_DeletedVerificationToken"
+  ADD COLUMN IF NOT EXISTS "key_rotation_email" TEXT;

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -732,9 +732,9 @@ async def _common_key_generation_helper(  # noqa: PLR0915
         request_type="key", **data_json, table_name="key"
     )
 
-    response[
-        "soft_budget"
-    ] = data.soft_budget  # include the user-input soft budget in the response
+    response["soft_budget"] = (
+        data.soft_budget
+    )  # include the user-input soft budget in the response
 
     response = GenerateKeyResponse(**response)
 
@@ -2783,6 +2783,7 @@ async def generate_key_helper_fn(  # noqa: PLR0915
     object_permission: Optional[LiteLLM_ObjectPermissionBase] = None,
     auto_rotate: Optional[bool] = None,
     rotation_interval: Optional[str] = None,
+    key_rotation_email: Optional[str] = None,
     router_settings: Optional[dict] = None,
     access_group_ids: Optional[list] = None,
 ):
@@ -2913,6 +2914,9 @@ async def generate_key_helper_fn(  # noqa: PLR0915
             rotation_interval=rotation_interval,
         )
 
+        if key_rotation_email is not None:
+            key_data["key_rotation_email"] = key_rotation_email
+
         if (
             get_secret("DISABLE_KEY_NAME", False) is True
         ):  # allow user to disable storing abbreviated key name (shown in UI, to help figure out which key spent how much)
@@ -3175,10 +3179,10 @@ async def delete_verification_tokens(
     try:
         if prisma_client:
             tokens = [_hash_token_if_needed(token=key) for key in tokens]
-            _keys_being_deleted: List[
-                LiteLLM_VerificationToken
-            ] = await prisma_client.db.litellm_verificationtoken.find_many(
-                where={"token": {"in": tokens}}
+            _keys_being_deleted: List[LiteLLM_VerificationToken] = (
+                await prisma_client.db.litellm_verificationtoken.find_many(
+                    where={"token": {"in": tokens}}
+                )
             )
 
             if len(_keys_being_deleted) == 0:
@@ -3378,9 +3382,9 @@ async def _rotate_master_key(  # noqa: PLR0915
     from litellm.proxy.proxy_server import proxy_config
 
     try:
-        models: Optional[
-            List
-        ] = await prisma_client.db.litellm_proxymodeltable.find_many()
+        models: Optional[List] = (
+            await prisma_client.db.litellm_proxymodeltable.find_many()
+        )
     except Exception:
         models = None
     # 2. process model table
@@ -4020,11 +4024,11 @@ async def validate_key_list_check(
             param="user_id",
             code=status.HTTP_403_FORBIDDEN,
         )
-    complete_user_info_db_obj: Optional[
-        BaseModel
-    ] = await prisma_client.db.litellm_usertable.find_unique(
-        where={"user_id": user_api_key_dict.user_id},
-        include={"organization_memberships": True},
+    complete_user_info_db_obj: Optional[BaseModel] = (
+        await prisma_client.db.litellm_usertable.find_unique(
+            where={"user_id": user_api_key_dict.user_id},
+            include={"organization_memberships": True},
+        )
     )
 
     if complete_user_info_db_obj is None:
@@ -4107,10 +4111,10 @@ async def _fetch_user_team_objects(
     if complete_user_info is None or not complete_user_info.teams:
         return []
 
-    teams: Optional[
-        List[BaseModel]
-    ] = await prisma_client.db.litellm_teamtable.find_many(
-        where={"team_id": {"in": complete_user_info.teams}}
+    teams: Optional[List[BaseModel]] = (
+        await prisma_client.db.litellm_teamtable.find_many(
+            where={"team_id": {"in": complete_user_info.teams}}
+        )
     )
     if teams is None:
         return []

ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx

@@ -93,6 +93,7 @@ export interface KeyResponse {
   access_group_ids?: string[];
   auto_rotate?: boolean;
   rotation_interval?: string;
+  key_rotation_email?: string;
   last_rotation_at?: string;
   key_rotation_at?: string;
   next_rotation_at?: string;

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -101,7 +101,7 @@ export function KeyEditView({
   const [selectedOrganizationId, setSelectedOrganizationId] = useState<string | null>(keyData.organization_id || null);
   const [autoRotationEnabled, setAutoRotationEnabled] = useState<boolean>(keyData.auto_rotate || false);
   const [rotationInterval, setRotationInterval] = useState<string>(keyData.rotation_interval || "");
-  const [keyRotationEmail, setKeyRotationEmail] = useState<string>((keyData as any).key_rotation_email || "");
+  const [keyRotationEmail, setKeyRotationEmail] = useState<string>(keyData.key_rotation_email || "");
   const [neverExpire, setNeverExpire] = useState<boolean>(!keyData.expires);
   const [isKeySaving, setIsKeySaving] = useState(false);
   const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();