From dbe067ec86f294fe49682a08d6efa644960829cf Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Tue, 2 Apr 2024 13:19:14 -0300
Subject: [PATCH] Add error handling on streaming and fix streaming bug on
 error

---
 .../IOModal/components/chatView/chatMessage/index.tsx    | 9 ++++++++-
 .../src/modals/IOModal/components/chatView/index.tsx     | 2 +-
 src/frontend/src/modals/shareModal/index.tsx             | 1 -
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
index 3d24353ce..b8c37f37b 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
@@ -13,6 +13,8 @@ import useFlowStore from "../../../../../stores/flowStore";
 import { chatMessagePropsType } from "../../../../../types/components";
 import { classNames, cn } from "../../../../../utils/utils";
 import FileCard from "../fileComponent";
+import useAlertStore from "../../../../../stores/alertStore";
+import { Message } from "@radix-ui/react-form";
 
 export default function ChatMessage({
   chat,
@@ -33,6 +35,7 @@ export default function ChatMessage({
   const [isStreaming, setIsStreaming] = useState(false);
   const eventSource = useRef<EventSource | undefined>(undefined);
   const updateFlowPool = useFlowStore((state) => state.updateFlowPool);
+  const setErrorData = useAlertStore((state) => state.setErrorData);
   const chatMessageRef = useRef(chatMessage);
 
   // Sync ref with state
@@ -53,10 +56,14 @@ export default function ChatMessage({
           setChatMessage((prev) => prev + parsedData.chunk);
         }
       };
-      eventSource.current.onerror = (event) => {
+      eventSource.current.onerror = (event:any) => {
         setIsStreaming(false);
         eventSource.current?.close();
         setStreamUrl(undefined);
+        if(JSON.parse(event.data)?.error){
+          setErrorData({title: "Error on Streaming", list: [JSON.parse(event.data)?.error]});
+        }
+        updateChat(chat, chatMessageRef.current);
         reject(new Error("Streaming failed"));
       };
       eventSource.current.addEventListener("close", (event) => {
diff --git a/src/frontend/src/modals/IOModal/components/chatView/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
index 36a17de47..046f26a74 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
@@ -124,7 +124,7 @@ export default function ChatView({
     message: string,
     stream_url?: string
   ) {
-    if (message === "") return;
+    // if (message === "") return;
     chat.message = message;
     // chat is one of the chatHistory
     updateFlowPool(chat.componentId, {
diff --git a/src/frontend/src/modals/shareModal/index.tsx b/src/frontend/src/modals/shareModal/index.tsx
index 83dbb7137..8272bddbe 100644
--- a/src/frontend/src/modals/shareModal/index.tsx
+++ b/src/frontend/src/modals/shareModal/index.tsx
@@ -143,7 +143,6 @@ export default function ShareModal({
         });
       });
   };
-  console.log("ShareModal");
 
   const handleUpdateComponent = () => {
     handleShareComponent(true);