add stream stop mechanism and unfuck clear conversation while streaming

Author: sutigit

src/client/components/ChatV2/ChatBox.tsx

@@ -28,6 +28,7 @@ export const ChatBox = ({
   handleContinue,
   handleSubmit,
   handleReset,
+  handleStop,
   isMobile,
 }: {
   disabled: boolean
@@ -41,6 +42,7 @@ export const ChatBox = ({
   handleContinue: (message: string) => void
   handleSubmit: (message: string) => void
   handleReset: () => void
+  handleStop: () => void
   isMobile: boolean
 }) => {
   const { courseId } = useParams()
@@ -229,9 +231,16 @@ export const ChatBox = ({
               {fileName && <Chip sx={{ borderRadius: 100 }} label={fileName} onDelete={handleDeleteFile} />}
             </Box>
             <Tooltip title={disabled ? t('chat:cancelResponse') : isShiftEnterSend ? t('chat:shiftEnterSend') : t('chat:enterSend')} arrow placement="top">
-              <IconButton type={disabled ? 'button' : 'submit'} ref={sendButtonRef} data-testid="send-chat-message">
-                {disabled ? <StopIcon /> : <Send />}
-              </IconButton>
+              {
+                disabled ?
+                  <IconButton onClick={handleStop}>
+                    <StopIcon />
+                  </IconButton>
+                  :
+                  <IconButton type={'submit'} ref={sendButtonRef} data-testid="send-chat-message">
+                    <Send />
+                  </IconButton>
+              }
             </Tooltip>
             <SendPreferenceConfiguratorModal
               open={sendPreferenceConfiguratorOpen}

src/client/components/ChatV2/ChatV2.tsx

@@ -29,7 +29,7 @@ import ToolResult from './ToolResult'
 import { OutlineButtonBlack } from './general/Buttons'
 import { ChatInfo } from './general/ChatInfo'
 import { SettingsModal } from './SettingsModal'
-import { useChatStream } from './useChatStream'
+import { StreamAbortReason, TypedAbortController, useChatStream } from './useChatStream'
 import { postCompletionStreamV3 } from './api'
 import PromptSelector from './PromptSelector'
 import ModelSelector from './ModelSelector'
@@ -129,7 +129,7 @@ const ChatV2Content = () => {
 
   const disclaimerInfo = infoTexts?.find((infoText) => infoText.name === 'disclaimer')?.text[i18n.language] ?? null
 
-  const { processStream, completion, isStreaming, setIsStreaming, toolCalls, streamController, generationInfo } = useChatStream({
+  const { processStream, completion, isStreaming, setIsStreaming, toolCalls, streamControllerRef, generationInfo } = useChatStream({
     onComplete: ({ message }) => {
       if (message.content.length > 0) {
         setMessages((prev: ChatMessage[]) => prev.concat(message))
@@ -165,6 +165,8 @@ const ChatV2Content = () => {
       return
     }
 
+    streamControllerRef.current = new TypedAbortController<StreamAbortReason>()
+
     const formData = new FormData()
 
     const file = fileInputRef.current?.files?.[0]
@@ -184,8 +186,8 @@ const ChatV2Content = () => {
     }
     setFileName('')
     setRetryTimeout(() => {
-      if (streamController) {
-        streamController.abort()
+      if (streamControllerRef.current) {
+        streamControllerRef.current.abort("timeout_error")
       }
     }, 5000)
 
@@ -209,7 +211,7 @@ const ChatV2Content = () => {
           },
           courseId,
         },
-        streamController,
+        streamControllerRef.current,
       )
 
       if (!stream && !tokenUsageAnalysis) {
@@ -240,6 +242,7 @@ const ChatV2Content = () => {
 
   const handleReset = () => {
     if (window.confirm(t('chat:emptyConfirm'))) {
+      streamControllerRef.current?.abort("conversation_cleared")
       setMessages([])
       setActiveToolResult(undefined)
       if (fileInputRef.current) {
@@ -248,11 +251,6 @@ const ChatV2Content = () => {
       setFileName('')
       setTokenUsageWarning('')
       setTokenUsageAlertOpen(false)
-      setRetryTimeout(() => {
-        if (streamController) {
-          streamController.abort()
-        }
-      }, 5000)
       clearRetryTimeout()
       dispatchAnalytics({ type: 'RESET_CHAT' })
     }
@@ -482,6 +480,7 @@ const ChatV2Content = () => {
               handleSendMessage(newMessage, false)
             }}
             handleReset={handleReset}
+            handleStop={() => streamControllerRef.current?.abort("user_aborted")}
             isMobile={isMobile}
           />
         </Box>

src/client/components/ChatV2/useChatStream.ts

@@ -1,6 +1,20 @@
-import { useState } from 'react'
+import { useRef, useState } from 'react'
 import type { AssistantMessage, ChatEvent, MessageGenerationInfo, ToolCallResultEvent, ToolCallStatusEvent } from '../../../shared/chat'
 
+export type StreamAbortReason = 'timeout_error' | 'user_aborted' | 'conversation_cleared' | 'error'
+export class TypedAbortController<Reason extends string> {
+  private inner = new AbortController()
+  signal: AbortSignal
+
+  constructor() {
+    this.signal = this.inner.signal
+  }
+
+  abort(reason: Reason) {
+    this.inner.abort(reason)
+  }
+}
+
 type ToolCallState = ToolCallStatusEvent
 
 export const useChatStream = ({
@@ -18,7 +32,7 @@ export const useChatStream = ({
   const [generationInfo, setGenerationInfo] = useState<MessageGenerationInfo | undefined>()
   const [isStreaming, setIsStreaming] = useState(false)
   const [toolCalls, setToolCalls] = useState<Record<string, ToolCallState>>({})
-  const [streamController, setStreamController] = useState<AbortController>()
+  const streamControllerRef = useRef<TypedAbortController<StreamAbortReason>>(null)
 
   const decoder = new TextDecoder()
 
@@ -84,15 +98,6 @@ export const useChatStream = ({
           }
         }
       }
-    } catch (err: unknown) {
-      error += '\nResponse stream was interrupted'
-      onError(err)
-    } finally {
-      setStreamController(undefined)
-      setCompletion('')
-      setToolCalls({})
-      setIsStreaming(false)
-      setGenerationInfo(undefined)
 
       onComplete({
         message: {
@@ -103,6 +108,42 @@ export const useChatStream = ({
           generationInfo: baseGenerationInfo,
         },
       })
+    } catch (err: unknown) {
+      if (err instanceof Error && err.name === 'AbortError') {
+        const reason = streamControllerRef.current?.signal.reason as StreamAbortReason | undefined
+
+        switch (reason) {
+          case 'timeout_error':
+            error += '\nTimeout error'
+            break
+
+          case 'user_aborted':
+            onComplete({
+              message: {
+                role: 'assistant',
+                content,
+                error: error.length > 0 ? error : undefined,
+                toolCalls: toolCallResultsAccum,
+                generationInfo: baseGenerationInfo,
+              },
+            })
+            return
+
+          case 'conversation_cleared':
+            setCompletion('')
+            return
+        }
+      } else {
+        error += '\nResponse stream was interrupted'
+      }
+
+      onError(err)
+    } finally {
+      streamControllerRef.current = null
+      setCompletion('')
+      setToolCalls({})
+      setIsStreaming(false)
+      setGenerationInfo(undefined)
     }
   }
 
@@ -113,6 +154,6 @@ export const useChatStream = ({
     isStreaming,
     setIsStreaming,
     toolCalls,
-    streamController,
+    streamControllerRef,
   }
 }