fixed onClose method websocket, streaming for text working

src/frontend/src/modals/chatModal/index.tsx

@@ -26,9 +26,15 @@ export default function ChatModal({
 	const [chatHistory, setChatHistory] = useState<ChatMessageType[]>([]);
 	const { reactFlowInstance } = useContext(typesContext);
 	const { setErrorData, setNoticeData } = useContext(alertContext);
-	const [isStreaming,setIsStreaming] = useState()
+	const ws = useRef<WebSocket|null>(null)
-	const [ws, setWs] = useState<WebSocket | null>(null);
 	const [lockChat, setLockChat] = useState(false);
+	const isOpen = useRef(open);
+
+	useEffect(() => {
+		isOpen.current = open;
+	}, [open]);
+	var isStream = false;
+
 	const addChatHistory = (
 		message: string,
 		isSend: boolean,
@@ -48,32 +54,29 @@ export default function ChatModal({
 		});
 	};
 	
-	function checkOpenFunction(){
+
-		return open
+	function updateLastMessage(str: string) {
+		console.log(str);
+		setChatHistory((old) => {
+			let newChat = [...old];
+			newChat[newChat.length - 1].message =
+				newChat[newChat.length - 1].message + str;
+			return newChat;
+		});
 	}
 
-	function connectWS() {
+	function handleOnClose(event: CloseEvent) {
-		try {
+		if (isOpen.current) {
-			const urlWs =
-				process.env.NODE_ENV === "development"
-					? `ws://localhost:7860/chat/${flow.id}`
-					: `${window.location.protocol === "https:" ? "wss" : "ws"}://${
-							window.location.host
-					  }/chat/${flow.id}`;
-
-			const newWs = new WebSocket(urlWs);
-			newWs.onopen = () => {
-				console.log("WebSocket connection established!");
-			};
-			console.log(flow.id)
-			newWs.onmessage = (event) => {
 			setLockChat(false);
-				const data = JSON.parse(event.data);
+			setTimeout(() => {
-				console.log("Received data:", data);
+				connectWS()
-				//get chat history
+			}, 1000);
-				if (Array.isArray(data)) {
+		}
-					console.log(data);
+	}
 
+	function handleWsMessage(data: any) {
+		if (Array.isArray(data)) {
+			//set chat history
 			setChatHistory((_) => {
 				let newChatHistory: ChatMessageType[] = [];
 				data.forEach(
@@ -105,36 +108,68 @@ export default function ChatModal({
 				return newChatHistory;
 			});
 		}
-				if (data.type === "end") {
+		if (data.type === "start") {
-					if (data.files) {
+			console.log("start");
-						addChatHistory(
+			addChatHistory("", false);
-							data.message,
+			isStream = true;
-							false,
-							data.intermediate_steps,
-							data.files
-						);
-					} else {
-						addChatHistory(data.message, false, data.intermediate_steps);
 		}
+		if (data.type === "end") {
+			setLockChat(false);
+			isStream = false;
+			// if (data.files) {
+			// 	addChatHistory(
+			// 		data.message,
+			// 		false,
+			// 		data.intermediate_steps,
+			// 		data.files
+			// 	);
+			// } else {
+			// 	addChatHistory(data.message, false, data.intermediate_steps);
+			// }
 		}
 		if (data.type === "file") {
 			console.log(data);
 		}
-			};
+		if (data.type === "stream" && isStream) {
-			newWs.onclose = (_) => {
+			updateLastMessage(data.message);
-				if(checkOpenFunction){
-					setLockChat(false);
-					setTimeout(() => {
-						connectWS();
-					}, 1000);
 		}
+	}
+
+	function connectWS() {
+		try {
+			const urlWs =
+				process.env.NODE_ENV === "development"
+					? `ws://localhost:7860/chat/${flow.id}`
+					: `${window.location.protocol === "https:" ? "wss" : "ws"}://${
+							window.location.host
+					  }/chat/${flow.id}`;
+
+			const newWs = new WebSocket(urlWs);
+			newWs.onopen = () => {
+				console.log("WebSocket connection established!");
+			};
+			console.log(flow.id);
+			newWs.onmessage = (event) => {
+				const data = JSON.parse(event.data);
+				console.log("Received data:", data);
+				handleWsMessage(data);
+				//get chat history
+			};
+			newWs.onclose = (event) => {
+				handleOnClose(event);
 			};
 			newWs.onerror = (ev) => {
 				console.log(ev, "error");
+				setErrorData({
+					title: "There was an error on web connection, please: ",
+					list: [
+						"Refresh the page",
+						"Use a new flow tab",
+						"Check if the backend is up",
+					],
+				});
 			};
-			setWs(newWs);
+			ws.current=newWs
-
-			return newWs;
 		} catch {
 			setErrorData({
 				title: "There was an error on web connection, please: ",
@@ -148,25 +183,21 @@ export default function ChatModal({
 	}
 
 	useEffect(() => {
-		let newWs = connectWS();
+		connectWS();
 		return () => {
-			newWs.close();
+			console.log("unmount")
+			console.log(ws)
+			if (ws) {
+				ws.current.close();
+			}
 		};
 	}, []);
 
-	useEffect(() => {
-		if (ws && (ws.readyState === ws.CLOSED || ws.readyState === ws.CLOSING)) {
-			let newWs = connectWS();
-			return () => {
-				newWs.close();
-			};
-		}
-	}, [lockChat]);
 
 	async function sendAll(data: sendAllProps) {
 		try {
 			if (ws) {
-				ws.send(JSON.stringify(data));
+				ws.current.send(JSON.stringify(data));
 			}
 		} catch (error) {
 			setErrorData({
@@ -182,12 +213,6 @@ export default function ChatModal({
 		if (ref.current) ref.current.scrollIntoView({ behavior: "smooth" });
 	}, [chatHistory]);
 
-	useEffect(() => {
-		if (ws && ws.readyState === ws.CLOSED) {
-			setLockChat(false);
-		}
-	}, [lockChat]);
-
 	function validateNode(n: NodeType): Array<string> {
 		if (!n.data?.node?.template || !Object.keys(n.data.node.template)) {
 			setNoticeData({
@@ -244,7 +269,6 @@ export default function ChatModal({
 				let message = chatValue;
 				setChatValue("");
 				addChatHistory(message, true);
-
 				sendAll({
 					...reactFlowInstance.toObject(),
 					message,
@@ -267,7 +291,7 @@ export default function ChatModal({
 	}
 	function clearChat() {
 		setChatHistory([]);
-		ws.send(JSON.stringify({ clear_history: true }));
+		ws.current.send(JSON.stringify({ clear_history: true }));
 	}
 
 	function setModalOpen(x: boolean) {