feat(ui): port UI slice to zod

Author: psychedelicious

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

@@ -3,55 +3,67 @@ import { createSelector, createSlice } from '@reduxjs/toolkit';
 import type { PersistConfig, RootState } from 'app/store/store';
 import { canvasReset } from 'features/controlLayers/store/actions';
 import { canvasSessionReset, generateSessionReset } from 'features/controlLayers/store/canvasStagingAreaSlice';
-import type { Dimensions } from 'features/controlLayers/store/types';
 import { workflowLoaded } from 'features/nodes/store/nodesSlice';
 import { atom } from 'nanostores';
 
-import type { CanvasRightPanelTabName, TabName, UIState } from './uiTypes';
-
-const initialUIState: UIState = {
-  _version: 3,
-  activeTab: 'canvas',
-  activeTabCanvasRightPanel: 'gallery',
-  shouldShowImageDetails: false,
-  shouldShowProgressInViewer: true,
-  accordions: {},
-  expanders: {},
-  textAreaSizes: {},
-  shouldShowNotificationV2: true,
-};
+import type { TabName, UIState } from './uiTypes';
+import { getInitialUIState } from './uiTypes';
 
 export const uiSlice = createSlice({
   name: 'ui',
-  initialState: initialUIState,
+  initialState: getInitialUIState(),
   reducers: {
-    setActiveTab: (state, action: PayloadAction<TabName>) => {
+    setActiveTab: (state, action: PayloadAction<UIState['activeTab']>) => {
       state.activeTab = action.payload;
     },
-    activeTabCanvasRightPanelChanged: (state, action: PayloadAction<CanvasRightPanelTabName>) => {
+    activeTabCanvasRightPanelChanged: (state, action: PayloadAction<UIState['activeTabCanvasRightPanel']>) => {
       state.activeTabCanvasRightPanel = action.payload;
     },
-    setShouldShowImageDetails: (state, action: PayloadAction<boolean>) => {
+    setShouldShowImageDetails: (state, action: PayloadAction<UIState['shouldShowImageDetails']>) => {
       state.shouldShowImageDetails = action.payload;
     },
-    setShouldShowProgressInViewer: (state, action: PayloadAction<boolean>) => {
+    setShouldShowProgressInViewer: (state, action: PayloadAction<UIState['shouldShowProgressInViewer']>) => {
       state.shouldShowProgressInViewer = action.payload;
     },
-    accordionStateChanged: (state, action: PayloadAction<{ id: string; isOpen: boolean }>) => {
+    accordionStateChanged: (
+      state,
+      action: PayloadAction<{
+        id: keyof UIState['accordions'];
+        isOpen: UIState['accordions'][keyof UIState['accordions']];
+      }>
+    ) => {
       const { id, isOpen } = action.payload;
       state.accordions[id] = isOpen;
     },
-    expanderStateChanged: (state, action: PayloadAction<{ id: string; isOpen: boolean }>) => {
+    expanderStateChanged: (
+      state,
+      action: PayloadAction<{
+        id: keyof UIState['expanders'];
+        isOpen: UIState['expanders'][keyof UIState['expanders']];
+      }>
+    ) => {
       const { id, isOpen } = action.payload;
       state.expanders[id] = isOpen;
     },
-    textAreaSizesStateChanged: (state, action: PayloadAction<{ id: string; size: Partial<Dimensions> }>) => {
+    textAreaSizesStateChanged: (
+      state,
+      action: PayloadAction<{
+        id: keyof UIState['textAreaSizes'];
+        size: UIState['textAreaSizes'][keyof UIState['textAreaSizes']];
+      }>
+    ) => {
       const { id, size } = action.payload;
       state.textAreaSizes[id] = size;
     },
-    shouldShowNotificationChanged: (state, action: PayloadAction<boolean>) => {
+    shouldShowNotificationChanged: (state, action: PayloadAction<UIState['shouldShowNotificationV2']>) => {
       state.shouldShowNotificationV2 = action.payload;
     },
+    showGenerateTabSplashScreenChanged: (state, action: PayloadAction<UIState['showGenerateTabSplashScreen']>) => {
+      state.showGenerateTabSplashScreen = action.payload;
+    },
+    showCanvasTabSplashScreenChanged: (state, action: PayloadAction<UIState['showCanvasTabSplashScreen']>) => {
+      state.showCanvasTabSplashScreen = action.payload;
+    },
   },
   extraReducers(builder) {
     builder.addCase(workflowLoaded, (state) => {
@@ -81,6 +93,8 @@ export const {
   expanderStateChanged,
   shouldShowNotificationChanged,
   textAreaSizesStateChanged,
+  showGenerateTabSplashScreenChanged,
+  showCanvasTabSplashScreenChanged,
 } = uiSlice.actions;
 
 export const selectUiSlice = (state: RootState) => state.ui;
@@ -103,7 +117,7 @@ const migrateUIState = (state: any): any => {
 
 export const uiPersistConfig: PersistConfig<UIState> = {
   name: uiSlice.name,
-  initialState: initialUIState,
+  initialState: getInitialUIState(),
   migrate: migrateUIState,
   persistDenylist: ['shouldShowImageDetails'],
 };

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

@@ -1,43 +1,30 @@
-import type { Dimensions } from 'features/controlLayers/store/types';
+import { deepClone } from 'common/util/deepClone';
+import { z } from 'zod';
 
-export type TabName = 'generate' | 'canvas' | 'upscaling' | 'workflows' | 'models' | 'queue';
-export type CanvasRightPanelTabName = 'layers' | 'gallery';
+const zTabName = z.enum(['generate', 'canvas', 'upscaling', 'workflows', 'models', 'queue']);
+export type TabName = z.infer<typeof zTabName>;
+const zCanvasRightPanelTabName = z.enum(['layers', 'gallery']);
+export type CanvasRightPanelTabName = z.infer<typeof zCanvasRightPanelTabName>;
 
-export interface UIState {
-  /**
-   * Slice schema version.
-   */
-  _version: 3;
-  /**
-   * The currently active tab.
-   */
-  activeTab: TabName;
-  /**
-   * The currently active right panel canvas tab
-   */
-  activeTabCanvasRightPanel: CanvasRightPanelTabName;
-  /**
-   * Whether or not to show image details, e.g. metadata, workflow, etc.
-   */
-  shouldShowImageDetails: boolean;
-  /**
-   * Whether or not to show progress in the viewer.
-   */
-  shouldShowProgressInViewer: boolean;
-  /**
-   * The state of accordions. The key is the id of the accordion, and the value is a boolean representing the open state.
-   */
-  accordions: Record<string, boolean>;
-  /**
-   * The state of expanders. The key is the id of the expander, and the value is a boolean representing the open state.
-   */
-  expanders: Record<string, boolean>;
-  /**
-   * The size of textareas. The key is the id of the text area, and the value is an object representing its width and/or height.
-   */
-  textAreaSizes: Record<string, Partial<Dimensions>>;
-  /**
-   * Whether or not to show the user the open notification. Bump version to reset users who may have closed previous version.
-   */
-  shouldShowNotificationV2: boolean;
-}
+const zPartialDimensions = z.object({
+  width: z.number().optional(),
+  height: z.number().optional(),
+});
+export type PartialDimensions = z.infer<typeof zPartialDimensions>;
+
+export const zUIState = z.object({
+  _version: z.literal(3).default(3),
+  activeTab: zTabName.default('canvas'),
+  activeTabCanvasRightPanel: zCanvasRightPanelTabName.default('gallery'),
+  shouldShowImageDetails: z.boolean().default(false),
+  shouldShowProgressInViewer: z.boolean().default(true),
+  accordions: z.record(z.string(), z.boolean()).default(() => ({})),
+  expanders: z.record(z.string(), z.boolean()).default(() => ({})),
+  textAreaSizes: z.record(z.string(), zPartialDimensions).default({}),
+  shouldShowNotificationV2: z.boolean().default(true),
+  showGenerateTabSplashScreen: z.boolean().default(true),
+  showCanvasTabSplashScreen: z.boolean().default(true),
+});
+const INITIAL_STATE = zUIState.parse({});
+export type UIState = z.infer<typeof zUIState>;
+export const getInitialUIState = (): UIState => deepClone(INITIAL_STATE);