From 43b7460968ee8fcbdfce692a15ddb9b0018d0207 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 14 Mar 2025 09:18:22 -0700
Subject: [PATCH] ui - instantly show new ui keys

---
 .../src/components/create_key_button.tsx      | 18 +++++++++++----
 .../components/key_team_helpers/key_list.tsx  | 23 +++++++++++++++----
 .../src/components/view_key_table.tsx         | 11 ++++++---
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx
index a641596cb3..cce52ab3d9 100644
--- a/ui/litellm-dashboard/src/components/create_key_button.tsx
+++ b/ui/litellm-dashboard/src/components/create_key_button.tsx
@@ -260,17 +260,25 @@ const CreateKey: React.FC<CreateKeyProps> = ({
       const response = await keyCreateCall(accessToken, userID, formValues);
 
       console.log("key create Response:", response);
-      setData((prevData) => (prevData ? [...prevData, response] : [response])); // Check if prevData is null
+      
+      // Update the data state in this component
+      setData((prevData) => (prevData ? [...prevData, response] : [response]));
+      
+      // Also directly update the keys list in AllKeysTable without an API call
+      if (window.addNewKeyToList) {
+        window.addNewKeyToList(response);
+      }
+      
       setApiKey(response["key"]);
       setSoftBudget(response["soft_budget"]);
       message.success("API Key Created");
       form.resetFields();
       localStorage.removeItem("userData" + userID);
       
-      // Add this line to refresh the keys list immediately
-      if (window.refreshKeysList) {
-        window.refreshKeysList();
-      }
+      // We don't need to call refresh anymore since we're directly updating the state
+      // if (window.refreshKeysList) {
+      //   window.refreshKeysList();
+      // }
     } catch (error) {
       console.log("error in create key:", error);
       message.error(`Error creating the key: ${error}`);
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
index fd9d1ef245..4c2a18d2b5 100644
--- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
@@ -100,6 +100,7 @@ isLoading: boolean;
 error: Error | null;
 pagination: PaginationData;
 refresh: (params?: Record<string, unknown>) => Promise<void>;
+setKeys: (newKeysOrUpdater: KeyResponse[] | ((prevKeys: KeyResponse[]) => KeyResponse[])) => void;
 }
 
 const useKeyList = ({
@@ -149,16 +150,30 @@ const useKeyList = ({
         console.log("selectedTeam", selectedTeam, "currentOrg", currentOrg, "accessToken", accessToken);
     }, [selectedTeam, currentOrg, accessToken]);
 
+    const setKeys = (newKeysOrUpdater: KeyResponse[] | ((prevKeys: KeyResponse[]) => KeyResponse[])) => {
+        setKeyData(prevData => {
+            const newKeys = typeof newKeysOrUpdater === 'function' 
+                ? newKeysOrUpdater(prevData.keys) 
+                : newKeysOrUpdater;
+                
+            return {
+                ...prevData,
+                keys: newKeys
+            };
+        });
+    };
+
     return {
         keys: keyData.keys,
         isLoading,
         error,
         pagination: {
-        currentPage: keyData.current_page,
-        totalPages: keyData.total_pages,
-        totalCount: keyData.total_count
+            currentPage: keyData.current_page,
+            totalPages: keyData.total_pages,
+            totalCount: keyData.total_count
         },
-        refresh: fetchKeys
+        refresh: fetchKeys,
+        setKeys
     };
 };
 
diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx
index 968b7cf4c8..f3661c8c64 100644
--- a/ui/litellm-dashboard/src/components/view_key_table.tsx
+++ b/ui/litellm-dashboard/src/components/view_key_table.tsx
@@ -176,15 +176,19 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   // Build a memoized filters object for the backend call.
 
   // Pass filters into the hook so the API call includes these query parameters.
-  const { keys, isLoading, error, pagination, refresh } = useKeyList({
+  const { keys, isLoading, error, pagination, refresh, setKeys } = useKeyList({
     selectedTeam,
     currentOrg,
     accessToken,
   });
 
-  // Make refresh function available globally so CreateKey can access it
+  // Make both refresh and addKey functions available globally
   if (typeof window !== 'undefined') {
     window.refreshKeysList = refresh;
+    window.addNewKeyToList = (newKey) => {
+      // Add the new key to the keys list without making an API call
+      setKeys((prevKeys) => [newKey, ...prevKeys]);
+    };
   }
 
   const handlePageChange = (newPage: number) => {
@@ -625,10 +629,11 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   );
 };
 
-// Add this type declaration at the top of the file to avoid TypeScript errors
+// Update the type declaration to include the new function
 declare global {
   interface Window {
     refreshKeysList?: () => void;
+    addNewKeyToList?: (newKey: any) => void;
   }
 }