fix: app does not gracefully stop a model

janhq · Apr 3, 2024 · 4507bb1 · 4507bb1
1 parent fe89901
commit 4507bb1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/server/index.ts b/server/index.ts
@@ -149,7 +149,7 @@ export const stopServer = async () => {
     // Log server stop
     if (isVerbose) logServer(`Debug: Server stopped`)
     // Stop the server
-    await server.close()
+    await server?.close()
   } catch (e) {
     // Log any errors
     if (isVerbose) logServer(`Error: ${e}`)

diff --git a/web/hooks/useActiveModel.ts b/web/hooks/useActiveModel.ts
@@ -86,24 +86,13 @@ export function useActiveModel() {
       ?.loadModel(model)
       .then(() => {
         setActiveModel(model)
-        setStateModel(() => ({
-          state: 'stop',
-          loading: false,
-          model: model.id,
-        }))
         toaster({
           title: 'Success!',
           description: `Model ${model.id} has been started.`,
           type: 'success',
         })
       })
       .catch((error) => {
-        setStateModel(() => ({
-          state: 'start',
-          loading: false,
-          model: model.id,
-        }))
-
         toaster({
           title: 'Failed!',
           description: `Model ${model.id} failed to start.`,
@@ -112,10 +101,17 @@ export function useActiveModel() {
         setLoadModelError(error)
         return Promise.reject(error)
       })
+      .finally(() => {
+        setStateModel(() => ({
+          state: 'start',
+          loading: false,
+          model: model.id,
+        }))
+      })
   }
 
   const stopModel = useCallback(async () => {
-    if (!activeModel) return
+    if (!activeModel || stateModel.state === 'stop') return
 
     setStateModel({ state: 'stop', loading: true, model: activeModel.id })
     const engine = EngineManager.instance().get(activeModel.engine)
@@ -126,7 +122,7 @@ export function useActiveModel() {
         setActiveModel(undefined)
         setStateModel({ state: 'start', loading: false, model: '' })
       })
-  }, [activeModel, setActiveModel, setStateModel])
+  }, [activeModel, stateModel, setActiveModel, setStateModel])
 
   return { activeModel, startModel, stopModel, stateModel }
 }