From 166a4fa44f473da8960a3cd2fd1268e883b491ff Mon Sep 17 00:00:00 2001
From: Danny Avila <110412045+danny-avila@users.noreply.github.com>
Date: Thu, 30 Nov 2023 10:23:57 -0500
Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20fix:=20Improve=20SSE=20?=
 =?UTF-8?q?Handling=20and=20Fix=20Typo=20in=20`sendEmail`=20Template=20(#1?=
 =?UTF-8?q?245)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix: typo for passwordReset.handlebars

* fix(useSSE): prevent unnecessary JSON.parse abort error, handle immediate abort-submit gracefully by reverting to previous state before immediate abort-submit, add showStopButton state to explicitly render disabled sendButton when message generation is cancelled, filter undefined messages and replace undefined convo for cancelHandler
---
 api/server/middleware/abortMiddleware.js      |  2 +-
 api/server/services/AuthService.js            |  2 +-
 client/src/components/Chat/Input/ChatForm.tsx |  8 +++---
 .../src/components/Chat/Input/StopButton.tsx  |  7 +++--
 client/src/hooks/useChatHelpers.ts            |  4 +++
 client/src/hooks/useSSE.ts                    | 27 ++++++++++++-------
 6 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js
index 1177d71d99..ea331bcbbd 100644
--- a/api/server/middleware/abortMiddleware.js
+++ b/api/server/middleware/abortMiddleware.js
@@ -8,7 +8,7 @@ async function abortMessage(req, res) {
   const { abortKey } = req.body;
 
   if (!abortControllers.has(abortKey) && !res.headersSent) {
-    return res.status(404).send('Request not found');
+    return res.status(404).send({ message: 'Request not found' });
   }
 
   const { abortController } = abortControllers.get(abortKey);
diff --git a/api/server/services/AuthService.js b/api/server/services/AuthService.js
index 06273826cc..363042a753 100644
--- a/api/server/services/AuthService.js
+++ b/api/server/services/AuthService.js
@@ -185,7 +185,7 @@ const resetPassword = async (userId, token, password) => {
     {
       name: user.name,
     },
-    'resetPassword.handlebars',
+    'passwordReset.handlebars',
   );
 
   await passwordResetToken.deleteOne();
diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx
index 062b5cd60c..632944eb8e 100644
--- a/client/src/components/Chat/Input/ChatForm.tsx
+++ b/client/src/components/Chat/Input/ChatForm.tsx
@@ -19,6 +19,8 @@ export default function ChatForm({ index = 0 }) {
     handleStopGenerating,
     filesLoading,
     setFilesLoading,
+    showStopButton,
+    setShowStopButton,
   } = useChatContext();
 
   const submitMessage = () => {
@@ -46,10 +48,10 @@ export default function ChatForm({ index = 0 }) {
               endpoint={conversation?.endpoint}
             />
             <AttachFile endpoint={conversation?.endpoint ?? ''} />
-            {isSubmitting ? (
-              <StopButton stop={handleStopGenerating} />
+            {isSubmitting && showStopButton ? (
+              <StopButton stop={handleStopGenerating} setShowStopButton={setShowStopButton} />
             ) : (
-              <SendButton text={text} disabled={filesLoading} />
+              <SendButton text={text} disabled={filesLoading || isSubmitting} />
             )}
           </div>
         </div>
diff --git a/client/src/components/Chat/Input/StopButton.tsx b/client/src/components/Chat/Input/StopButton.tsx
index 0ce8416ffa..69303adb97 100644
--- a/client/src/components/Chat/Input/StopButton.tsx
+++ b/client/src/components/Chat/Input/StopButton.tsx
@@ -1,4 +1,4 @@
-export default function StopButton({ stop }) {
+export default function StopButton({ stop, setShowStopButton }) {
   return (
     <div className="absolute bottom-0 right-2 top-0 p-1 md:right-3 md:p-2">
       <div className="flex h-full">
@@ -7,7 +7,10 @@ export default function StopButton({ stop }) {
             type="button"
             className="border-gizmo-gray-950 rounded-full border-2 p-1 dark:border-gray-200"
             aria-label="Stop generating"
-            onClick={stop}
+            onClick={(e) => {
+              setShowStopButton(false);
+              stop(e);
+            }}
           >
             <svg
               xmlns="http://www.w3.org/2000/svg"
diff --git a/client/src/hooks/useChatHelpers.ts b/client/src/hooks/useChatHelpers.ts
index a70fd7bb48..ef559aa776 100644
--- a/client/src/hooks/useChatHelpers.ts
+++ b/client/src/hooks/useChatHelpers.ts
@@ -25,6 +25,7 @@ import store from '~/store';
 // this to be set somewhere else
 export default function useChatHelpers(index = 0, paramId: string | undefined) {
   const [files, setFiles] = useRecoilState(store.filesByIndex(index));
+  const [showStopButton, setShowStopButton] = useState(true);
   const [filesLoading, setFilesLoading] = useState(false);
   const setFilesToDelete = useSetFilesToDelete();
 
@@ -130,6 +131,7 @@ export default function useChatHelpers(index = 0, paramId: string | undefined) {
       isEdited = false,
     } = {},
   ) => {
+    setShowStopButton(true);
     if (!!isSubmitting || text === '') {
       return;
     }
@@ -369,5 +371,7 @@ export default function useChatHelpers(index = 0, paramId: string | undefined) {
     invalidateConvos,
     filesLoading,
     setFilesLoading,
+    showStopButton,
+    setShowStopButton,
   };
 }
diff --git a/client/src/hooks/useSSE.ts b/client/src/hooks/useSSE.ts
index 908bf9afd7..e50070c145 100644
--- a/client/src/hooks/useSSE.ts
+++ b/client/src/hooks/useSSE.ts
@@ -17,7 +17,7 @@ import useChatHelpers from './useChatHelpers';
 import useSetStorage from './useSetStorage';
 
 type TResData = {
-  plugin: TResPlugin;
+  plugin?: TResPlugin;
   final?: boolean;
   initial?: boolean;
   requestMessage: TMessage;
@@ -91,16 +91,19 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
     const { requestMessage, responseMessage, conversation } = data;
     const { messages, isRegenerate = false } = submission;
 
+    const convoUpdate = conversation ?? submission.conversation;
+
     // update the messages
     if (isRegenerate) {
-      setMessages([...messages, responseMessage]);
+      const messagesUpdate = [...messages, responseMessage].filter((msg) => msg);
+      setMessages(messagesUpdate);
     } else {
-      setMessages([...messages, requestMessage, responseMessage]);
+      const messagesUpdate = [...messages, requestMessage, responseMessage].filter((msg) => msg);
+      setMessages(messagesUpdate);
     }
-    setIsSubmitting(false);
 
     // refresh title
-    if (requestMessage.parentMessageId == '00000000-0000-0000-0000-000000000000') {
+    if (requestMessage?.parentMessageId == '00000000-0000-0000-0000-000000000000') {
       setTimeout(() => {
         invalidateConvos();
       }, 2000);
@@ -114,12 +117,14 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
     setConversation((prevState) => {
       const update = {
         ...prevState,
-        ...conversation,
+        ...convoUpdate,
       };
 
       setStorage(update);
       return update;
     });
+
+    setIsSubmitting(false);
   };
 
   const createdHandler = (data: TResData, submission: TSubmission) => {
@@ -176,7 +181,6 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
     } else {
       setMessages([...messages, requestMessage, responseMessage]);
     }
-    setIsSubmitting(false);
 
     // refresh title
     if (requestMessage.parentMessageId == '00000000-0000-0000-0000-000000000000') {
@@ -199,6 +203,8 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
       setStorage(update);
       return update;
     });
+
+    setIsSubmitting(false);
   };
 
   const errorHandler = (data: TResData, submission: TSubmission) => {
@@ -210,11 +216,13 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
       error: true,
       parentMessageId: message?.messageId,
     });
-    setIsSubmitting(false);
+
     setMessages([...messages, message, errorResponse]);
     if (data.conversationId && paramId === 'new') {
       newConversation({ template: { conversationId: data.conversationId } });
     }
+
+    setIsSubmitting(false);
     return;
   };
 
@@ -240,7 +248,7 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
       .catch((error) => {
         console.error('Error aborting request');
         console.error(error);
-        // errorHandler({ text: 'Error aborting request' }, { ...submission, message });
+        setIsSubmitting(false);
       });
     return;
   };
@@ -324,7 +332,6 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
         const e = new Event('cancel');
         events.dispatchEvent(e);
       }
-      setIsSubmitting(false);
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [submission]);