🏄‍♂️ fix: Handle SSE Stream Edge Case (#8556)

* refactor: Move draft-related utilities to a new `drafts.ts` file

* refactor: auto-save draft logic to use new get/set functions

* fix: Ensure `getDraft` properly decodes stored draft values

* fix: Handle edge case where stream is cancelled before any response, which creates a blank page

client/src/hooks/Input/useAutoSave.ts

@@ -4,36 +4,11 @@ import { useState, useEffect, useMemo, useCallback, useRef } from 'react';
 import { LocalStorageKeys, Constants } from 'librechat-data-provider';
 import type { TFile } from 'librechat-data-provider';
 import type { ExtendedFile } from '~/common';
+import { clearDraft, getDraft, setDraft } from '~/utils';
 import { useChatFormContext } from '~/Providers';
 import { useGetFiles } from '~/data-provider';
 import store from '~/store';
 
-const clearDraft = debounce((id?: string | null) => {
-  localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${id ?? ''}`);
-}, 2500);
-
-const encodeBase64 = (plainText: string): string => {
-  try {
-    const textBytes = new TextEncoder().encode(plainText);
-    return btoa(String.fromCharCode(...textBytes));
-  } catch (e) {
-    return '';
-  }
-};
-
-const decodeBase64 = (base64String: string): string => {
-  try {
-    const bytes = atob(base64String);
-    const uint8Array = new Uint8Array(bytes.length);
-    for (let i = 0; i < bytes.length; i++) {
-      uint8Array[i] = bytes.charCodeAt(i);
-    }
-    return new TextDecoder().decode(uint8Array);
-  } catch (e) {
-    return '';
-  }
-};
-
 export const useAutoSave = ({
   isSubmitting,
   conversationId: _conversationId,
@@ -98,8 +73,11 @@ export const useAutoSave = ({
 
   const restoreText = useCallback(
     (id: string) => {
-      const savedDraft = (localStorage.getItem(`${LocalStorageKeys.TEXT_DRAFT}${id}`) ?? '') || '';
-      setValue('text', decodeBase64(savedDraft));
+      const savedDraft = getDraft(id);
+      if (!savedDraft) {
+        return;
+      }
+      setValue('text', savedDraft);
     },
     [setValue],
   );
@@ -113,10 +91,7 @@ export const useAutoSave = ({
       if (textAreaRef.current.value === '' || textAreaRef.current.value.length === 1) {
         clearDraft(id);
       } else {
-        localStorage.setItem(
-          `${LocalStorageKeys.TEXT_DRAFT}${id}`,
-          encodeBase64(textAreaRef.current.value),
-        );
+        setDraft({ id, value: textAreaRef.current.value });
       }
     },
     [textAreaRef],
@@ -130,16 +105,7 @@ export const useAutoSave = ({
       return;
     }
 
-    const handleInput = debounce((value: string) => {
-      if (value && value.length > 1) {
-        localStorage.setItem(
-          `${LocalStorageKeys.TEXT_DRAFT}${conversationId}`,
-          encodeBase64(value),
-        );
-      } else {
-        localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`);
-      }
-    }, 750);
+    const handleInput = debounce((value: string) => setDraft({ id: conversationId, value }), 750);
 
     const eventListener = (e: Event) => {
       const target = e.target as HTMLTextAreaElement;
@@ -194,10 +160,7 @@ export const useAutoSave = ({
         if (pendingDraft) {
           localStorage.setItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`, pendingDraft);
         } else if (textAreaRef?.current?.value) {
-          localStorage.setItem(
-            `${LocalStorageKeys.TEXT_DRAFT}${conversationId}`,
-            encodeBase64(textAreaRef.current.value),
-          );
+          setDraft({ id: conversationId, value: textAreaRef.current.value });
         }
         const pendingFileDraft = localStorage.getItem(
           `${LocalStorageKeys.FILES_DRAFT}${Constants.PENDING_CONVO}`,

client/src/hooks/SSE/useEventHandlers.ts

@@ -7,10 +7,10 @@ import {
   QueryKeys,
   Constants,
   EndpointURLs,
+  ContentTypes,
   tPresetSchema,
   tMessageSchema,
   tConvoUpdateSchema,
-  ContentTypes,
   isAssistantsEndpoint,
 } from 'librechat-data-provider';
 import type { TMessage, TConversation, EventSubmission } from 'librechat-data-provider';
@@ -21,6 +21,7 @@ import type { SetterOrUpdater, Resetter } from 'recoil';
 import type { ConversationCursorData } from '~/utils';
 import {
   logger,
+  setDraft,
   scrollToEnd,
   getAllContentText,
   addConvoToAllQueries,
@@ -457,6 +458,38 @@ export default function useEventHandlers({
       announcePolite({ message: 'end', isStatus: true });
       announcePolite({ message: getAllContentText(responseMessage) });
 
+      const isNewConvo = conversation.conversationId !== submissionConvo.conversationId;
+
+      const setFinalMessages = (id: string | null, _messages: TMessage[]) => {
+        setMessages(_messages);
+        queryClient.setQueryData<TMessage[]>([QueryKeys.messages, id], _messages);
+      };
+
+      /** Handle edge case where stream is cancelled before any response, which creates a blank page */
+      if (
+        !conversation.conversationId &&
+        responseMessage?.content?.[0]?.['text']?.value ===
+          submission.initialResponse?.content?.[0]?.['text']?.value
+      ) {
+        const currentConvoId =
+          (submissionConvo.conversationId ?? conversation.conversationId) || Constants.NEW_CONVO;
+        if (isNewConvo && submissionConvo.conversationId) {
+          removeConvoFromAllQueries(queryClient, submissionConvo.conversationId);
+        }
+
+        const isNewChat =
+          location.pathname === `/c/${Constants.NEW_CONVO}` &&
+          currentConvoId === Constants.NEW_CONVO;
+
+        setFinalMessages(currentConvoId, isNewChat ? [] : [...messages]);
+        setDraft({ id: currentConvoId, value: requestMessage?.text });
+        setIsSubmitting(false);
+        if (isNewChat) {
+          navigate(`/c/${Constants.NEW_CONVO}`, { replace: true, state: { focusChat: true } });
+        }
+        return;
+      }
+
       /* Update messages; if assistants endpoint, client doesn't receive responseMessage */
       let finalMessages: TMessage[] = [];
       if (runMessages) {
@@ -467,11 +500,7 @@ export default function useEventHandlers({
         finalMessages = [...messages, requestMessage, responseMessage];
       }
       if (finalMessages.length > 0) {
-        setMessages(finalMessages);
-        queryClient.setQueryData<TMessage[]>(
-          [QueryKeys.messages, conversation.conversationId],
-          finalMessages,
-        );
+        setFinalMessages(conversation.conversationId, finalMessages);
       } else if (
         isAssistantsEndpoint(submissionConvo.endpoint) &&
         (!submissionConvo.conversationId || submissionConvo.conversationId === Constants.NEW_CONVO)
@@ -482,9 +511,8 @@ export default function useEventHandlers({
         );
       }
 
-      const isNewConvo = conversation.conversationId !== submissionConvo.conversationId;
-      if (isNewConvo) {
-        removeConvoFromAllQueries(queryClient, submissionConvo.conversationId as string);
+      if (isNewConvo && submissionConvo.conversationId) {
+        removeConvoFromAllQueries(queryClient, submissionConvo.conversationId);
       }
 
       /* Refresh title */
@@ -527,7 +555,7 @@ export default function useEventHandlers({
           );
         }
 
-        if (location.pathname === '/c/new') {
+        if (location.pathname === `/c/${Constants.NEW_CONVO}`) {
           navigate(`/c/${conversation.conversationId}`, { replace: true });
         }
       }

client/src/utils/drafts.ts

@@ -0,0 +1,39 @@
+import debounce from 'lodash/debounce';
+import { LocalStorageKeys } from 'librechat-data-provider';
+
+export const clearDraft = debounce((id?: string | null) => {
+  localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${id ?? ''}`);
+}, 2500);
+
+export const encodeBase64 = (plainText: string): string => {
+  try {
+    const textBytes = new TextEncoder().encode(plainText);
+    return btoa(String.fromCharCode(...textBytes));
+  } catch {
+    return '';
+  }
+};
+
+export const decodeBase64 = (base64String: string): string => {
+  try {
+    const bytes = atob(base64String);
+    const uint8Array = new Uint8Array(bytes.length);
+    for (let i = 0; i < bytes.length; i++) {
+      uint8Array[i] = bytes.charCodeAt(i);
+    }
+    return new TextDecoder().decode(uint8Array);
+  } catch {
+    return '';
+  }
+};
+
+export const setDraft = ({ id, value }: { id: string; value?: string }) => {
+  if (value && value.length > 1) {
+    localStorage.setItem(`${LocalStorageKeys.TEXT_DRAFT}${id}`, encodeBase64(value));
+    return;
+  }
+  localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${id}`);
+};
+
+export const getDraft = (id?: string): string | null =>
+  decodeBase64((localStorage.getItem(`${LocalStorageKeys.TEXT_DRAFT}${id ?? ''}`) ?? '') || '');

client/src/utils/index.ts

@@ -6,6 +6,7 @@ export * from './files';
 export * from './latex';
 export * from './theme';
 export * from './forms';
+export * from './drafts';
 export * from './convos';
 export * from './presets';
 export * from './prompts';