From 902309e2d932af961a69b11d383d06ed2d635845 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 11 Apr 2024 13:23:41 +0700
Subject: [PATCH] fix: cannot reset data while starting model

---
 web/hooks/useActiveModel.ts  | 40 +++++++++++++++---------------------
 web/hooks/useFactoryReset.ts | 16 +++++++--------
 web/utils/retry.ts           | 24 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 31 deletions(-)
 create mode 100644 web/utils/retry.ts

diff --git a/web/hooks/useActiveModel.ts b/web/hooks/useActiveModel.ts
index 34ffd1af73..1e648f60ee 100644
--- a/web/hooks/useActiveModel.ts
+++ b/web/hooks/useActiveModel.ts
@@ -126,33 +126,27 @@ export function useActiveModel() {
       })
   }
 
-  const stopModel = useCallback(
-    async (model?: Model) => {
-      const stoppingModel = activeModel || model
-      if (
-        !stoppingModel ||
-        (!model && stateModel.state === 'stop' && stateModel.loading)
-      )
-        return
-
-      setStateModel({ state: 'stop', loading: true, model: stoppingModel })
-      const engine = EngineManager.instance().get(stoppingModel.engine)
-      return engine
-        ?.unloadModel(stoppingModel)
-        .catch()
-        .then(() => {
-          setActiveModel(undefined)
-          setStateModel({ state: 'start', loading: false, model: undefined })
-          loadModelController?.abort()
-        })
-    },
-    [activeModel, setActiveModel, setStateModel, stateModel]
-  )
+  const stopModel = useCallback(async () => {
+    const stoppingModel = activeModel || stateModel.model
+    if (!stoppingModel || (stateModel.state === 'stop' && stateModel.loading))
+      return
+
+    setStateModel({ state: 'stop', loading: true, model: stoppingModel })
+    const engine = EngineManager.instance().get(stoppingModel.engine)
+    return engine
+      ?.unloadModel(stoppingModel)
+      .catch()
+      .then(() => {
+        setActiveModel(undefined)
+        setStateModel({ state: 'start', loading: false, model: undefined })
+        loadModelController?.abort()
+      })
+  }, [activeModel, setActiveModel, setStateModel, stateModel])
 
   const stopInference = useCallback(async () => {
     // Loading model
     if (stateModel.loading) {
-      stopModel(stateModel.model)
+      stopModel()
       return
     }
     if (!activeModel) return
diff --git a/web/hooks/useFactoryReset.ts b/web/hooks/useFactoryReset.ts
index 878461ef12..4411c553d4 100644
--- a/web/hooks/useFactoryReset.ts
+++ b/web/hooks/useFactoryReset.ts
@@ -3,6 +3,8 @@ import { useCallback } from 'react'
 import { fs, AppConfiguration } from '@janhq/core'
 import { atom, useAtomValue, useSetAtom } from 'jotai'
 
+import { PromiseRetry } from '@/utils/retry'
+
 import { useActiveModel } from './useActiveModel'
 
 import { defaultJanDataFolderAtom } from '@/helpers/atoms/App.atom'
@@ -19,7 +21,7 @@ export const factoryResetStateAtom = atom(FactoryResetState.Idle)
 
 export default function useFactoryReset() {
   const defaultJanDataFolder = useAtomValue(defaultJanDataFolderAtom)
-  const { activeModel, stopModel } = useActiveModel()
+  const { stopModel } = useActiveModel()
   const setFactoryResetState = useSetAtom(factoryResetStateAtom)
 
   const resetAll = useCallback(
@@ -44,14 +46,12 @@ export default function useFactoryReset() {
         await window.core?.api?.updateAppConfiguration(configuration)
       }
 
-      if (activeModel) {
-        setFactoryResetState(FactoryResetState.StoppingModel)
-        await stopModel()
-        await new Promise((resolve) => setTimeout(resolve, 4000))
-      }
+      setFactoryResetState(FactoryResetState.StoppingModel)
+      await stopModel()
 
       setFactoryResetState(FactoryResetState.DeletingData)
-      await fs.rm(janDataFolderPath)
+      // Try to delete the jan data folder
+      await PromiseRetry(fs.rm(janDataFolderPath), 10, 1000)
 
       setFactoryResetState(FactoryResetState.ClearLocalStorage)
       // reset the localStorage
@@ -59,7 +59,7 @@ export default function useFactoryReset() {
 
       await window.core?.api?.relaunch()
     },
-    [defaultJanDataFolder, activeModel, stopModel, setFactoryResetState]
+    [defaultJanDataFolder, stopModel, setFactoryResetState]
   )
 
   return {
diff --git a/web/utils/retry.ts b/web/utils/retry.ts
new file mode 100644
index 0000000000..8a968db317
--- /dev/null
+++ b/web/utils/retry.ts
@@ -0,0 +1,24 @@
+export function PromiseRetry<T>(
+  promise: Promise<T>,
+  times: number = 5,
+  delay: number = 500
+): Promise<T> {
+  return new Promise<T>((resolve, reject) => {
+    const retry = async () => {
+      let lastError: Error | undefined
+      for (let i = 0; i < times; i++) {
+        try {
+          const res = await promise
+          resolve(res)
+          return
+        } catch (err) {
+          lastError = err as Error
+          console.warn(err)
+          await new Promise((resolve) => setTimeout(resolve, delay))
+        }
+      }
+      reject(lastError)
+    }
+    retry()
+  })
+}