Ensure continuous generating state during chat/tool call transitions

Author: aiqlcom

src/renderer/composables/chatCompletions.ts

@@ -28,6 +28,8 @@ interface ChatRequestBody {
   tools?: Tool[]
 }
 
+export type ChatProcessResult = 'aborted' | 'error' | 'done'
+
 type RequestMessageType = ChatCompletionRequestMessage | McpSamplingMessage
 
 const THINK_OPEN = '<think>'
@@ -113,7 +115,7 @@ export const createCompletion = async (
   rawconversation: RequestMessageType[],
   sessionId: string,
   sampling: any = null
-) => {
+): Promise<ChatProcessResult> => {
   const snackbarStore = useSnackbarStore()
 
   const messageStore = useMessageStore()
@@ -250,7 +252,7 @@ export const createCompletion = async (
         }
       } finally {
         snackbarStore.showErrorMessage(errMessage)
-        return
+        return 'error'
       }
     }
 
@@ -262,7 +264,7 @@ export const createCompletion = async (
     const reader = completion.body?.getReader()
     if (!reader) {
       snackbarStore.showErrorMessage('snackbar.parse-stream-fail')
-      return
+      return 'error'
     }
 
     // Add the bot message
@@ -284,10 +286,8 @@ export const createCompletion = async (
   } catch (error: any) {
     snackbarStore.showErrorMessage(error?.message)
   } finally {
-    if (sessionId in messageStore.generating) {
-      messageStore.generating[sessionId].abort()
-      delete messageStore.generating[sessionId]
-    }
+    const result = messageStore.delete(sessionId) ? 'done' : 'aborted'
+    return result
   }
 }
 
@@ -310,10 +310,10 @@ const read = async (
 
   // If the stream is done reading, release the lock on the reader
   if (done) {
-    if (sessionId in messageStore.generating) {
-      messageStore.generating[sessionId].abort()
-      delete messageStore.generating[sessionId]
-    }
+    // if (sessionId in messageStore.generating) {
+    //   messageStore.generating[sessionId].abort()
+    //   delete messageStore.generating[sessionId]
+    // }
     return reader.releaseLock()
   }
 

src/renderer/store/message.ts

@@ -2,7 +2,11 @@ import { defineStore } from 'pinia'
 import { useSnackbarStore } from '@/renderer/store/snackbar'
 import { useMcpStore } from '@/renderer/store/mcp'
 import { useHistoryStore } from '@/renderer/store/history'
-import { createCompletion, isEmptyTools } from '@/renderer/composables/chatCompletions'
+import {
+  createCompletion,
+  isEmptyTools,
+  ChatProcessResult
+} from '@/renderer/composables/chatCompletions'
 
 import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
 
@@ -15,7 +19,7 @@ interface MessageStoreState {
   conversation: ChatConversationMessage[]
   historyId: string
   base64: string
-  generating: Record<string, AbortController>
+  generating: Record<string, 'prepare' | AbortController | 'toolcall'>
 }
 
 export const useMessageStore = defineStore('messageStore', {
@@ -41,12 +45,21 @@ export const useMessageStore = defineStore('messageStore', {
       this.historyId = ''
     },
     stop() {
-      const id = this.historyId
+      this.delete(this.historyId)
       const snackbarStore = useSnackbarStore()
-      this.generating[id].abort()
-      delete this.generating[id]
       snackbarStore.showInfoMessage('snackbar.stopped')
     },
+    delete(id: string) {
+      if (id in this.generating) {
+        if (this.generating[id] instanceof AbortController) {
+          this.generating[id].abort()
+        }
+        delete this.generating[id]
+        return true
+      } else {
+        return false
+      }
+    },
     clear() {
       this.userMessage = ''
     },
@@ -120,13 +133,17 @@ export const useMessageStore = defineStore('messageStore', {
       const historyId = this.syncHistory()
       this.clear()
 
-      createCompletion(this.conversation, historyId).then(() => {
-        this.postToolCall()
+      this.generating[historyId] = 'prepare'
+
+      createCompletion(this.conversation, historyId).then((reason: ChatProcessResult) => {
+        if (reason === 'done') {
+          this.postToolCall(historyId)
+        }
       })
 
       return historyId
     },
-    postToolCall: async function () {
+    postToolCall: async function (historyId: string) {
       const mcpStore = useMcpStore()
       const last = this.conversation.at(-1)
 
@@ -141,6 +158,7 @@ export const useMessageStore = defineStore('messageStore', {
 
       let toolCalled = false
       console.log(last.tool_calls)
+      this.generating[historyId] = 'toolcall'
 
       for (const toolCall of last.tool_calls) {
         let result: CallToolResult
@@ -160,7 +178,7 @@ export const useMessageStore = defineStore('messageStore', {
         }
       }
 
-      if (toolCalled) {
+      if (this.delete(historyId) && toolCalled) {
         this.startInference()
       }
     },