SliceConfig refactored

Author: Attila Cseh

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

@@ -17,9 +17,6 @@ import { addModelSelectedListener } from 'app/store/middleware/listenerMiddlewar
 import { addModelsLoadedListener } from 'app/store/middleware/listenerMiddleware/listeners/modelsLoaded';
 import { addSetDefaultSettingsListener } from 'app/store/middleware/listenerMiddleware/listeners/setDefaultSettings';
 import { addSocketConnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketConnected';
-import { deepClone } from 'common/util/deepClone';
-import { merge } from 'es-toolkit';
-import { omit, pick } from 'es-toolkit/compat';
 import { changeBoardModalSliceConfig } from 'features/changeBoardModal/store/slice';
 import { canvasSliceConfig } from 'features/controlLayers/store/canvasSlice';
 import { lorasSliceConfig } from 'features/controlLayers/store/lorasSlice';
@@ -117,21 +114,14 @@ const unserialize: UnserializeFunction = (data, key) => {
   const { getInitialState, persistConfig } = sliceConfig;
   let state;
   try {
-    const initialState = getInitialState();
-    const parsed = JSON.parse(data);
-
-    // We need to inject non-persisted values from initial state into the rehydrated state. These values always are
-    // required to be in the state, but won't be in the persisted data. Build an object that consists of only these
-    // values, then merge it with the rehydrated state.
-    const nonPersistedSubsetOfState = pick(initialState, persistConfig.persistDenylist ?? []);
-    const stateToMigrate = merge(deepClone(parsed), nonPersistedSubsetOfState);
+    const parsedState = JSON.parse(data);
 
     // Run migrations to bring old state up to date with the current version.
-    const migrated = persistConfig.migrate(stateToMigrate);
+    const migrated = persistConfig.migrate(parsedState);
 
     log.debug(
       {
-        persistedData: parsed as JsonObject,
+        persistedData: parsedState as JsonObject,
         rehydratedData: migrated as JsonObject,
         diff: diff(data, migrated) as JsonObject,
       },
@@ -146,7 +136,7 @@ const unserialize: UnserializeFunction = (data, key) => {
     state = getInitialState();
   }
 
-  return persistConfig.wrapState ? persistConfig.wrapState(state) : state;
+  return persistConfig.deserialize ? persistConfig.deserialize(state) : state;
 };
 
 const serialize: SerializeFunction = (data, key) => {
@@ -155,12 +145,9 @@ const serialize: SerializeFunction = (data, key) => {
     throw new Error(`No persist config for slice "${key}"`);
   }
 
-  const result = omit(
-    sliceConfig.persistConfig.unwrapState ? sliceConfig.persistConfig.unwrapState(data) : data,
-    sliceConfig.persistConfig.persistDenylist ?? []
-  );
+  const state = sliceConfig.persistConfig.serialize ? sliceConfig.persistConfig.serialize(data) : data;
 
-  return JSON.stringify(result);
+  return JSON.stringify(state);
 };
 
 const PERSISTED_KEYS = Object.values(SLICE_CONFIGS)

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

@@ -2,6 +2,7 @@ import type { Slice } from '@reduxjs/toolkit';
 import type { ZodType } from 'zod';
 
 type StateFromSlice<T extends Slice> = T extends Slice<infer U> ? U : never;
+export type SerializedStateFromDenyList<S, T extends readonly (keyof S)[]> = Omit<S, T[number]>;
 
 export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSerializedState = StateFromSlice<T>> = {
   /**
@@ -29,22 +30,18 @@ export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSe
      */
     migrate: (state: unknown) => TSerializedState;
     /**
-     * Keys to omit from the persisted state.
-     */
-    persistDenylist?: (keyof StateFromSlice<T>)[];
-    /**
-     * Wraps state into state with history
+     * Serializes the state
      *
-     * @param state The state without history
-     * @returns The state with history
+     * @param state The internal state
+     * @returns The serialized state
      */
-    wrapState?: (state: unknown) => TInternalState;
+    serialize?: (state: TInternalState) => TSerializedState;
     /**
-     * Unwraps state with history
+     * Deserializes the state
      *
-     * @param state The state with history
-     * @returns The state without history
+     * @param state The serialized state
+     * @returns The internal state
      */
-    unwrapState?: (state: TInternalState) => TSerializedState;
+    deserialize?: (state: unknown) => TInternalState;
   };
 };

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

@@ -1938,7 +1938,7 @@ export const isCanvasInstanceAction = (action: UnknownAction) =>
   isTabParamsStateAction(action) ||
   isCanvasSettingsStateAction(action) ||
   isCanvasStagingAreaStateAction(action);
-export const isCanvasEntityStateAction = isAnyOf(...Object.values(canvasEntityState.actions), canvasReset);
+const isCanvasEntityStateAction = isAnyOf(...Object.values(canvasEntityState.actions), canvasReset);
 
 export const {
   // Canvas
@@ -2102,28 +2102,34 @@ export const canvasSliceConfig: SliceConfig<typeof canvasSlice, CanvasStateWithH
       }
       return zCanvasStateWithoutHistory.parse(state);
     },
-    wrapState: (state) => {
-      const canvasState = state as CanvasState;
-
+    serialize: (state) => {
       return {
-        _version: canvasState._version,
-        activeCanvasId: canvasState.activeCanvasId,
+        _version: state._version,
+        activeCanvasId: state.activeCanvasId,
         canvases: Object.fromEntries(
-          Object.entries(canvasState.canvases).map(([canvasId, instance]) => [
+          Object.entries(state.canvases).map(([canvasId, instance]) => [
             canvasId,
-            { ...instance, canvas: newHistory([], instance.canvas, []) },
+            {
+              ...instance,
+              canvas: instance.canvas.present,
+            },
           ])
         ),
       };
     },
-    unwrapState: (state) => {
+    deserialize: (state) => {
+      const canvasState = state as CanvasState;
+
       return {
-        _version: state._version,
-        activeCanvasId: state.activeCanvasId,
+        _version: canvasState._version,
+        activeCanvasId: canvasState.activeCanvasId,
         canvases: Object.fromEntries(
-          Object.entries(state.canvases).map(([canvasId, instance]) => [
+          Object.entries(canvasState.canvases).map(([canvasId, instance]) => [
             canvasId,
-            { ...instance, canvas: instance.canvas.present },
+            {
+              ...instance,
+              canvas: newHistory([], instance.canvas, []),
+            },
           ])
         ),
       };

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

@@ -3,8 +3,8 @@ import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSelector, createSlice } from '@reduxjs/toolkit';
 import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import type { RootState } from 'app/store/store';
-import type { SliceConfig } from 'app/store/types';
-import { clamp } from 'es-toolkit/compat';
+import type { SerializedStateFromDenyList, SliceConfig } from 'app/store/types';
+import { clamp, merge, omit } from 'es-toolkit/compat';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
 import type {
   CroppableImageWithDims,
@@ -279,13 +279,21 @@ export const {
   refImagesRecalled,
 } = slice.actions;
 
-export const refImagesSliceConfig: SliceConfig<typeof slice> = {
+const denyList = ['selectedEntityId', 'isPanelOpen'] as const;
+type SerializedRefImagesState = SerializedStateFromDenyList<RefImagesState, typeof denyList>;
+
+export const refImagesSliceConfig: SliceConfig<typeof slice, RefImagesState, SerializedRefImagesState> = {
   slice,
   schema: zRefImagesState,
   getInitialState: getInitialRefImagesState,
   persistConfig: {
     migrate: (state) => zRefImagesState.parse(state),
-    persistDenylist: ['selectedEntityId', 'isPanelOpen'],
+    serialize: (state) => omit(state, denyList),
+    deserialize: (state) => {
+      const refImagesState = state as SerializedRefImagesState;
+
+      return merge(refImagesState, getInitialRefImagesState());
+    },
   },
 };
 

invokeai/frontend/web/src/features/dynamicPrompts/store/dynamicPromptsSlice.ts

@@ -1,9 +1,10 @@
 import type { PayloadAction, Selector } from '@reduxjs/toolkit';
 import { createSelector, createSlice } from '@reduxjs/toolkit';
 import type { RootState } from 'app/store/store';
-import type { SliceConfig } from 'app/store/types';
+import type { SerializedStateFromDenyList, SliceConfig } from 'app/store/types';
 import { buildZodTypeGuard } from 'common/util/zodUtils';
 import { isPlainObject } from 'es-toolkit';
+import { merge, omit } from 'es-toolkit/compat';
 import { assert } from 'tsafe';
 import { z } from 'zod';
 
@@ -69,21 +70,30 @@ export const {
   seedBehaviourChanged,
 } = slice.actions;
 
-export const dynamicPromptsSliceConfig: SliceConfig<typeof slice> = {
-  slice,
-  schema: zDynamicPromptsState,
-  getInitialState,
-  persistConfig: {
-    migrate: (state) => {
-      assert(isPlainObject(state));
-      if (!('_version' in state)) {
-        state._version = 1;
-      }
-      return zDynamicPromptsState.parse(state);
+const denyList = ['prompts', 'parsingError', 'isError', 'isLoading'] as const;
+type SerializedDynamicPromptsState = SerializedStateFromDenyList<DynamicPromptsState, typeof denyList>;
+
+export const dynamicPromptsSliceConfig: SliceConfig<typeof slice, DynamicPromptsState, SerializedDynamicPromptsState> =
+  {
+    slice,
+    schema: zDynamicPromptsState,
+    getInitialState,
+    persistConfig: {
+      migrate: (state) => {
+        assert(isPlainObject(state));
+        if (!('_version' in state)) {
+          state._version = 1;
+        }
+        return zDynamicPromptsState.parse(state);
+      },
+      serialize: (state) => omit(state, denyList),
+      deserialize: (state) => {
+        const dynamicPromptsState = state as SerializedDynamicPromptsState;
+
+        return merge(dynamicPromptsState, getInitialState());
+      },
     },
-    persistDenylist: ['prompts', 'parsingError', 'isError', 'isLoading'],
-  },
-};
+  };
 
 export const selectDynamicPromptsSlice = (state: RootState) => state.dynamicPrompts;
 const createDynamicPromptsSelector = <T>(selector: Selector<DynamicPromptsState, T>) =>

invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts

@@ -1,8 +1,9 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
 import type { RootState } from 'app/store/store';
-import type { SliceConfig } from 'app/store/types';
+import type { SerializedStateFromDenyList, SliceConfig } from 'app/store/types';
 import { isPlainObject } from 'es-toolkit';
+import { merge, omit } from 'es-toolkit/compat';
 import type { BoardRecordOrderBy } from 'services/api/types';
 import { assert } from 'tsafe';
 
@@ -176,7 +177,10 @@ export const {
 
 export const selectGallerySlice = (state: RootState) => state.gallery;
 
-export const gallerySliceConfig: SliceConfig<typeof slice> = {
+const denyList = ['selection', 'selectedBoardId', 'galleryView', 'imageToCompare'] as const;
+type SerializedGalleryState = SerializedStateFromDenyList<GalleryState, typeof denyList>;
+
+export const gallerySliceConfig: SliceConfig<typeof slice, GalleryState, SerializedGalleryState> = {
   slice,
   schema: zGalleryState,
   getInitialState,
@@ -188,6 +192,11 @@ export const gallerySliceConfig: SliceConfig<typeof slice> = {
       }
       return zGalleryState.parse(state);
     },
-    persistDenylist: ['selection', 'selectedBoardId', 'galleryView', 'imageToCompare'],
+    serialize: (state) => omit(state, denyList),
+    deserialize: (state) => {
+      const galleryState = state as SerializedGalleryState;
+
+      return merge(galleryState, getInitialState());
+    },
   },
 };

invokeai/frontend/web/src/features/modelManagerV2/store/modelManagerV2Slice.ts

@@ -1,8 +1,9 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSelector, createSlice } from '@reduxjs/toolkit';
 import type { RootState } from 'app/store/store';
-import type { SliceConfig } from 'app/store/types';
+import type { SerializedStateFromDenyList, SliceConfig } from 'app/store/types';
 import { isPlainObject } from 'es-toolkit';
+import { merge, omit } from 'es-toolkit/compat';
 import { zModelType } from 'features/nodes/types/common';
 import { assert } from 'tsafe';
 import z from 'zod';
@@ -67,7 +68,10 @@ export const {
   shouldInstallInPlaceChanged,
 } = slice.actions;
 
-export const modelManagerSliceConfig: SliceConfig<typeof slice> = {
+const denyList = ['selectedModelKey', 'selectedModelMode', 'filteredModelType', 'searchTerm'] as const;
+type SerializedModelManagerState = SerializedStateFromDenyList<ModelManagerState, typeof denyList>;
+
+export const modelManagerSliceConfig: SliceConfig<typeof slice, ModelManagerState, SerializedModelManagerState> = {
   slice,
   schema: zModelManagerState,
   getInitialState,
@@ -79,7 +83,12 @@ export const modelManagerSliceConfig: SliceConfig<typeof slice> = {
       }
       return zModelManagerState.parse(state);
     },
-    persistDenylist: ['selectedModelKey', 'selectedModelMode', 'filteredModelType', 'searchTerm'],
+    serialize: (state) => omit(state, denyList),
+    deserialize: (state) => {
+      const modelManagerState = state as SerializedModelManagerState;
+
+      return merge(modelManagerState, getInitialState());
+    },
   },
 };
 

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

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

invokeai/frontend/web/src/features/ui/store/uiSlice.ts

@@ -1,8 +1,9 @@
 import type { PayloadAction } from '@reduxjs/toolkit';
 import { createSlice } from '@reduxjs/toolkit';
 import type { RootState } from 'app/store/store';
-import type { SliceConfig } from 'app/store/types';
+import type { SerializedStateFromDenyList, SliceConfig } from 'app/store/types';
 import { isPlainObject } from 'es-toolkit';
+import { merge, omit } from 'es-toolkit/compat';
 import { assert } from 'tsafe';
 
 import { getInitialUIState, type UIState, zUIState } from './uiTypes';
@@ -87,7 +88,10 @@ export const {
 
 export const selectUiSlice = (state: RootState) => state.ui;
 
-export const uiSliceConfig: SliceConfig<typeof slice> = {
+const denyList = ['shouldShowItemDetails'] as const;
+type SerializedUIState = SerializedStateFromDenyList<UIState, typeof denyList>;
+
+export const uiSliceConfig: SliceConfig<typeof slice, UIState, SerializedUIState> = {
   slice,
   schema: zUIState,
   getInitialState: getInitialUIState,
@@ -111,6 +115,11 @@ export const uiSliceConfig: SliceConfig<typeof slice> = {
       }
       return zUIState.parse(state);
     },
-    persistDenylist: ['shouldShowItemDetails'],
+    serialize: (state) => omit(state, denyList),
+    deserialize: (state) => {
+      const uiState = state as SerializedUIState;
+
+      return merge(uiState, getInitialUIState());
+    },
   },
 };