Merge pull request #9257 from BerriAI/litellm_ui_fixes_key_create

[UI] Fix 1 - instantly show newly create keys on Admin UI (don't require refresh)

ui/litellm-dashboard/src/components/all_keys_table.tsx

@@ -32,6 +32,7 @@ interface AllKeysTableProps {
   userRole: string | null;
   organizations: Organization[] | null;
   setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
+  refresh?: () => void;
 }
 
 // Define columns similar to our logs table
@@ -98,6 +99,7 @@ export function AllKeysTable({
   userRole,
   organizations,
   setCurrentOrg,
+  refresh,
 }: AllKeysTableProps) {
   const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
   const [userList, setUserList] = useState<UserResponse[]>([]);
@@ -131,6 +133,22 @@ export function AllKeysTable({
     }
   }, [accessToken, keys]);
 
+  // Add a useEffect to call refresh when a key is created
+  useEffect(() => {
+    if (refresh) {
+      const handleStorageChange = () => {
+        refresh();
+      };
+      
+      // Listen for storage events that might indicate a key was created
+      window.addEventListener('storage', handleStorageChange);
+      
+      return () => {
+        window.removeEventListener('storage', handleStorageChange);
+      };
+    }
+  }, [refresh]);
+
   const columns: ColumnDef<KeyResponse>[] = [
     {
       id: "expander",

ui/litellm-dashboard/src/components/create_key_button.tsx

@@ -264,12 +264,21 @@ 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);
+
     } catch (error) {
       console.log("error in create key:", error);
       message.error(`Error creating the key: ${error}`);

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
     };
 };
 

ui/litellm-dashboard/src/components/view_key_table.tsx

@@ -176,12 +176,21 @@ 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 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) => {
     refresh({ page: newPage });
   };
@@ -421,6 +430,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         userRole={userRole}
         organizations={organizations}
         setCurrentOrg={setCurrentOrg}
+        refresh={refresh}
       />
 
       {isDeleteModalOpen && (
@@ -619,4 +629,12 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   );
 };
 
+// Update the type declaration to include the new function
+declare global {
+  interface Window {
+    refreshKeysList?: () => void;
+    addNewKeyToList?: (newKey: any) => void;
+  }
+}
+
 export default ViewKeyTable;