From 26b123dd0398cc5174e3d2cdf60dd1b59df034b9 Mon Sep 17 00:00:00 2001
From: Yu Chao <yuchao.k@bytedance.com>
Date: Mon, 28 Apr 2025 14:16:31 +0800
Subject: [PATCH] feat: handle exception

---
 web/src/core/store/store.ts | 48 ++++++++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/web/src/core/store/store.ts b/web/src/core/store/store.ts
index 9ebdaa2..94f3762 100644
--- a/web/src/core/store/store.ts
+++ b/web/src/core/store/store.ts
@@ -26,6 +26,9 @@ export const useStore = create<{
   ongoingResearchId: string | null;
   openResearchId: string | null;
 
+  appendMessage: (message: Message) => void;
+  updateMessage: (message: Message) => void;
+  updateMessages: (messages: Message[]) => void;
   setOngoingResearchId: (value: string | null) => void;
 }>((set) => ({
   responding: false,
@@ -39,6 +42,22 @@ export const useStore = create<{
   ongoingResearchId: null,
   openResearchId: null,
 
+  appendMessage(message: Message) {
+    set((state) => ({
+      messageIds: [...state.messageIds, message.id],
+      messages: new Map(state.messages).set(message.id, message),
+    }));
+  },
+  updateMessage(message: Message) {
+    set(state => ({ messages: new Map(state.messages).set(message.id, message) }));
+  },
+  updateMessages(messages: Message[]) {
+    set((state) => {
+      const newMessages = new Map(state.messages);
+      messages.forEach((m) => newMessages.set(m.id, m));
+      return { messages: newMessages };
+    });
+  },
   setOngoingResearchId(value: string | null) {
     set({ ongoingResearchId: value });
   }
@@ -78,10 +97,11 @@ export async function sendMessage(
   );
 
   setResponding(true);
+  let messageId: string | undefined;
   try {
     for await (const event of stream) {
       const { type, data } = event;
-      const messageId = data.id;
+      messageId = data.id;
       let message: Message | undefined;
       if (type === "tool_call_result") {
         message = findMessageByToolCallId(data.tool_call_id);
@@ -105,7 +125,16 @@ export async function sendMessage(
       }
     }
   } catch {
-    setOngoingResearchId(null);
+    // Update message status.
+    // TODO: const isAborted = (error as Error).name === "AbortError";
+    if (messageId != null) {
+      const message = getMessage(messageId);
+      if (message?.isStreaming) {
+        message.isStreaming = false;
+        useStore.getState().updateMessage(message);
+      }
+    }
+    useStore.getState().setOngoingResearchId(null);
   } finally {
     setResponding(false);
   }
@@ -195,10 +224,7 @@ function appendMessage(message: Message) {
   ) {
     appendResearchActivity(message);
   }
-  useStore.setState({
-    messageIds: [...useStore.getState().messageIds, message.id],
-    messages: new Map(useStore.getState().messages).set(message.id, message),
-  });
+  useStore.getState().appendMessage(message);
 }
 
 function updateMessage(message: Message) {
@@ -219,21 +245,15 @@ function updateMessage(message: Message) {
     message.agent === "reporter" &&
     !message.isStreaming
   ) {
-    setOngoingResearchId(null);
+    useStore.getState().setOngoingResearchId(null);
   }
-  useStore.setState({
-    messages: new Map(useStore.getState().messages).set(message.id, message),
-  });
+  useStore.getState().updateMessage(message);
 }
 
 function getOngoingResearchId() {
   return useStore.getState().ongoingResearchId;
 }
 
-function setOngoingResearchId(value: string | null) {
-  useStore.getState().setOngoingResearchId(value);
-}
-
 function appendResearch(researchId: string) {
   let planMessage: Message | undefined;
   const reversedMessageIds = [...useStore.getState().messageIds].reverse();