From 2b732bc9761cfd8847e27f368a51fde38b34f311 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 17 Apr 2025 21:49:32 +0700
Subject: [PATCH] fix: clean up threads persistence and fix assistant update
 issue

---
 src-tauri/src/core/threads.rs          |  2 +-
 web/containers/ModelDropdown/index.tsx | 44 +++++++++++---------------
 web/hooks/useCreateNewThread.ts        |  2 +-
 web/hooks/useUpdateModelParameters.ts  |  8 ++++-
 4 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/src-tauri/src/core/threads.rs b/src-tauri/src/core/threads.rs
index 34b0bbd77..051837992 100644
--- a/src-tauri/src/core/threads.rs
+++ b/src-tauri/src/core/threads.rs
@@ -465,7 +465,7 @@ pub async fn modify_thread_assistant<R: Runtime>(
     {
         if let Some(index) = assistants
             .iter()
-            .position(|a| a.get("id").and_then(|v| v.as_str()) == Some(assistant_id))
+            .position(|a| a.get("assistant_id").and_then(|v| v.as_str()) == Some(assistant_id))
         {
             assistants[index] = assistant.clone();
             let data = serde_json::to_string_pretty(&thread).map_err(|e| e.to_string())?;
diff --git a/web/containers/ModelDropdown/index.tsx b/web/containers/ModelDropdown/index.tsx
index d95a114c4..0d9f862df 100644
--- a/web/containers/ModelDropdown/index.tsx
+++ b/web/containers/ModelDropdown/index.tsx
@@ -232,26 +232,6 @@ const ModelDropdown = ({
       stopModel()
 
       if (activeThread) {
-        // Change assistand tools based on model support RAG
-        updateThreadMetadata({
-          ...activeThread,
-          assistants: [
-            {
-              ...activeAssistant,
-              tools: [
-                {
-                  type: 'retrieval',
-                  enabled: model?.engine === InferenceEngine.cortex,
-                  settings: {
-                    ...(activeAssistant.tools &&
-                      activeAssistant.tools[0]?.settings),
-                  },
-                },
-              ],
-            },
-          ],
-        })
-
         const contextLength = model?.settings.ctx_len
           ? Math.min(8192, model?.settings.ctx_len ?? 8192)
           : undefined
@@ -273,11 +253,25 @@ const ModelDropdown = ({
 
         // Update model parameter to the thread file
         if (model)
-          updateModelParameter(activeThread, {
-            params: modelParams,
-            modelId: model.id,
-            engine: model.engine,
-          })
+          updateModelParameter(
+            activeThread,
+            {
+              params: modelParams,
+              modelId: model.id,
+              engine: model.engine,
+            },
+            // Update tools
+            [
+              {
+                type: 'retrieval',
+                enabled: model?.engine === InferenceEngine.cortex,
+                settings: {
+                  ...(activeAssistant.tools &&
+                    activeAssistant.tools[0]?.settings),
+                },
+              },
+            ]
+          )
       }
     },
     [
diff --git a/web/hooks/useCreateNewThread.ts b/web/hooks/useCreateNewThread.ts
index 57ceeb385..6f1efc04d 100644
--- a/web/hooks/useCreateNewThread.ts
+++ b/web/hooks/useCreateNewThread.ts
@@ -180,7 +180,7 @@ export const useCreateNewThread = () => {
       updateThreadCallback(thread)
       if (thread.assistants && thread.assistants?.length > 0) {
         setActiveAssistant(thread.assistants[0])
-        updateAssistantCallback(thread.id, thread.assistants[0])
+        return updateAssistantCallback(thread.id, thread.assistants[0])
       }
     },
     [
diff --git a/web/hooks/useUpdateModelParameters.ts b/web/hooks/useUpdateModelParameters.ts
index dab2f6e28..8bab0c357 100644
--- a/web/hooks/useUpdateModelParameters.ts
+++ b/web/hooks/useUpdateModelParameters.ts
@@ -1,6 +1,7 @@
 import { useCallback } from 'react'
 
 import {
+  AssistantTool,
   ConversationalExtension,
   ExtensionTypeEnum,
   InferenceEngine,
@@ -51,7 +52,11 @@ export default function useUpdateModelParameters() {
   )
 
   const updateModelParameter = useCallback(
-    async (thread: Thread, settings: UpdateModelParameter) => {
+    async (
+      thread: Thread,
+      settings: UpdateModelParameter,
+      tools?: AssistantTool[]
+    ) => {
       if (!activeAssistant) return
 
       const toUpdateSettings = processStopWords(settings.params ?? {})
@@ -70,6 +75,7 @@ export default function useUpdateModelParameters() {
       const settingParams = extractModelLoadParams(updatedModelParams)
       const assistantInfo = {
         ...activeAssistant,
+        tools: tools ?? activeAssistant.tools,
         model: {
           ...activeAssistant?.model,
           parameters: runtimeParams,