From 4507bb164288fdf122bfbaccd6fc2d2099d45e8b Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Wed, 3 Apr 2024 11:59:22 +0700
Subject: [PATCH] fix: app does not gracefully stop a model

---
 server/index.ts             |  2 +-
 web/hooks/useActiveModel.ts | 22 +++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/server/index.ts b/server/index.ts
index 98cc8385d1..7e02847f2a 100644
--- 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
index d869afd2fe..10da808318 100644
--- a/web/hooks/useActiveModel.ts
+++ b/web/hooks/useActiveModel.ts
@@ -86,11 +86,6 @@ 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.`,
@@ -98,12 +93,6 @@ export function useActiveModel() {
         })
       })
       .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 }
 }