undo/redo refinements

Author: Attila Cseh

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

@@ -22,7 +22,7 @@ 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, undoableCanvasSliceReducer } from 'features/controlLayers/store/canvasSlice';
+import { canvasSliceConfig, undoableCanvasesReducer } 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';
@@ -90,7 +90,7 @@ const ALL_REDUCERS = {
   [api.reducerPath]: api.reducer,
   [canvasSessionSliceConfig.slice.reducerPath]: canvasSessionSliceConfig.slice.reducer,
   [canvasSettingsSliceConfig.slice.reducerPath]: canvasSettingsSliceConfig.slice.reducer,
-  [canvasSliceConfig.slice.reducerPath]: undoableCanvasSliceReducer,
+  [canvasSliceConfig.slice.reducerPath]: undoableCanvasesReducer,
   [changeBoardModalSliceConfig.slice.reducerPath]: changeBoardModalSliceConfig.slice.reducer,
   [configSliceConfig.slice.reducerPath]: configSliceConfig.slice.reducer,
   [dynamicPromptsSliceConfig.slice.reducerPath]: dynamicPromptsSliceConfig.slice.reducer,
@@ -119,7 +119,7 @@ const unserialize: UnserializeFunction = (data, key) => {
   if (!sliceConfig?.persistConfig) {
     throw new Error(`No persist config for slice "${key}"`);
   }
-  const { getInitialState, persistConfig, undoableConfig } = sliceConfig;
+  const { getInitialState, persistConfig } = sliceConfig;
   let state;
   try {
     const initialState = getInitialState();
@@ -151,12 +151,7 @@ const unserialize: UnserializeFunction = (data, key) => {
     state = getInitialState();
   }
 
-  // Undoable slices must be wrapped in a history!
-  if (undoableConfig) {
-    return undoableConfig.wrapState(state);
-  } else {
-    return state;
-  }
+  return persistConfig.wrapState ? persistConfig.wrapState(state) : state;
 };
 
 const serialize: SerializeFunction = (data, key) => {
@@ -166,7 +161,7 @@ const serialize: SerializeFunction = (data, key) => {
   }
 
   const result = omit(
-    sliceConfig.undoableConfig ? sliceConfig.undoableConfig.unwrapState(data) : data,
+    sliceConfig.persistConfig.unwrapState ? sliceConfig.persistConfig.unwrapState(data) : data,
     sliceConfig.persistConfig.persistDenylist ?? []
   );
 

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

@@ -32,24 +32,19 @@ export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSe
      * Keys to omit from the persisted state.
      */
     persistDenylist?: (keyof StateFromSlice<T>)[];
-  };
-  /**
-   * The optional undoable configuration for this slice. If omitted, the slice will not be undoable.
-   */
-  undoableConfig?: {
     /**
      * Wraps state into state with history
      *
      * @param state The state without history
      * @returns The state with history
      */
-    wrapState: (state: unknown) => TInternalState;
+    wrapState?: (state: unknown) => TInternalState;
     /**
      * Unwraps state with history
      *
      * @param state The state with history
      * @returns The state without history
      */
-    unwrapState: (state: TInternalState) => TSerializedState;
+    unwrapState?: (state: TInternalState) => TSerializedState;
   };
 };

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

@@ -109,8 +109,6 @@ import {
   makeDefaultRasterLayerAdjustments,
 } from './util';
 
-type CanvasDeletedPayloadAction = PayloadAction<{ id: string }>;
-
 const getInitialCanvasState = (id: string, name: string): CanvasState => ({
   id,
   name,
@@ -156,7 +154,7 @@ const getInitialCanvasesHistoryState = (): CanvasesStateWithHistory => {
   };
 };
 
-const slice = createSlice({
+const canvasesSlice = createSlice({
   name: 'canvas',
   initialState: getInitialCanvasesHistoryState(),
   reducers: {
@@ -198,7 +196,7 @@ const slice = createSlice({
 
       canvas.name = name;
     },
-    canvasDeleted: (state, action: CanvasDeletedPayloadAction) => {
+    canvasDeleted: (state, action: PayloadAction<{ id: string }>) => {
       const { id } = action.payload;
 
       if (state.canvases.length === 1) {
@@ -1869,7 +1867,7 @@ export const {
   canvasSelected,
   canvasNameChanged,
   canvasDeleted,
-} = slice.actions;
+} = canvasesSlice.actions;
 
 export const {
   canvasMetadataRecalled,
@@ -1969,6 +1967,8 @@ export const {
   // inpaintMaskRecalled,
 } = canvasSlice.actions;
 
+const isCanvasSliceAction = isAnyOf(...Object.values(canvasSlice.actions));
+
 let filter = true;
 
 const canvasUndoableConfig: UndoableOptions<CanvasState, UnknownAction> = {
@@ -1978,7 +1978,7 @@ const canvasUndoableConfig: UndoableOptions<CanvasState, UnknownAction> = {
   clearHistoryType: canvasClearHistory.type,
   filter: (action, _state, _history) => {
     // Ignore both all actions from other slices and canvas management actions
-    if (!action.type.startsWith(slice.name)) {
+    if (!action.type.startsWith(canvasSlice.name)) {
       return false;
     }
     // Throttle rapid actions of the same type
@@ -1991,11 +1991,15 @@ const canvasUndoableConfig: UndoableOptions<CanvasState, UnknownAction> = {
 
 const undoableCanvasReducer = undoable(canvasSlice.reducer, canvasUndoableConfig);
 
-export const undoableCanvasSliceReducer = (
+export const undoableCanvasesReducer = (
   state: CanvasesStateWithHistory,
   action: UnknownAction
 ): CanvasesStateWithHistory => {
-  state = slice.reducer(state, action);
+  state = canvasesSlice.reducer(state, action);
+
+  if (!isCanvasSliceAction(action)) {
+    return state;
+  }
 
   return {
     ...state,
@@ -2005,28 +2009,14 @@ export const undoableCanvasSliceReducer = (
   };
 };
 
-export const canvasSliceConfig: SliceConfig<typeof slice, CanvasesStateWithHistory, CanvasesStateWithoutHistory> = {
-  slice,
+export const canvasSliceConfig: SliceConfig<
+  typeof canvasesSlice,
+  CanvasesStateWithHistory,
+  CanvasesStateWithoutHistory
+> = {
+  slice: canvasesSlice,
   getInitialState: getInitialCanvasesState,
   schema: zCanvasesStateWithHistory,
-  undoableConfig: {
-    unwrapState: (state) => {
-      return {
-        _version: state._version,
-        selectedCanvasId: state.selectedCanvasId,
-        canvases: state.canvases.map((canvas) => canvas.present),
-      };
-    },
-    wrapState: (state) => {
-      const canvasesState = state as CanvasesStateWithoutHistory;
-
-      return {
-        _version: canvasesState._version,
-        selectedCanvasId: canvasesState.selectedCanvasId,
-        canvases: canvasesState.canvases.map((canvas) => newHistory([], canvas, [])),
-      };
-    },
-  },
   persistConfig: {
     migrate: (state) => {
       assert(isPlainObject(state));
@@ -2049,6 +2039,22 @@ export const canvasSliceConfig: SliceConfig<typeof slice, CanvasesStateWithHisto
       }
       return zCanvasesStateWithoutHistory.parse(state);
     },
+    wrapState: (state) => {
+      const canvasesState = state as CanvasesStateWithoutHistory;
+
+      return {
+        _version: canvasesState._version,
+        selectedCanvasId: canvasesState.selectedCanvasId,
+        canvases: canvasesState.canvases.map((canvas) => newHistory([], canvas, [])),
+      };
+    },
+    unwrapState: (state) => {
+      return {
+        _version: state._version,
+        selectedCanvasId: state.selectedCanvasId,
+        canvases: state.canvases.map((canvas) => canvas.present),
+      };
+    },
   },
 };
 

invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts

@@ -819,14 +819,12 @@ export const nodesSliceConfig: SliceConfig<typeof slice, StateWithHistory<NodesS
       }
       return zNodesState.parse(state);
     },
-  },
-  undoableConfig: {
-    unwrapState: (state) => state.present,
     wrapState: (state) => {
       const nodesState = state as NodesState;
 
       return newHistory([], nodesState, []);
     },
+    unwrapState: (state) => state.present,
   },
 };