diff --git a/src/frontend/src/modals/chatModal/index.tsx b/src/frontend/src/modals/chatModal/index.tsx
index 7f6c61cde..b310a36f9 100644
--- a/src/frontend/src/modals/chatModal/index.tsx
+++ b/src/frontend/src/modals/chatModal/index.tsx
@@ -1,5 +1,8 @@
 import { Dialog, Transition } from "@headlessui/react";
-import { ChatBubbleOvalLeftEllipsisIcon, XMarkIcon } from "@heroicons/react/24/outline";
+import {
+  ChatBubbleOvalLeftEllipsisIcon,
+  XMarkIcon,
+} from "@heroicons/react/24/outline";
 import { Fragment, useContext, useEffect, useRef, useState } from "react";
 import { FlowType, NodeType } from "../../types/flow";
 import { alertContext } from "../../contexts/alertContext";
@@ -37,8 +40,7 @@ export default function ChatModal({
   }, [open]);
   useEffect(() => {
     id.current = flow.id;
-  },[flow.id])
-
+  }, [flow.id]);
 
   var isStream = false;
 
@@ -122,16 +124,16 @@ export default function ChatModal({
               newChatHistory.push(
                 chatItem.files
                   ? {
-                    isSend: !chatItem.is_bot,
-                    message: chatItem.message,
-                    thought: chatItem.intermediate_steps,
-                    files: chatItem.files,
-                  }
+                      isSend: !chatItem.is_bot,
+                      message: chatItem.message,
+                      thought: chatItem.intermediate_steps,
+                      files: chatItem.files,
+                    }
                   : {
-                    isSend: !chatItem.is_bot,
-                    message: chatItem.message,
-                    thought: chatItem.intermediate_steps,
-                  }
+                      isSend: !chatItem.is_bot,
+                      message: chatItem.message,
+                      thought: chatItem.intermediate_steps,
+                    }
               );
             }
           }
@@ -144,6 +146,9 @@ export default function ChatModal({
       isStream = true;
     }
     if (data.type === "end") {
+      if (data.message) {
+        updateLastMessage({ str: data.message, end: true });
+      }
       if (data.intermediate_steps) {
         updateLastMessage({
           str: data.message,
@@ -171,8 +176,9 @@ export default function ChatModal({
       const urlWs =
         process.env.NODE_ENV === "development"
           ? `ws://localhost:7860/chat/${id.current}`
-          : `${window.location.protocol === "https:" ? "wss" : "ws"}://${window.location.host
-          }/chat/${id.current}`;
+          : `${window.location.protocol === "https:" ? "wss" : "ws"}://${
+              window.location.host
+            }/chat/${id.current}`;
       const newWs = new WebSocket(urlWs);
       newWs.onopen = () => {
         console.log("WebSocket connection established!");
@@ -189,10 +195,9 @@ export default function ChatModal({
       };
       newWs.onerror = (ev) => {
         console.log(ev, "error");
-        if(flow.id===""){
+        if (flow.id === "") {
           connectWS();
-        }
-        else{
+        } else {
           setErrorData({
             title: "There was an error on web connection, please: ",
             list: [
@@ -205,10 +210,9 @@ export default function ChatModal({
       };
       ws.current = newWs;
     } catch {
-      if(flow.id===""){
+      if (flow.id === "") {
         connectWS();
-      }
-      else{
+      } else {
         setErrorData({
           title: "There was an error on web connection, please: ",
           list: [
@@ -279,11 +283,12 @@ export default function ChatModal({
                   e.targetHandle.split("|")[2] === n.id
               )
             ? [
-              `${type} is missing ${template.display_name
-                ? template.display_name
-                : toNormalCase(template[t].name)
-              }.`,
-            ]
+                `${type} is missing ${
+                  template.display_name
+                    ? template.display_name
+                    : toNormalCase(template[t].name)
+                }.`,
+              ]
             : []
         ),
       [] as string[]
@@ -382,7 +387,9 @@ export default function ChatModal({
                 </div>
                 <div className="w-full h-full bg-white dark:bg-gray-800 border-t dark:border-t-gray-600 flex-col flex items-center overflow-scroll scrollbar-hide">
                   {chatHistory.length > 0 ? (
-                    chatHistory.map((c, i) => <ChatMessage lockChat={lockChat} chat={c} key={i} />)
+                    chatHistory.map((c, i) => (
+                      <ChatMessage lockChat={lockChat} chat={c} key={i} />
+                    ))
                   ) : (
                     <div className="flex flex-col h-full text-center justify-center w-full items-center align-middle">
                       <span>