From 9ea9b7d87d65d6bd9e00496547992252c4555a0e Mon Sep 17 00:00:00 2001
From: Dinh Long Nguyen <dinhlongviolin1@gmail.com>
Date: Tue, 19 Aug 2025 14:45:57 +0700
Subject: [PATCH] handle abort properly + finally clause to resolve (#6227)

---
 web-app/src/hooks/useChat.ts | 92 +++++++++++++++++++++++-------------
 1 file changed, 60 insertions(+), 32 deletions(-)

diff --git a/web-app/src/hooks/useChat.ts b/web-app/src/hooks/useChat.ts
index d77d507c5..e90d5b1c2 100644
--- a/web-app/src/hooks/useChat.ts
+++ b/web-app/src/hooks/useChat.ts
@@ -309,7 +309,7 @@ export const useChat = () => {
               let pendingDeltaCount = 0
               const reasoningProcessor = new ReasoningProcessor()
               const scheduleFlush = () => {
-                if (rafScheduled) return
+                if (rafScheduled || abortController.signal.aborted) return
                 rafScheduled = true
                 const doSchedule = (cb: () => void) => {
                   if (typeof requestAnimationFrame !== 'undefined') {
@@ -321,6 +321,12 @@ export const useChat = () => {
                   }
                 }
                 doSchedule(() => {
+                  // Check abort status before executing the scheduled callback
+                  if (abortController.signal.aborted) {
+                    rafScheduled = false
+                    return
+                  }
+
                   const currentContent = newAssistantThreadContent(
                     activeThread.id,
                     accumulatedText,
@@ -367,41 +373,63 @@ export const useChat = () => {
                 pendingDeltaCount = 0
                 rafScheduled = false
               }
-              for await (const part of completion) {
-                // Error message
-                if (!part.choices) {
-                  throw new Error(
-                    'message' in part
-                      ? (part.message as string)
-                      : (JSON.stringify(part) ?? '')
-                  )
+              try {
+                for await (const part of completion) {
+                  // Check if aborted before processing each part
+                  if (abortController.signal.aborted) {
+                    break
+                  }
+
+                  // Error message
+                  if (!part.choices) {
+                    throw new Error(
+                      'message' in part
+                        ? (part.message as string)
+                        : (JSON.stringify(part) ?? '')
+                    )
+                  }
+
+                  if (part.choices[0]?.delta?.tool_calls) {
+                    extractToolCall(part, currentCall, toolCalls)
+                    // Schedule a flush to reflect tool update
+                    scheduleFlush()
+                  }
+                  const deltaReasoning =
+                    reasoningProcessor.processReasoningChunk(part)
+                  if (deltaReasoning) {
+                    accumulatedText += deltaReasoning
+                    pendingDeltaCount += 1
+                    // Schedule flush for reasoning updates
+                    scheduleFlush()
+                  }
+                  const deltaContent = part.choices[0]?.delta?.content || ''
+                  if (deltaContent) {
+                    accumulatedText += deltaContent
+                    pendingDeltaCount += 1
+                    // Batch UI update on next animation frame
+                    scheduleFlush()
+                  }
+                }
+              } finally {
+                // Always clean up scheduled RAF when stream ends (either normally or via abort)
+                if (rafHandle !== undefined) {
+                  if (typeof cancelAnimationFrame !== 'undefined') {
+                    cancelAnimationFrame(rafHandle)
+                  } else {
+                    clearTimeout(rafHandle)
+                  }
+                  rafHandle = undefined
+                  rafScheduled = false
                 }
 
-                if (part.choices[0]?.delta?.tool_calls) {
-                  extractToolCall(part, currentCall, toolCalls)
-                  // Schedule a flush to reflect tool update
-                  scheduleFlush()
-                }
-                const deltaReasoning =
-                  reasoningProcessor.processReasoningChunk(part)
-                if (deltaReasoning) {
-                  accumulatedText += deltaReasoning
-                  pendingDeltaCount += 1
-                  // Schedule flush for reasoning updates
-                  scheduleFlush()
-                }
-                const deltaContent = part.choices[0]?.delta?.content || ''
-                if (deltaContent) {
-                  accumulatedText += deltaContent
-                  pendingDeltaCount += 1
-                  // Batch UI update on next animation frame
-                  scheduleFlush()
+                // Only finalize and flush if not aborted
+                if (!abortController.signal.aborted) {
+                  // Finalize reasoning (close any open think tags)
+                  accumulatedText += reasoningProcessor.finalize()
+                  // Ensure any pending buffered content is rendered at the end
+                  flushIfPending()
                 }
               }
-              // Finalize reasoning (close any open think tags)
-              accumulatedText += reasoningProcessor.finalize()
-              // Ensure any pending buffered content is rendered at the end
-              flushIfPending()
             }
           } catch (error) {
             const errorMessage =