feat(ui): more workflow loading standardization

There is now a single entrypoint for loading a workflow - `useLoadWorkflowWithDialog`.

The hook:
Handles loading workflows from various sources. If there are unsaved changes, the user will be prompted to confirm before loading the workflow.

It returns  a function that:
Loads a workflow from various sources. If there are unsaved changes, the user will be prompted to confirm before loading the workflow. The workflow will be loaded immediately if there are no unsaved changes. On success, error or completion, the corresponding callback will be called.

WHEW

invokeai/frontend/web/src/features/workflowLibrary/components/LoadWorkflowConfirmationAlertDialog.tsx

@@ -1,81 +1,121 @@
 import { ConfirmationAlertDialog, Flex, Text } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { useAppSelector } from 'app/store/storeHooks';
 import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
 import { useWorkflowLibraryModal } from 'features/nodes/store/workflowLibraryModal';
-import { selectWorkflowIsTouched, workflowModeChanged } from 'features/nodes/store/workflowSlice';
+import { selectWorkflowIsTouched } from 'features/nodes/store/workflowSlice';
 import type { WorkflowV3 } from 'features/nodes/types/workflow';
 import { useLoadWorkflowFromFile } from 'features/workflowLibrary/hooks/useLoadWorkflowFromFile';
+import { useLoadWorkflowFromImage } from 'features/workflowLibrary/hooks/useLoadWorkflowFromImage';
 import { useLoadWorkflowFromLibrary } from 'features/workflowLibrary/hooks/useLoadWorkflowFromLibrary';
-import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
+import { useLoadWorkflowFromObject } from 'features/workflowLibrary/hooks/useLoadWorkflowFromObject';
 import { atom } from 'nanostores';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 
-type LoadLibraryWorkflowData = {
+type Callbacks = {
+  onSuccess?: (workflow: WorkflowV3) => void;
+  onError?: () => void;
+  onCompleted?: () => void;
+};
+
+type LoadLibraryWorkflowData = Callbacks & {
   type: 'library';
-  workflowId: string;
-  mode: 'view' | 'edit';
+  data: string;
 };
 
-type LoadDirectWorkflowData = {
+type LoadWorkflowFromObjectData = Callbacks & {
   type: 'direct';
-  workflow: WorkflowV3;
-  mode: 'view' | 'edit';
+  data: WorkflowV3;
 };
 
-type LoadFileWorkflowData = {
+type LoadWorkflowFromFileData = Callbacks & {
   type: 'file';
-  file: File;
-  mode: 'view' | 'edit';
+  data: File;
+};
+
+type LoadWorkflowFromImageData = Callbacks & {
+  type: 'image';
+  data: string;
+};
+
+type DialogStateExtra = {
+  isOpen: boolean;
 };
 
 const $dialogState = atom<
-  | (LoadLibraryWorkflowData & { isOpen: boolean })
-  | (LoadDirectWorkflowData & { isOpen: boolean })
-  | (LoadFileWorkflowData & { isOpen: boolean })
+  | (LoadLibraryWorkflowData & DialogStateExtra)
+  | (LoadWorkflowFromObjectData & DialogStateExtra)
+  | (LoadWorkflowFromFileData & DialogStateExtra)
+  | (LoadWorkflowFromImageData & DialogStateExtra)
   | null
 >(null);
 const cleanup = () => $dialogState.set(null);
 
-export const useLoadWorkflow = () => {
-  const dispatch = useAppDispatch();
+const useLoadImmediate = () => {
   const workflowLibraryModal = useWorkflowLibraryModal();
   const loadWorkflowFromLibrary = useLoadWorkflowFromLibrary();
   const loadWorkflowFromFile = useLoadWorkflowFromFile();
-  const validatedAndLoadWorkflow = useValidateAndLoadWorkflow();
-
-  const isTouched = useAppSelector(selectWorkflowIsTouched);
+  const loadWorkflowFromImage = useLoadWorkflowFromImage();
+  const loadWorkflowFromObject = useLoadWorkflowFromObject();
 
   const loadImmediate = useCallback(async () => {
-    const data = $dialogState.get();
-    if (!data) {
+    const dialogState = $dialogState.get();
+    if (!dialogState) {
       return;
     }
-    if (data.type === 'direct') {
-      const validatedWorkflow = await validatedAndLoadWorkflow(data.workflow);
-      if (validatedWorkflow) {
-        dispatch(workflowModeChanged(data.mode));
-      }
-    } else if (data.type === 'file') {
-      await loadWorkflowFromFile(data.file, {
-        onSuccess: () => {
-          dispatch(workflowModeChanged(data.mode));
-        },
-      });
-    } else {
-      await loadWorkflowFromLibrary(data.workflowId, {
-        onSuccess: () => {
-          dispatch(workflowModeChanged(data.mode));
-        },
-      });
+    const { type, data, onSuccess, onError, onCompleted } = dialogState;
+    const options = {
+      onSuccess,
+      onError,
+      onCompleted,
+    };
+    if (type === 'direct') {
+      await loadWorkflowFromObject(data, options);
+    } else if (type === 'file') {
+      await loadWorkflowFromFile(data, options);
+    } else if (type === 'library') {
+      await loadWorkflowFromLibrary(data, options);
+    } else if (type === 'image') {
+      await loadWorkflowFromImage(data, options);
     }
     cleanup();
     workflowLibraryModal.close();
-  }, [dispatch, loadWorkflowFromFile, loadWorkflowFromLibrary, validatedAndLoadWorkflow, workflowLibraryModal]);
+  }, [
+    loadWorkflowFromFile,
+    loadWorkflowFromImage,
+    loadWorkflowFromLibrary,
+    loadWorkflowFromObject,
+    workflowLibraryModal,
+  ]);
 
-  const loadWithDialog = useCallback(
-    (data: LoadLibraryWorkflowData | LoadDirectWorkflowData | LoadFileWorkflowData) => {
+  return loadImmediate;
+};
+
+/**
+ * Handles loading workflows from various sources. If there are unsaved changes, the user will be prompted to confirm
+ * before loading the workflow.
+ */
+export const useLoadWorkflowWithDialog = () => {
+  const isTouched = useAppSelector(selectWorkflowIsTouched);
+  const loadImmediate = useLoadImmediate();
+
+  const loadWorkflowWithDialog = useCallback(
+    /**
+     * Loads a workflow from various sources. If there are unsaved changes, the user will be prompted to confirm before
+     * loading the workflow. The workflow will be loaded immediately if there are no unsaved changes. On success, error
+     * or completion, the corresponding callback will be called.
+     *
+     * @param data - The data to load the workflow from.
+     * @param data.type - The type of data to load the workflow from.
+     * @param data.data - The data to load the workflow from. The type of this data depends on the `type` field.
+     * @param data.onSuccess - A callback to call when the workflow is successfully loaded.
+     * @param data.onError - A callback to call when an error occurs while loading the workflow.
+     * @param data.onCompleted - A callback to call when the loading process is completed (both success and error).
+     */
+    (
+      data: LoadLibraryWorkflowData | LoadWorkflowFromObjectData | LoadWorkflowFromFileData | LoadWorkflowFromImageData
+    ) => {
       if (!isTouched) {
         $dialogState.set({ ...data, isOpen: false });
         loadImmediate();
@@ -86,24 +126,21 @@ export const useLoadWorkflow = () => {
     [loadImmediate, isTouched]
   );
 
-  return {
-    loadImmediate,
-    loadWithDialog,
-  } as const;
+  return loadWorkflowWithDialog;
 };
 
 export const LoadWorkflowConfirmationAlertDialog = memo(() => {
   useAssertSingleton('LoadWorkflowConfirmationAlertDialog');
   const { t } = useTranslation();
   const workflow = useStore($dialogState);
-  const loadWorkflow = useLoadWorkflow();
+  const loadImmediate = useLoadImmediate();
 
   return (
     <ConfirmationAlertDialog
       isOpen={!!workflow?.isOpen}
       onClose={cleanup}
       title={t('nodes.loadWorkflow')}
-      acceptCallback={loadWorkflow.loadImmediate}
+      acceptCallback={loadImmediate}
       useInert={false}
       acceptButtonText={t('common.load')}
     >

invokeai/frontend/web/src/features/workflowLibrary/components/LoadWorkflowFromGraphModal/LoadWorkflowFromGraphModal.tsx

@@ -15,7 +15,7 @@ import {
 } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { graphToWorkflow } from 'features/nodes/util/workflow/graphToWorkflow';
-import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
+import { useLoadWorkflowFromObject } from 'features/workflowLibrary/hooks/useLoadWorkflowFromObject';
 import { atom } from 'nanostores';
 import type { ChangeEvent } from 'react';
 import { useCallback, useState } from 'react';
@@ -37,8 +37,8 @@ export const useLoadWorkflowFromGraphModal = () => {
 
 export const LoadWorkflowFromGraphModal = () => {
   const { t } = useTranslation();
-  const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
   const { isOpen, onClose } = useLoadWorkflowFromGraphModal();
+  const loadWorkflowFromObject = useLoadWorkflowFromObject();
   const [graphRaw, setGraphRaw] = useState<string>('');
   const [unvalidatedWorkflow, setUnvalidatedWorkflow] = useState<string>('');
   const [shouldAutoLayout, setShouldAutoLayout] = useState(true);
@@ -57,9 +57,9 @@ export const LoadWorkflowFromGraphModal = () => {
     setUnvalidatedWorkflow(JSON.stringify(workflow, null, 2));
   }, [graphRaw, shouldAutoLayout]);
   const loadWorkflow = useCallback(async () => {
-    await validateAndLoadWorkflow(unvalidatedWorkflow);
+    await loadWorkflowFromObject(unvalidatedWorkflow);
     onClose();
-  }, [validateAndLoadWorkflow, onClose, unvalidatedWorkflow]);
+  }, [loadWorkflowFromObject, unvalidatedWorkflow, onClose]);
   return (
     <Modal isOpen={isOpen} onClose={onClose} isCentered useInert={false}>
       <ModalOverlay />

invokeai/frontend/web/src/features/workflowLibrary/components/UploadWorkflowButton.tsx

@@ -1,5 +1,5 @@
 import { Button } from '@invoke-ai/ui-library';
-import { useLoadWorkflow } from 'features/workflowLibrary/components/LoadWorkflowConfirmationAlertDialog';
+import { useLoadWorkflowWithDialog } from 'features/workflowLibrary/components/LoadWorkflowConfirmationAlertDialog';
 import { memo, useCallback } from 'react';
 import { useDropzone } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
@@ -7,20 +7,19 @@ import { PiUploadSimpleBold } from 'react-icons/pi';
 
 export const UploadWorkflowButton = memo(() => {
   const { t } = useTranslation();
-  const loadWorkflow = useLoadWorkflow();
+  const loadWorkflowWithDialog = useLoadWorkflowWithDialog();
 
   const onDropAccepted = useCallback(
     ([file]: File[]) => {
       if (!file) {
         return;
       }
-      loadWorkflow.loadWithDialog({
+      loadWorkflowWithDialog({
         type: 'file',
-        file,
-        mode: 'edit',
+        data: file,
       });
     },
-    [loadWorkflow]
+    [loadWorkflowWithDialog]
   );
 
   const { getInputProps, getRootProps } = useDropzone({

invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/UploadWorkflowMenuItem.tsx

@@ -1,5 +1,5 @@
 import { MenuItem } from '@invoke-ai/ui-library';
-import { useLoadWorkflow } from 'features/workflowLibrary/components/LoadWorkflowConfirmationAlertDialog';
+import { useLoadWorkflowWithDialog } from 'features/workflowLibrary/components/LoadWorkflowConfirmationAlertDialog';
 import { memo, useCallback } from 'react';
 import { useDropzone } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
@@ -7,20 +7,19 @@ import { PiUploadSimpleBold } from 'react-icons/pi';
 
 const UploadWorkflowMenuItem = () => {
   const { t } = useTranslation();
-  const loadWorkflow = useLoadWorkflow();
+  const loadWorkflowWithDialog = useLoadWorkflowWithDialog();
 
   const onDropAccepted = useCallback(
     ([file]: File[]) => {
       if (!file) {
         return;
       }
-      loadWorkflow.loadWithDialog({
+      loadWorkflowWithDialog({
         type: 'file',
-        file,
-        mode: 'edit',
+        data: file,
       });
     },
-    [loadWorkflow]
+    [loadWorkflowWithDialog]
   );
 
   const { getRootProps, getInputProps } = useDropzone({

invokeai/frontend/web/src/features/workflowLibrary/hooks/useLoadWorkflowFromFile.tsx

@@ -4,6 +4,12 @@ import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useVa
 import { workflowLoadedFromFile } from 'features/workflowLibrary/store/actions';
 import { useCallback } from 'react';
 
+/**
+ * Loads a workflow from a file.
+ *
+ * You probably should instead use `useLoadWorkflowWithDialog`, which opens a dialog to prevent loss of unsaved changes
+ * and handles the loading process.
+ */
 export const useLoadWorkflowFromFile = () => {
   const dispatch = useAppDispatch();
   const validatedAndLoadWorkflow = useValidateAndLoadWorkflow();
@@ -13,30 +19,37 @@ export const useLoadWorkflowFromFile = () => {
       options: {
         onSuccess?: (workflow: WorkflowV3) => void;
         onError?: () => void;
+        onCompleted?: () => void;
       } = {}
     ) => {
-      const reader = new FileReader();
-      reader.onload = async () => {
-        const rawJSON = reader.result;
-        const { onSuccess, onError } = options;
-        try {
-          const unvalidatedWorkflow = JSON.parse(rawJSON as string);
-          const validatedWorkflow = await validatedAndLoadWorkflow(unvalidatedWorkflow);
+      return new Promise<WorkflowV3 | void>((resolve, reject) => {
+        const reader = new FileReader();
+        reader.onload = async () => {
+          const rawJSON = reader.result;
+          const { onSuccess, onError, onCompleted } = options;
+          try {
+            const unvalidatedWorkflow = JSON.parse(rawJSON as string);
+            const validatedWorkflow = await validatedAndLoadWorkflow(unvalidatedWorkflow);
 
-          if (!validatedWorkflow) {
-            reader.abort();
+            if (!validatedWorkflow) {
+              reader.abort();
+              onError?.();
+              return;
+            }
+            dispatch(workflowLoadedFromFile());
+            onSuccess?.(validatedWorkflow);
+            resolve(validatedWorkflow);
+          } catch {
+            // This is catching the error from the parsing the JSON file
             onError?.();
-            return;
+            reject();
+          } finally {
+            onCompleted?.();
           }
-          dispatch(workflowLoadedFromFile());
-          onSuccess?.(validatedWorkflow);
-        } catch {
-          // This is catching the error from the parsing the JSON file
-          onError?.();
-        }
-      };
+        };
 
-      reader.readAsText(file);
+        reader.readAsText(file);
+      });
     },
     [validatedAndLoadWorkflow, dispatch]
   );

invokeai/frontend/web/src/features/workflowLibrary/hooks/useLoadWorkflowFromImage.ts

@@ -8,19 +8,26 @@ import { useLazyGetImageWorkflowQuery } from 'services/api/endpoints/images';
 import type { NonNullableGraph } from 'services/api/types';
 import { assert } from 'tsafe';
 
+/**
+ * Loads a workflow from an image.
+ *
+ * You probably should instead use `useLoadWorkflowWithDialog`, which opens a dialog to prevent loss of unsaved changes
+ * and handles the loading process.
+ */
 export const useLoadWorkflowFromImage = () => {
   const { t } = useTranslation();
-  const [getWorkflowAndGraphFromImage, result] = useLazyGetImageWorkflowQuery();
+  const [getWorkflowAndGraphFromImage] = useLazyGetImageWorkflowQuery();
   const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
-  const getAndLoadEmbeddedWorkflow = useCallback(
+  const loadWorkflowFromImage = useCallback(
     async (
       imageName: string,
       options: {
         onSuccess?: (workflow: WorkflowV3) => void;
         onError?: () => void;
+        onCompleted?: () => void;
       } = {}
     ) => {
-      const { onSuccess, onError } = options;
+      const { onSuccess, onError, onCompleted } = options;
       try {
         const { workflow, graph } = await getWorkflowAndGraphFromImage(imageName).unwrap();
 
@@ -52,11 +59,12 @@ export const useLoadWorkflowFromImage = () => {
           status: 'error',
         });
         onError?.();
-        return;
+      } finally {
+        onCompleted?.();
       }
     },
     [getWorkflowAndGraphFromImage, validateAndLoadWorkflow, t]
   );
 
-  return [getAndLoadEmbeddedWorkflow, result] as const;
+  return loadWorkflowFromImage;
 };

invokeai/frontend/web/src/features/workflowLibrary/hooks/useLoadWorkflowFromLibrary.ts

@@ -5,6 +5,12 @@ import { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useLazyGetWorkflowQuery, useUpdateOpenedAtMutation, workflowsApi } from 'services/api/endpoints/workflows';
 
+/**
+ * Loads a workflow from the library.
+ *
+ * You probably should instead use `useLoadWorkflowWithDialog`, which opens a dialog to prevent loss of unsaved changes
+ * and handles the loading process.
+ */
 export const useLoadWorkflowFromLibrary = () => {
   const toast = useToast();
   const { t } = useTranslation();
@@ -17,9 +23,10 @@ export const useLoadWorkflowFromLibrary = () => {
       options: {
         onSuccess?: (workflow: WorkflowV3) => void;
         onError?: () => void;
+        onCompleted?: () => void;
       } = {}
     ) => {
-      const { onSuccess, onError } = options;
+      const { onSuccess, onError, onCompleted } = options;
       try {
         const res = await getWorkflow(workflowId).unwrap();
 
@@ -39,6 +46,8 @@ export const useLoadWorkflowFromLibrary = () => {
           status: 'error',
         });
         onError?.();
+      } finally {
+        onCompleted?.();
       }
     },
     [getWorkflow, validateAndLoadWorkflow, updateOpenedAt, toast, t]

invokeai/frontend/web/src/features/workflowLibrary/hooks/useLoadWorkflowFromObject.ts

@@ -0,0 +1,39 @@
+import type { WorkflowV3 } from 'features/nodes/types/workflow';
+import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
+import { useCallback } from 'react';
+
+/**
+ * Loads a workflow from an object.
+ *
+ * You probably should instead use `useLoadWorkflowWithDialog`, which opens a dialog to prevent loss of unsaved changes
+ * and handles the loading process.
+ */
+export const useLoadWorkflowFromObject = () => {
+  const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
+  const loadWorkflowFromObject = useCallback(
+    async (
+      unvalidatedWorkflow: unknown,
+      options: {
+        onSuccess?: (workflow: WorkflowV3) => void;
+        onError?: () => void;
+        onCompleted?: () => void;
+      } = {}
+    ) => {
+      const { onSuccess, onError, onCompleted } = options;
+      try {
+        const validatedWorkflow = await validateAndLoadWorkflow(unvalidatedWorkflow);
+
+        if (!validatedWorkflow) {
+          onError?.();
+          return;
+        }
+        onSuccess?.(validatedWorkflow);
+      } finally {
+        onCompleted?.();
+      }
+    },
+    [validateAndLoadWorkflow]
+  );
+
+  return loadWorkflowFromObject;
+};

invokeai/frontend/web/src/features/workflowLibrary/hooks/useValidateAndLoadWorkflow.ts

@@ -17,6 +17,21 @@ import { fromZodError } from 'zod-validation-error';
 
 const log = logger('workflows');
 
+/**
+ * This hook manages the lower-level workflow validation and loading process.
+ *
+ * You probably should instead use `useLoadWorkflowWithDialog`, which opens a dialog to prevent loss of unsaved changes
+ * and handles the loading process.
+ *
+ * Internally, `useLoadWorkflowWithDialog` uses these hooks...
+ *
+ * - `useLoadWorkflowFromFile`
+ * - `useLoadWorkflowFromImage`
+ * - `useLoadWorkflowFromLibrary`
+ * - `useLoadWorkflowFromObject`
+ *
+ * ...each of which internally uses hook.
+ */
 export const useValidateAndLoadWorkflow = () => {
   const { t } = useTranslation();
   const dispatch = useAppDispatch();