fix: message should only be interrupted when i start another thread (#2053)

* fix: message should only be interrupted when i start another thread

* fix: thread lost message streaming if navigate to another thread

* fix: state issue with useThreads

web/containers/Providers/EventHandler.tsx

@@ -177,7 +177,6 @@ export default function EventHandler({ children }: { children: ReactNode }) {
       )
       if (message.status === MessageStatus.Pending) {
         if (message.content.length) {
-          updateThreadWaiting(message.thread_id, false)
           setIsGeneratingResponse(false)
         }
         return

web/helpers/atoms/ChatMessage.atom.ts

@@ -16,6 +16,8 @@ import {
  */
 export const chatMessages = atom<Record<string, ThreadMessage[]>>({})
 
+export const readyThreadsMessagesAtom = atom<Record<string, boolean>>({})
+
 /**
  * Return the chat messages for the current active conversation
  */
@@ -34,6 +36,10 @@ export const setConvoMessagesAtom = atom(
     }
     newData[threadId] = messages
     set(chatMessages, newData)
+    set(readyThreadsMessagesAtom, {
+      ...get(readyThreadsMessagesAtom),
+      [threadId]: true,
+    })
   }
 )
 

web/hooks/useCreateNewThread.ts

@@ -9,6 +9,8 @@ import {
   ThreadState,
   Model,
   AssistantTool,
+  events,
+  InferenceEvent,
 } from '@janhq/core'
 import { atom, useAtomValue, useSetAtom } from 'jotai'
 
@@ -30,6 +32,7 @@ import {
   threadStatesAtom,
   updateThreadAtom,
   setThreadModelParamsAtom,
+  isGeneratingResponseAtom,
 } from '@/helpers/atoms/Thread.atom'
 
 const createNewThreadAtom = atom(null, (get, set, newThread: Thread) => {
@@ -57,6 +60,7 @@ export const useCreateNewThread = () => {
   const setSelectedModel = useSetAtom(selectedModelAtom)
   const setThreadModelParams = useSetAtom(setThreadModelParamsAtom)
   const { experimentalFeature } = useContext(FeatureToggleContext)
+  const setIsGeneratingResponse = useSetAtom(isGeneratingResponseAtom)
 
   const { recommendedModel, downloadedModels } = useRecommendedModel()
 
@@ -66,6 +70,10 @@ export const useCreateNewThread = () => {
     assistant: Assistant,
     model?: Model | undefined
   ) => {
+    // Stop generating if any
+    setIsGeneratingResponse(false)
+    events.emit(InferenceEvent.OnInferenceStopped, {})
+
     const defaultModel = model ?? recommendedModel ?? downloadedModels[0]
 
     // check last thread message, if there empty last message use can not create thread

web/hooks/useSetActiveThread.ts

@@ -1,20 +1,14 @@
-import { useCallback } from 'react'
+import { ExtensionTypeEnum, Thread, ConversationalExtension } from '@janhq/core'
 
-import {
-  InferenceEvent,
-  ExtensionTypeEnum,
-  Thread,
-  events,
-  ConversationalExtension,
-} from '@janhq/core'
-
-import { useSetAtom } from 'jotai'
+import { useAtomValue, useSetAtom } from 'jotai'
 
 import { extensionManager } from '@/extension'
-import { setConvoMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
+import {
+  readyThreadsMessagesAtom,
+  setConvoMessagesAtom,
+} from '@/helpers/atoms/ChatMessage.atom'
 import {
   ModelParams,
-  isGeneratingResponseAtom,
   setActiveThreadIdAtom,
   setThreadModelParamsAtom,
 } from '@/helpers/atoms/Thread.atom'
@@ -23,31 +17,22 @@ export default function useSetActiveThread() {
   const setActiveThreadId = useSetAtom(setActiveThreadIdAtom)
   const setThreadMessage = useSetAtom(setConvoMessagesAtom)
   const setThreadModelParams = useSetAtom(setThreadModelParamsAtom)
-  const setIsGeneratingResponse = useSetAtom(isGeneratingResponseAtom)
+  const readyMessageThreads = useAtomValue(readyThreadsMessagesAtom)
 
-  const setActiveThread = useCallback(
-    async (thread: Thread) => {
-      setIsGeneratingResponse(false)
-      events.emit(InferenceEvent.OnInferenceStopped, thread.id)
-
-      // load the corresponding messages
+  const setActiveThread = async (thread: Thread) => {
+    // Load local messages only if there are no messages in the state
+    if (!readyMessageThreads[thread.id]) {
       const messages = await getLocalThreadMessage(thread.id)
       setThreadMessage(thread.id, messages)
+    }
 
-      setActiveThreadId(thread.id)
-      const modelParams: ModelParams = {
-        ...thread.assistants[0]?.model?.parameters,
-        ...thread.assistants[0]?.model?.settings,
-      }
-      setThreadModelParams(thread.id, modelParams)
-    },
-    [
-      setActiveThreadId,
-      setThreadMessage,
-      setThreadModelParams,
-      setIsGeneratingResponse,
-    ]
-  )
+    setActiveThreadId(thread.id)
+    const modelParams: ModelParams = {
+      ...thread.assistants[0]?.model?.parameters,
+      ...thread.assistants[0]?.model?.settings,
+    }
+    setThreadModelParams(thread.id, modelParams)
+  }
 
   return { setActiveThread }
 }

web/hooks/useThreads.ts

@@ -9,8 +9,6 @@ import {
 
 import { useSetAtom } from 'jotai'
 
-import useSetActiveThread from './useSetActiveThread'
-
 import { extensionManager } from '@/extension/ExtensionManager'
 import {
   ModelParams,
@@ -24,7 +22,6 @@ const useThreads = () => {
   const setThreadStates = useSetAtom(threadStatesAtom)
   const setThreads = useSetAtom(threadsAtom)
   const setThreadModelRuntimeParams = useSetAtom(threadModelParamsAtom)
-  const { setActiveThread } = useSetActiveThread()
   const setThreadDataReady = useSetAtom(threadDataReadyAtom)
 
   useEffect(() => {
@@ -56,16 +53,11 @@ const useThreads = () => {
       setThreadStates(localThreadStates)
       setThreads(localThreads)
       setThreadModelRuntimeParams(threadModelParams)
-
-      if (localThreads.length > 0) {
-        setActiveThread(localThreads[0])
-      }
       setThreadDataReady(true)
     }
 
     getThreads()
   }, [
-    setActiveThread,
     setThreadModelRuntimeParams,
     setThreadStates,
     setThreads,

web/screens/Chat/ChatInput/index.tsx

@@ -38,6 +38,8 @@ import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
 import {
   activeThreadAtom,
   getActiveThreadIdAtom,
+  isGeneratingResponseAtom,
+  threadStatesAtom,
   waitingToSendMessage,
 } from '@/helpers/atoms/Thread.atom'
 
@@ -57,6 +59,12 @@ const ChatInput: React.FC = () => {
   const imageInputRef = useRef<HTMLInputElement>(null)
   const [showAttacmentMenus, setShowAttacmentMenus] = useState(false)
   const { experimentalFeature } = useContext(FeatureToggleContext)
+  const isGeneratingResponse = useAtomValue(isGeneratingResponseAtom)
+  const threadStates = useAtomValue(threadStatesAtom)
+
+  const isStreamingResponse = Object.values(threadStates).some(
+    (threadState) => threadState.waitingForResponse
+  )
 
   const onPromptChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
     setCurrentPrompt(e.target.value)
@@ -235,7 +243,9 @@ const ChatInput: React.FC = () => {
         accept="application/pdf"
       />
 
-      {messages[messages.length - 1]?.status !== MessageStatus.Pending ? (
+      {messages[messages.length - 1]?.status !== MessageStatus.Pending &&
+      !isGeneratingResponse &&
+      !isStreamingResponse ? (
         <Button
           size="lg"
           disabled={

web/screens/Chat/ThreadList/index.tsx

@@ -49,8 +49,17 @@ export default function ThreadList() {
   useEffect(() => {
     if (threadDataReady && assistants.length > 0 && threads.length === 0) {
       requestCreateNewThread(assistants[0])
+    } else if (threadDataReady && !activeThreadId) {
+      setActiveThread(threads[0])
     }
-  }, [assistants, threads, threadDataReady, requestCreateNewThread])
+  }, [
+    assistants,
+    threads,
+    threadDataReady,
+    requestCreateNewThread,
+    activeThreadId,
+    setActiveThread,
+  ])
 
   return (
     <div className="px-3 py-4">