fix: hydrated

xinyin025 · Sep 3, 2024 · ed9aae5 · ed9aae5
1 parent 886ffc0
commit ed9aae5
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/app/store/prompt.ts b/app/store/prompt.ts
@@ -179,6 +179,8 @@ export const usePromptStore = createPersistStore(
             res.en.length + res.cn.length + res.tw.length;
           SearchService.init(allPromptsForSearch, userPrompts);
         });
+
+      return () => state.setHasHydrated(true);
     },
   },
 );
diff --git a/app/utils/indexedDB-storage.ts b/app/utils/indexedDB-storage.ts
@@ -5,11 +5,7 @@ class IndexedDBStorage implements StateStorage {
   public async getItem(name: string): Promise<string | null> {
     try {
       const value = (await get(name)) || localStorage.getItem(name);
-      const _value = JSON.parse(value);
-      if (_value?.state) {
-        _value.state._hasHydrated = true;
-      }
-      return JSON.stringify(_value);
+      return value;
     } catch (error) {
       return localStorage.getItem(name);
     }

diff --git a/app/utils/store.ts b/app/utils/store.ts
@@ -14,9 +14,11 @@ type SecondParam<T> = T extends (
 
 type MakeUpdater<T> = {
   lastUpdateTime: number;
+  _hasHydrated: boolean;
 
   markUpdate: () => void;
   update: Updater<T>;
+  setHasHydrated: (state: boolean) => void;
 };
 
 type SetStoreState<T> = (
@@ -33,12 +35,18 @@ export function createPersistStore<T extends object, M>(
   persistOptions: SecondParam<typeof persist<T & M & MakeUpdater<T>>>,
 ) {
   persistOptions.storage = createJSONStorage(() => indexedDBStorage);
+  persistOptions.onRehydrateStorage = persistOptions.onRehydrateStorage
+    ? persistOptions.onRehydrateStorage
+    : (state) => {
+        return () => state.setHasHydrated(true);
+      };
   return create(
     persist(
       combine(
         {
           ...state,
           lastUpdateTime: 0,
+          _hasHydrated: false,
         },
         (set, get) => {
           return {
@@ -57,6 +65,9 @@ export function createPersistStore<T extends object, M>(
                 lastUpdateTime: Date.now(),
               });
             },
+            setHasHydrated: (state: boolean) => {
+              set({ _hasHydrated: state } as Partial<T & M & MakeUpdater<T>>);
+            },
           } as M & MakeUpdater<T>;
         },
       ),