Fix multiple hydration error

Author: tommasoberlose

src/reducer.ts

@@ -78,6 +78,11 @@ export const createPersistedReducer: <
 ) => {
   // Subscribe the reducer to be persisted
   Settings.subscribeSlice(reducerName);
+  /**
+   * A flag to ensure the rehydration process runs only once.
+   * @internal
+   */
+  let isHydrated = false;
 
   /**
    * Creates a typed instance of the listener middleware's startListening function.
@@ -96,8 +101,10 @@ export const createPersistedReducer: <
   const reducer = createReducer(initialState, builder => {
     const b = new Builder(builder, UpdatedAtHelper.onStateChange.bind(null, reducerName));
     // Add a case to handle the rehydration of state from storage.
-    b.builder.addCase(REHYDRATE.toString(), (_state, action: PayloadAction<Record<ReducerName, S> | null>): void | S => {
-      if (action.payload?.[reducerName]) return action.payload[reducerName];
+    b.addCase(REHYDRATE.toString(), (_state, action: PayloadAction<Record<ReducerName, S> | null>): void | S => {
+      if (!isHydrated) {
+        if (action.payload?.[reducerName]) return action.payload[reducerName];
+      }
     });
     mapOrBuilderCallback(b);
   });
@@ -121,14 +128,15 @@ export const createPersistedReducer: <
   });
 
   /**
-   * Listens for the rehydration action to update the local timestamp,
-   * ensuring the state isn't immediately re-saved.
+   * Listens for the rehydration action to update the local timestamp and
+   * set the `isHydrated` flag to prevent subsequent rehydrations.
    * @internal
    */
   startAppListening({
     actionCreator: REHYDRATE,
     effect: () => {
       UpdatedAtHelper.onStateChange(reducerName);
+      isHydrated = true;
     },
   });
 

src/slice.ts

@@ -60,6 +60,11 @@ export const createPersistedSlice: <
 ) => {
   // Subscribe the slice to be persisted
   Settings.subscribeSlice(sliceOptions.name);
+  /**
+   * A flag to ensure the rehydration process runs only once per slice.
+   * @internal
+   */
+  let isHydrated = false;
 
   /**
    * Creates a typed instance of the listener middleware's startListening function.
@@ -80,8 +85,10 @@ export const createPersistedSlice: <
     extraReducers: builder => {
       const b = new Builder(builder, UpdatedAtHelper.onStateChange.bind(null, sliceOptions.name));
       // Add a case to handle the rehydration of state from storage.
-      b.builder.addCase(REHYDRATE.toString(), (_state, action: PayloadAction<Record<Name, SliceState> | null>): void | SliceState => {
-        if (action.payload?.[sliceOptions.name]) return action.payload[sliceOptions.name];
+      b.addCase(REHYDRATE.toString(), (_state, action: PayloadAction<Record<Name, SliceState> | null>): void | SliceState => {
+        if (!isHydrated) {
+          if (action.payload?.[sliceOptions.name]) return action.payload[sliceOptions.name];
+        }
       });
       // Allow the user to add their own extra reducers.
       sliceOptions.extraReducers?.(b);
@@ -124,13 +131,14 @@ export const createPersistedSlice: <
   });
 
   /**
-   * Listens for the rehydration action to update the local timestamp,
-   * ensuring the state isn't immediately re-saved.
+   * Listens for the rehydration action to update the local timestamp and
+   * set the `isHydrated` flag to prevent subsequent rehydrations.
    * @internal
    */
   startAppListening({
     actionCreator: REHYDRATE,
     effect: () => {
+      isHydrated = true;
       UpdatedAtHelper.onStateChange(sliceOptions.name);
     },
   });

src/store.ts

@@ -137,6 +137,17 @@ export const configurePersistedStore: <
     }));
   }
 
+  /**
+   * Overrides the original `replaceReducer` function to automatically
+   * trigger rehydration after a new reducer has been injected.
+   * @internal
+   */
+  const _replaceReducer = persistedStore.replaceReducer;
+  persistedStore.replaceReducer = (nR) => {
+    _replaceReducer.call(persistedStore, nR);
+    rehydrate();
+  }
+
   return new Promise<PersistedStore<S, A, M, E>>(async (resolve) => {
     try {
       const m = createListenerMiddleware();