fix: clean up stuck streaming messages after refresh

mlc-ai · May 22, 2024 · 33e62aa · 33e62aa
1 parent cfb78d2
commit 33e62aa
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 14 deletions.
diff --git a/app/client/webllm.ts b/app/client/webllm.ts
@@ -168,4 +168,4 @@ export class WebLLMApi implements LLMApi {
   }
 }
 
-export const WebLLMContext = createContext<WebLLMApi | null>(null);
+export const WebLLMContext = createContext<WebLLMApi | undefined>(undefined);
diff --git a/app/components/chat.tsx b/app/components/chat.tsx
@@ -669,7 +669,7 @@ function _Chat() {
     }
   };
 
-  const doSubmit = (userInput: string) => {
+  const onSubmit = (userInput: string) => {
     if (userInput.trim() === "") return;
 
     const matchCommand = chatCommands.match(userInput);
@@ -681,6 +681,7 @@ function _Chat() {
     }
 
     if (isStreaming) return;
+
     setIsLoading(true);
     chatStore
       .onUserInput(userInput, webllm, attachImages)
@@ -740,7 +741,7 @@ function _Chat() {
           }
         }
       });
-      session.messages.filter((m) => m.content.length > 0);
+      session.messages = session.messages.filter((m) => m.content.length > 0);
     });
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
@@ -758,7 +759,7 @@ function _Chat() {
       return;
     }
     if (shouldSubmit(e) && promptHints.length === 0) {
-      doSubmit(userInput);
+      onSubmit(userInput);
       e.preventDefault();
     }
   };
@@ -947,7 +948,7 @@ function _Chat() {
   useCommand({
     fill: setUserInput,
     submit: (text) => {
-      doSubmit(text);
+      onSubmit(text);
     },
   });
 
@@ -1432,7 +1433,7 @@ function _Chat() {
               text={Locale.Chat.Send}
               className={styles["chat-input-send"]}
               type="primary"
-              onClick={() => doSubmit(userInput)}
+              onClick={() => onSubmit(userInput)}
             />
           )}
         </label>

diff --git a/app/components/home.tsx b/app/components/home.tsx
@@ -28,6 +28,7 @@ import { useAppConfig } from "../store/config";
 import { getClientConfig } from "../config/client";
 import { WebLLMApi, WebLLMContext } from "../client/webllm";
 import Locale from "../locales";
+import { useChatStore } from "../store";
 
 export function Loading(props: { noLogo?: boolean }) {
   return (
@@ -181,7 +182,7 @@ function Screen() {
 }
 
 const useWebLLM = () => {
-  const [webllm, setWebLLM] = useState<WebLLMApi | null>(null);
+  const [webllm, setWebLLM] = useState<WebLLMApi | undefined>(undefined);
   const [isSWAlive, setSWAlive] = useState(true);
 
   useEffect(() => {
@@ -232,6 +233,15 @@ const useLoadUrlParam = () => {
   }, []);
 };
 
+const useStopStreamingMessages = () => {
+  const chatStore = useChatStore();
+
+  // Clean up bad chat messages due to refresh during generating
+  useEffect(() => {
+    chatStore.stopStreaming();
+  }, []);
+};
+
 export function Home() {
   const hasHydrated = useHasHydrated();
   const isServiceWorkerReady = useServiceWorkerReady();
@@ -240,6 +250,7 @@ export function Home() {
   useSwitchTheme();
   useHtmlLang();
   useLoadUrlParam();
+  useStopStreamingMessages();
 
   if (!hasHydrated || !isServiceWorkerReady) {
     return <Loading />;

diff --git a/app/store/chat.ts b/app/store/chat.ts
@@ -630,14 +630,24 @@ export const useChatStore = createPersistStore(
       },
 
       stopStreaming() {
-        get().updateCurrentSession((session) => {
-          session.messages = session.messages
-            .map((m: ChatMessage) => ({
-              ...m,
-              streaming: false,
-            }))
-            .filter((m: ChatMessage) => !!m.content);
+        const sessions = get().sessions;
+        sessions.forEach((session) => {
+          if (session.messages.length === 0) {
+            return;
+          }
+          const lastMessage = session.messages[session.messages.length - 1];
+          if (
+            lastMessage.role === "assistant" &&
+            lastMessage.streaming &&
+            lastMessage.content.length === 0
+          ) {
+            // This message generation is interrupted by refresh and is stuck
+            session.messages.splice(session.messages.length - 1, 1);
+          }
+          // Reset streaming status for all messages
+          session.messages.forEach((m) => (m.streaming = false));
         });
+        set(() => ({ sessions }));
       },
 
       updateStat(message: ChatMessage) {