From ed9aae531e0191d8b7fcbe594e0dc4e6176450da Mon Sep 17 00:00:00 2001
From: Dogtiti <499960698@qq.com>
Date: Tue, 3 Sep 2024 20:29:01 +0800
Subject: [PATCH] fix: hydrated

---
 app/store/prompt.ts            |  2 ++
 app/utils/indexedDB-storage.ts |  6 +-----
 app/utils/store.ts             | 11 +++++++++++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/app/store/prompt.ts b/app/store/prompt.ts
index a25cda5813a..29a97abba5f 100644
--- 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
index b80c84ce932..86b30477788 100644
--- 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
index 13bef6d5dfe..02c8139b0ad 100644
--- 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>;
         },
       ),