undo/redo handles multiple canvases

Author: Attila Cseh

invokeai/frontend/web/src/app/store/store.ts

@@ -22,15 +22,15 @@ import { merge } from 'es-toolkit';
 import { omit, pick } from 'es-toolkit/compat';
 import { changeBoardModalSliceConfig } from 'features/changeBoardModal/store/slice';
 import { canvasSettingsSliceConfig } from 'features/controlLayers/store/canvasSettingsSlice';
-import { canvasSliceConfig } from 'features/controlLayers/store/canvasSlice';
+import { canvasSliceConfig, undoableCanvasSliceReducer } from 'features/controlLayers/store/canvasSlice';
 import { canvasSessionSliceConfig } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { lorasSliceConfig } from 'features/controlLayers/store/lorasSlice';
 import { paramsSliceConfig } from 'features/controlLayers/store/paramsSlice';
 import { refImagesSliceConfig } from 'features/controlLayers/store/refImagesSlice';
 import { dynamicPromptsSliceConfig } from 'features/dynamicPrompts/store/dynamicPromptsSlice';
 import { gallerySliceConfig } from 'features/gallery/store/gallerySlice';
 import { modelManagerSliceConfig } from 'features/modelManagerV2/store/modelManagerV2Slice';
-import { nodesSliceConfig } from 'features/nodes/store/nodesSlice';
+import { nodesSliceConfig, undoableNodesSliceReducer } from 'features/nodes/store/nodesSlice';
 import { workflowLibrarySliceConfig } from 'features/nodes/store/workflowLibrarySlice';
 import { workflowSettingsSliceConfig } from 'features/nodes/store/workflowSettingsSlice';
 import { upscaleSliceConfig } from 'features/parameters/store/upscaleSlice';
@@ -44,7 +44,6 @@ import { diff } from 'jsondiffpatch';
 import dynamicMiddlewares from 'redux-dynamic-middlewares';
 import type { SerializeFunction, UnserializeFunction } from 'redux-remember';
 import { REMEMBER_REHYDRATED, rememberEnhancer, rememberReducer } from 'redux-remember';
-import undoable, { newHistory } from 'redux-undo';
 import { serializeError } from 'serialize-error';
 import { api } from 'services/api';
 import { authToastMiddleware } from 'services/api/authToastMiddleware';
@@ -91,22 +90,14 @@ const ALL_REDUCERS = {
   [api.reducerPath]: api.reducer,
   [canvasSessionSliceConfig.slice.reducerPath]: canvasSessionSliceConfig.slice.reducer,
   [canvasSettingsSliceConfig.slice.reducerPath]: canvasSettingsSliceConfig.slice.reducer,
-  // Undoable!
-  [canvasSliceConfig.slice.reducerPath]: undoable(
-    canvasSliceConfig.slice.reducer,
-    canvasSliceConfig.undoableConfig?.reduxUndoOptions
-  ),
+  [canvasSliceConfig.slice.reducerPath]: undoableCanvasSliceReducer,
   [changeBoardModalSliceConfig.slice.reducerPath]: changeBoardModalSliceConfig.slice.reducer,
   [configSliceConfig.slice.reducerPath]: configSliceConfig.slice.reducer,
   [dynamicPromptsSliceConfig.slice.reducerPath]: dynamicPromptsSliceConfig.slice.reducer,
   [gallerySliceConfig.slice.reducerPath]: gallerySliceConfig.slice.reducer,
   [lorasSliceConfig.slice.reducerPath]: lorasSliceConfig.slice.reducer,
   [modelManagerSliceConfig.slice.reducerPath]: modelManagerSliceConfig.slice.reducer,
-  // Undoable!
-  [nodesSliceConfig.slice.reducerPath]: undoable(
-    nodesSliceConfig.slice.reducer,
-    nodesSliceConfig.undoableConfig?.reduxUndoOptions
-  ),
+  [nodesSliceConfig.slice.reducerPath]: undoableNodesSliceReducer,
   [paramsSliceConfig.slice.reducerPath]: paramsSliceConfig.slice.reducer,
   [queueSliceConfig.slice.reducerPath]: queueSliceConfig.slice.reducer,
   [refImagesSliceConfig.slice.reducerPath]: refImagesSliceConfig.slice.reducer,
@@ -162,7 +153,7 @@ const unserialize: UnserializeFunction = (data, key) => {
 
   // Undoable slices must be wrapped in a history!
   if (undoableConfig) {
-    return newHistory([], state, []);
+    return undoableConfig.wrapState(state);
   } else {
     return state;
   }
@@ -175,7 +166,7 @@ const serialize: SerializeFunction = (data, key) => {
   }
 
   const result = omit(
-    sliceConfig.undoableConfig ? data.present : data,
+    sliceConfig.undoableConfig ? sliceConfig.undoableConfig.unwrapState(data) : data,
     sliceConfig.persistConfig.persistDenylist ?? []
   );
 

invokeai/frontend/web/src/app/store/types.ts

@@ -1,10 +1,9 @@
 import type { Slice } from '@reduxjs/toolkit';
-import type { UndoableOptions } from 'redux-undo';
 import type { ZodType } from 'zod';
 
 type StateFromSlice<T extends Slice> = T extends Slice<infer U> ? U : never;
 
-export type SliceConfig<T extends Slice> = {
+export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSerializedState = StateFromSlice<T>> = {
   /**
    * The redux slice (return of createSlice).
    */
@@ -16,7 +15,7 @@ export type SliceConfig<T extends Slice> = {
   /**
    * A function that returns the initial state of the slice.
    */
-  getInitialState: () => StateFromSlice<T>;
+  getInitialState: () => TSerializedState;
   /**
    * The optional persist configuration for this slice. If omitted, the slice will not be persisted.
    */
@@ -28,7 +27,7 @@ export type SliceConfig<T extends Slice> = {
      * @param state The rehydrated state.
      * @returns A correctly-shaped state.
      */
-    migrate: (state: unknown) => StateFromSlice<T>;
+    migrate: (state: unknown) => TSerializedState;
     /**
      * Keys to omit from the persisted state.
      */
@@ -39,8 +38,18 @@ export type SliceConfig<T extends Slice> = {
    */
   undoableConfig?: {
     /**
-     * The options to be passed into redux-undo.
+     * Wraps state into state with history
+     *
+     * @param state The state without history
+     * @returns The state with history
+     */
+    wrapState: (state: unknown) => TInternalState;
+    /**
+     * Unwraps state with history
+     *
+     * @param state The state with history
+     * @returns The state without history
      */
-    reduxUndoOptions: UndoableOptions<StateFromSlice<T>>;
+    unwrapState: (state: TInternalState) => TSerializedState;
   };
 };