deserialization fixed

Author: Attila Cseh

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

@@ -106,12 +106,13 @@ const unserialize: UnserializeFunction = (data, key) => {
     throw new Error(`No persist config for slice "${key}"`);
   }
   const { getInitialState, persistConfig } = sliceConfig;
-  let state;
+
   try {
     const parsedState = JSON.parse(data);
+    const stateToMigrate = persistConfig.deserialize ? persistConfig.deserialize(parsedState) : parsedState;
 
     // Run migrations to bring old state up to date with the current version.
-    const migrated = persistConfig.migrate(parsedState);
+    const migrated = persistConfig.migrate(stateToMigrate);
 
     log.debug(
       {
@@ -121,16 +122,17 @@ const unserialize: UnserializeFunction = (data, key) => {
       },
       `Rehydrated slice "${key}"`
     );
-    state = migrated;
+
+    return migrated;
   } catch (err) {
     log.warn(
       { error: serializeError(err as Error) },
       `Error rehydrating slice "${key}", falling back to default initial state`
     );
-    state = getInitialState();
-  }
+    const initialState = getInitialState();
 
-  return persistConfig.deserialize ? persistConfig.deserialize(state) : state;
+    return persistConfig.deserialize ? persistConfig.deserialize(initialState) : initialState;
+  }
 };
 
 const serialize: SerializeFunction = (data, key) => {

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

@@ -1,5 +1,6 @@
 import type { Slice } from '@reduxjs/toolkit';
 import type { ZodType } from 'zod';
+import { z } 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]>;
@@ -28,7 +29,7 @@ export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSe
      * @param state The rehydrated state.
      * @returns A correctly-shaped state.
      */
-    migrate: (state: unknown) => TSerializedState;
+    migrate: (state: unknown) => TInternalState;
     /**
      * Serializes the state
      *
@@ -45,3 +46,14 @@ export type SliceConfig<T extends Slice, TInternalState = StateFromSlice<T>, TSe
     deserialize?: (state: unknown) => TInternalState;
   };
 };
+
+export const zStateWithHistory = <T extends z.ZodTypeAny>(stateSchema: T) =>
+  z.object({
+    past: z.array(stateSchema),
+    present: stateSchema,
+    future: z.array(stateSchema),
+    _latestUnfiltered: stateSchema.optional(),
+    group: z.unknown().optional(),
+    index: z.number().optional(),
+    limit: z.number().optional(),
+  });

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

@@ -101,7 +101,6 @@ import {
   isRegionalGuidanceFLUXReduxConfig,
   isRegionalGuidanceIPAdapterConfig,
   zCanvasStateWithHistory,
-  zCanvasStateWithoutHistory,
 } from './types';
 import {
   converters,
@@ -2102,7 +2101,7 @@ export const canvasSliceConfig: SliceConfig<typeof canvasSlice, CanvasStateWithH
           },
         };
       }
-      return zCanvasStateWithoutHistory.parse(state);
+      return zCanvasStateWithHistory.parse(state);
     },
     serialize: (state) => {
       return {

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

@@ -1,3 +1,4 @@
+import { zStateWithHistory } from 'app/store/types';
 import { deepClone } from 'common/util/deepClone';
 import type { CanvasEntityAdapter } from 'features/controlLayers/konva/CanvasEntity/types';
 import { zMainModelBase, zModelIdentifierField } from 'features/nodes/types/common';
@@ -874,16 +875,6 @@ const zCanvasStagingAreaState = z.object({
 });
 export type CanvasStagingAreaState = z.infer<typeof zCanvasStagingAreaState>;
 
-const zStateWithHistory = <T extends z.ZodTypeAny>(stateSchema: T) =>
-  z.object({
-    past: z.array(stateSchema),
-    present: stateSchema,
-    future: z.array(stateSchema),
-    _latestUnfiltered: stateSchema.optional(),
-    group: z.string().optional(),
-    index: z.number().optional(),
-    limit: z.number().optional(),
-  });
 const zCanvasEntity = z.object({
   selectedEntityIdentifier: zCanvasEntityIdentifer.nullable(),
   bookmarkedEntityIdentifier: zCanvasEntityIdentifer.nullable(),
@@ -916,9 +907,9 @@ const zCanvasState = <T extends z.ZodTypeAny>(canvasInstanceSchema: T) =>
     activeCanvasId: zId,
     canvases: z.record(zId, canvasInstanceSchema),
   });
-export const zCanvasStateWithoutHistory = zCanvasState(zCanvasInstanceStateWithoutHistory);
+const _zCanvasStateWithoutHistory = zCanvasState(zCanvasInstanceStateWithoutHistory);
 export const zCanvasStateWithHistory = zCanvasState(zCanvasInstanceStateWithHistory);
-export type CanvasState = z.infer<typeof zCanvasStateWithoutHistory>;
+export type CanvasState = z.infer<typeof _zCanvasStateWithoutHistory>;
 export type CanvasStateWithHistory = z.infer<typeof zCanvasStateWithHistory>;
 
 export const zCanvasReferenceImageState_OLD = zCanvasEntityBase.extend({

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

@@ -19,7 +19,7 @@ import {
   removeElement,
   reparentElement,
 } from 'features/nodes/components/sidePanel/builder/form-manipulation';
-import { type NodesState, zNodesState } from 'features/nodes/store/types';
+import { type NodesState, zNodesState, zNodesStateWithHistory } from 'features/nodes/store/types';
 import { SHARED_NODE_PROPERTIES } from 'features/nodes/types/constants';
 import type {
   BoardFieldValue,
@@ -85,12 +85,12 @@ import {
 } from 'features/nodes/types/workflow';
 import { atom, computed } from 'nanostores';
 import type { MouseEvent } from 'react';
-import type { StateWithHistory, UndoableOptions } from 'redux-undo';
+import type { UndoableOptions } from 'redux-undo';
 import undoable, { newHistory } from 'redux-undo';
 import { assert } from 'tsafe';
 import type { z } from 'zod';
 
-import type { PendingConnection, Templates } from './types';
+import type { NodesStateWithHistory, PendingConnection, Templates } from './types';
 
 export const getInitialWorkflow = (): Omit<NodesState, 'mode' | 'formFieldInitialValues' | '_version'> => {
   return {
@@ -807,7 +807,7 @@ const reduxUndoOptions: UndoableOptions<NodesState, UnknownAction> = {
 
 export const undoableNodesSliceReducer = undoable(slice.reducer, reduxUndoOptions);
 
-export const nodesSliceConfig: SliceConfig<typeof slice, StateWithHistory<NodesState>, NodesState> = {
+export const nodesSliceConfig: SliceConfig<typeof slice, NodesStateWithHistory, NodesState> = {
   slice,
   schema: zNodesState,
   getInitialState,
@@ -817,7 +817,7 @@ export const nodesSliceConfig: SliceConfig<typeof slice, StateWithHistory<NodesS
       if (!('_version' in state)) {
         state._version = 1;
       }
-      return zNodesState.parse(state);
+      return zNodesStateWithHistory.parse(state);
     },
     serialize: (state) => state.present,
     deserialize: (state) => {

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

@@ -1,4 +1,5 @@
 import type { HandleType } from '@xyflow/react';
+import { zStateWithHistory } from 'app/store/types';
 import { type FieldInputTemplate, type FieldOutputTemplate, zStatefulFieldValue } from 'features/nodes/types/field';
 import { type InvocationTemplate, type NodeExecutionState, zAnyEdge, zAnyNode } from 'features/nodes/types/invocation';
 import { zWorkflowV3 } from 'features/nodes/types/workflow';
@@ -23,4 +24,6 @@ export const zNodesState = z.object({
   formFieldInitialValues: z.record(z.string(), zStatefulFieldValue),
   ...zWorkflowV3.omit({ nodes: true, edges: true, is_published: true }).shape,
 });
+export const zNodesStateWithHistory = zStateWithHistory(zNodesState);
 export type NodesState = z.infer<typeof zNodesState>;
+export type NodesStateWithHistory = z.infer<typeof zNodesStateWithHistory>;