From 91a9361dde873a2aab462c0764445c614f32f4f7 Mon Sep 17 00:00:00 2001
From: Kwanghyun On <kwanghyun.on@gmail.com>
Date: Sun, 17 Nov 2024 22:26:22 -0500
Subject: [PATCH 1/2] feat: Add stop generation button during AI responses

---
 src/components/chat-view/Chat.tsx | 14 ++++++++++++--
 styles.css                        | 12 ++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/components/chat-view/Chat.tsx b/src/components/chat-view/Chat.tsx
index 3c8f07e..4acfb3d 100644
--- a/src/components/chat-view/Chat.tsx
+++ b/src/components/chat-view/Chat.tsx
@@ -1,5 +1,5 @@
 import { useMutation } from '@tanstack/react-query'
-import { History, Plus } from 'lucide-react'
+import { CircleStop, History, Plus } from 'lucide-react'
 import { App, Notice } from 'obsidian'
 import {
   forwardRef,
@@ -118,9 +118,10 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
   const [queryProgress, setQueryProgress] = useState<QueryProgressState>({
     type: 'idle',
   })
+  const [isStreaming, setIsStreaming] = useState(false)
+
   const preventAutoScrollRef = useRef(false)
   const lastProgrammaticScrollRef = useRef<number>(0)
-
   const activeStreamAbortControllersRef = useRef<AbortController[]>([])
   const chatUserInputRefs = useRef<Map<string, ChatUserInputRef>>(new Map())
   const chatMessagesRef = useRef<HTMLDivElement>(null)
@@ -171,6 +172,7 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
       abortController.abort()
     }
     activeStreamAbortControllersRef.current = []
+    setIsStreaming(false)
   }
 
   const handleLoadConversation = async (conversationId: string) => {
@@ -241,6 +243,7 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
       try {
         const abortController = new AbortController()
         activeStreamAbortControllersRef.current.push(abortController)
+        setIsStreaming(true)
 
         const { requestMessages, compiledMessages } =
           await promptGenerator.generateRequestMessages({
@@ -280,6 +283,7 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
             handleScrollToBottom()
           }
         }
+        setIsStreaming(false)
       } catch (error) {
         if (error.name === 'AbortError') {
           return
@@ -594,6 +598,12 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
           ),
         )}
         <QueryProgress state={queryProgress} />
+        {isStreaming && (
+          <button onClick={abortActiveStreams} className="smtcmp-stop-gen-btn">
+            <CircleStop size={16} />
+            <div>Stop Generation</div>
+          </button>
+        )}
       </div>
       <ChatUserInput
         key={inputMessage.id} // this is needed to clear the editor when the user submits a new message
diff --git a/styles.css b/styles.css
index 5a6724f..b248fac 100644
--- a/styles.css
+++ b/styles.css
@@ -126,8 +126,20 @@ button:not(.clickable-icon).smtcmp-chat-list-dropdown {
 
 .smtcmp-chat-container {
   display: flex;
+  position: relative;
   flex-direction: column;
   height: 100%;
+
+  .smtcmp-stop-gen-btn {
+    z-index: 1000;
+    position: absolute;
+    bottom: 160px;
+    left: 50%;
+    transform: translateX(-50%);
+    display: flex;
+    align-items: center;
+    gap: var(--size-4-1);
+  }
 }
 
 .smtcmp-chat-messages {

From 3274ca62d0bd7a0339bccd51ed0c08bea3bb9ef0 Mon Sep 17 00:00:00 2001
From: Kevin On <40454531+kevin-on@users.noreply.github.com>
Date: Wed, 20 Nov 2024 14:47:57 +0900
Subject: [PATCH 2/2] fix: use submitMutation.isPending instead of separate
 streaming state

---
 src/components/chat-view/Chat.tsx | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/components/chat-view/Chat.tsx b/src/components/chat-view/Chat.tsx
index 4acfb3d..ea0eb57 100644
--- a/src/components/chat-view/Chat.tsx
+++ b/src/components/chat-view/Chat.tsx
@@ -118,7 +118,6 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
   const [queryProgress, setQueryProgress] = useState<QueryProgressState>({
     type: 'idle',
   })
-  const [isStreaming, setIsStreaming] = useState(false)
 
   const preventAutoScrollRef = useRef(false)
   const lastProgrammaticScrollRef = useRef<number>(0)
@@ -172,7 +171,6 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
       abortController.abort()
     }
     activeStreamAbortControllersRef.current = []
-    setIsStreaming(false)
   }
 
   const handleLoadConversation = async (conversationId: string) => {
@@ -243,7 +241,6 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
       try {
         const abortController = new AbortController()
         activeStreamAbortControllersRef.current.push(abortController)
-        setIsStreaming(true)
 
         const { requestMessages, compiledMessages } =
           await promptGenerator.generateRequestMessages({
@@ -283,7 +280,6 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
             handleScrollToBottom()
           }
         }
-        setIsStreaming(false)
       } catch (error) {
         if (error.name === 'AbortError') {
           return
@@ -598,7 +594,7 @@ const Chat = forwardRef<ChatRef, ChatProps>((props, ref) => {
           ),
         )}
         <QueryProgress state={queryProgress} />
-        {isStreaming && (
+        {submitMutation.isPending && (
           <button onClick={abortActiveStreams} className="smtcmp-stop-gen-btn">
             <CircleStop size={16} />
             <div>Stop Generation</div>