feat(hydration): implement guard for store hydration management

src/providers/HydrationGuard.tsx

@@ -0,0 +1,38 @@
+import type { ReactNode, FC } from "react";
+
+import { useGeneratorStore, useHistoryStore, useStartupStore } from "@/stores";
+
+/**
+ * Props for the HydrationGuard component
+ * @interface HydrationGuardProps
+ * @property {ReactNode} children - Child components to render once stores are hydrated
+ */
+interface HydrationGuardProps {
+  children: ReactNode;
+}
+
+/**
+ * HydrationGuard ensures that all required Zustand stores are hydrated before rendering children.
+ * This prevents hydration mismatches and ensures consistent state management across the application.
+ *
+ * @component
+ * @example
+ * ```tsx
+ * <HydrationGuard>
+ *   <App />
+ * </HydrationGuard>
+ * ```
+ *
+ * If any store is not yet hydrated, it returns null to prevent premature rendering.
+ */
+export const HydrationGuard: FC<HydrationGuardProps> = ({ children }) => {
+  const isGeneratorReady = useGeneratorStore((state) => state._hasHydrated);
+  const isHistoryReady = useHistoryStore((state) => state._hasHydrated);
+  const isStartupReady = useStartupStore((state) => state._hasHydrated);
+
+  if (!isGeneratorReady || !isHistoryReady || !isStartupReady) {
+    return null;
+  }
+
+  return <>{children}</>;
+};

src/providers/index.tsx

@@ -1,5 +1,6 @@
-import { type ReactNode, type FC } from "react";
+import type { ReactNode, FC } from "react";
 
+import { HydrationGuard } from "./HydrationGuard";
 import { ThemeProvider } from "./theme";
 import { CacheProvider } from "./cache";
 
@@ -34,7 +35,9 @@ interface AppProvidersProps {
  * ```
  */
 export const AppProviders: FC<AppProvidersProps> = ({ children }) => (
-  <CacheProvider>
-    <ThemeProvider>{children}</ThemeProvider>
-  </CacheProvider>
+  <HydrationGuard>
+    <CacheProvider>
+      <ThemeProvider>{children}</ThemeProvider>
+    </CacheProvider>
+  </HydrationGuard>
 );

src/stores/HydrationService.ts

@@ -0,0 +1,74 @@
+import type { HydrationState, WithHydration } from "./types";
+
+/**
+ * Service class for managing hydration state in Zustand stores
+ * Provides utilities for handling store rehydration from persistent storage
+ */
+export class HydrationService {
+  /**
+   * Returns the hydration configuration for persist middleware
+   * Includes callback to mark store as hydrated after rehydration
+   *
+   * @template T - Type of the store state
+   * @returns Configuration object with onRehydrateStorage callback
+   *
+   * @example
+   * ```typescript
+   * export const useStore = create<YourStoreType>()(
+   *   persist(
+   *     (set) => ({
+   *       // Store state
+   *     }),
+   *     {
+   *       ...HydrationService.getHydrationConfig(),
+   *       storage: customStorage,
+   *     }
+   *   )
+   * );
+   * ```
+   */
+  static getHydrationConfig<T>() {
+    return {
+      onRehydrateStorage: (state: WithHydration<T>) => {
+        return () => state.setHasHydrated(true);
+      },
+    };
+  }
+
+  /**
+   * Returns state object with hydration flag set to true
+   * Used to mark store as hydrated
+   *
+   * @returns Hydration state object with _hasHydrated set to true
+   *
+   * @example
+   * ```typescript
+   * setHasHydrated: () => set(() => HydrationService.setStateAsHydrated())
+   * ```
+   */
+  static setStateAsHydrated(): HydrationState {
+    return {
+      _hasHydrated: true,
+    };
+  }
+
+  /**
+   * Returns initial hydration state
+   * Used when initializing store
+   *
+   * @returns Initial hydration state with _hasHydrated set to false
+   *
+   * @example
+   * ```typescript
+   * export const useStore = create()((set) => ({
+   *   ...HydrationService.getInitialState(),
+   *   // Other initial state
+   * }));
+   * ```
+   */
+  static getInitialState(): HydrationState {
+    return {
+      _hasHydrated: false,
+    };
+  }
+}

src/stores/direction/direction.store.ts

@@ -1,9 +1,7 @@
-import { persist } from "zustand/middleware";
 import { create } from "zustand";
 
-import type { DirectionState, DirectionStore } from "./direction.types";
+import type { DirectionStore } from "./direction.types";
 
-import { createWxtStorage } from "../persistMiddleware";
 import { DirectionService } from "./direction.service";
 
 /**
@@ -23,44 +21,27 @@ import { DirectionService } from "./direction.service";
  * useDirectionStore.getState().initializeDirection();
  * ```
  */
-export const useDirectionStore = create<DirectionStore>()(
-  persist(
-    (set) => ({
-      // Initialize with default direction state
-      ...DirectionService.getInitialState(),
+export const useDirectionStore = create<DirectionStore>()((set) => ({
+  // Initialize with default direction state
+  ...DirectionService.getInitialState(),
 
-      /**
-       * Initializes direction based on browser language
-       */
-      initializeDirection: async () => {
-        const direction = await DirectionService.detectDirection();
+  /**
+   * Initializes direction based on browser language
+   */
+  initializeDirection: async () => {
+    const direction = await DirectionService.detectDirection();
 
-        set(() => DirectionService.setDirection(direction));
-      },
+    set(() => DirectionService.setDirection(direction));
+  },
 
-      /**
-       * Sets the current direction and updates HTML dir attribute
-       * @param direction - Direction to set ('ltr' or 'rtl')
-       */
-      setDirection: (direction) => {
-        set(() => DirectionService.setDirection(direction));
-      },
-    }),
-    {
-      /**
-       * Specifies which parts of the state should be persisted
-       * @param state - Current direction state
-       * @returns Partial state containing only persistent values
-       */
-      partialize: (state) => ({
-        direction: state.direction,
-      }),
-      // Use Wxt storage for persistence
-      storage: createWxtStorage<DirectionState>(),
-      name: "direction-storage",
-    },
-  ),
-);
+  /**
+   * Sets the current direction and updates HTML dir attribute
+   * @param direction - Direction to set ('ltr' or 'rtl')
+   */
+  setDirection: (direction) => {
+    set(() => DirectionService.setDirection(direction));
+  },
+}));
 
 /**
  * Initialize browser language direction detection when the store is created

src/stores/generator/generator.store.ts

@@ -6,9 +6,11 @@ import type {
   GeneratorState,
   GeneratorStore,
 } from "./generator.types";
+import type { WithHydration } from "../types";
 
 import { createWxtStorage } from "../persistMiddleware";
 import { GeneratorService } from "./generator.service";
+import { HydrationService } from "../HydrationService";
 
 /**
  * Zustand store for managing password/passphrase generation state and options
@@ -27,11 +29,12 @@ import { GeneratorService } from "./generator.service";
  * useGeneratorStore.getState().updatePasswordOptions({ length: 16 });
  * ```
  */
-export const useGeneratorStore = create<GeneratorStore>()(
+export const useGeneratorStore = create<WithHydration<GeneratorStore>>()(
   persist(
     (set, get) => ({
       // Initialize with default generator state
       ...GeneratorService.getInitialState(),
+      ...HydrationService.getInitialState(),
 
       /**
        * Updates passphrase generation options
@@ -61,6 +64,14 @@ export const useGeneratorStore = create<GeneratorStore>()(
         set((state) => GeneratorService.updatePasswordOptions(state, updates));
       },
 
+      /**
+       * Sets the hydration status to true
+       * Called automatically by persist middleware after rehydration
+       */
+      setHasHydrated: () => {
+        set(() => HydrationService.setStateAsHydrated());
+      },
+
       /**
        * Sets the currently displayed password/passphrase
        * @param value - Password or passphrase to set
@@ -91,6 +102,7 @@ export const useGeneratorStore = create<GeneratorStore>()(
       // Use Wxt storage for persistence
       storage: createWxtStorage<Omit<GeneratorState, "password">>(),
       name: "generator-storage",
+      ...HydrationService.getHydrationConfig(),
     },
   ),
 );

src/stores/history/core/history.store.ts

@@ -1,3 +1,6 @@
+import type { WithHydration } from "@/stores/types";
+
+import { HydrationService } from "@/stores/HydrationService";
 import { persist } from "zustand/middleware";
 import { create } from "zustand";
 
@@ -29,11 +32,12 @@ import { HistoryService } from "./history.service";
  * console.log(passwords);
  * ```
  */
-export const useHistoryStore = create<HistoryStore>()(
+export const useHistoryStore = create<WithHydration<HistoryStore>>()(
   persist(
     (set) => ({
       // Initialize with default history state
       ...HistoryService.getInitialState(),
+      ...HydrationService.getInitialState(),
 
       /**
        * Adds a new password to the history
@@ -51,6 +55,14 @@ export const useHistoryStore = create<HistoryStore>()(
         set((state) => HistoryService.removePassword(state, id));
       },
 
+      /**
+       * Sets the hydration status to true
+       * Called automatically by persist middleware after rehydration
+       */
+      setHasHydrated: () => {
+        set(() => HydrationService.setStateAsHydrated());
+      },
+
       /**
        * Clears entire password history
        * Resets to initial empty state
@@ -71,6 +83,7 @@ export const useHistoryStore = create<HistoryStore>()(
       // Use Wxt storage for persistence
       storage: createWxtStorage<HistoryState>(),
       name: "history-storage",
+      ...HydrationService.getHydrationConfig(),
     },
   ),
 );

src/stores/startup/startup.store.ts

@@ -2,8 +2,10 @@ import { persist } from "zustand/middleware";
 import { create } from "zustand";
 
 import type { StartupState, StartupStore } from "./startup.types";
+import type { WithHydration } from "../types";
 
 import { createWxtStorage } from "../persistMiddleware";
+import { HydrationService } from "../HydrationService";
 import { StartupService } from "./startup.service";
 
 /**
@@ -25,11 +27,12 @@ import { StartupService } from "./startup.service";
  * }
  * ```
  */
-export const useStartupStore = create<StartupStore>()(
+export const useStartupStore = create<WithHydration<StartupStore>>()(
   persist(
     (set) => ({
       // Initialize with default startup state
       ...StartupService.getInitialState(),
+      ...HydrationService.getInitialState(),
 
       /**
        * Updates the welcome screen viewed status
@@ -38,6 +41,14 @@ export const useStartupStore = create<StartupStore>()(
       setHasSeenWelcome: (value) => {
         set(() => StartupService.updateStatus(value));
       },
+
+      /**
+       * Sets the hydration status to true
+       * Called automatically by persist middleware after rehydration
+       */
+      setHasHydrated: () => {
+        set(() => HydrationService.setStateAsHydrated());
+      },
     }),
     {
       /**
@@ -51,6 +62,7 @@ export const useStartupStore = create<StartupStore>()(
       // Use Wxt storage for persistence
       storage: createWxtStorage<StartupState>(),
       name: "startup-storage",
+      ...HydrationService.getHydrationConfig(),
     },
   ),
 );