refactor: migrate chat lock state to useFlowStore

src/frontend/src/modals/IOModal/components/chat-view-wrapper.tsx

@@ -19,8 +19,6 @@ export const ChatViewWrapper = ({
   messagesFetched,
   sessionId,
   sendMessage,
-  lockChat,
-  setLockChat,
   canvasOpen,
   setOpen,
 }: ChatViewWrapperProps) => {
@@ -93,8 +91,6 @@ export const ChatViewWrapper = ({
           <ChatView
             focusChat={sessionId}
             sendMessage={sendMessage}
-            lockChat={lockChat}
-            setLockChat={setLockChat}
             visibleSession={visibleSession}
             closeChat={
               !canvasOpen

src/frontend/src/modals/IOModal/components/chatView/chat-view.tsx

@@ -30,8 +30,6 @@ const MemoizedChatMessage = memo(ChatMessage, (prevProps, nextProps) => {
 
 export default function ChatView({
   sendMessage,
-  lockChat,
-  setLockChat,
   visibleSession,
   focusChat,
   closeChat,
@@ -51,6 +49,8 @@ export default function ChatView({
     (state) => state.displayLoadingMessage,
   );
 
+  const isBuilding = useFlowStore((state) => state.isBuilding);
+
   const inputTypes = inputs.map((obj) => obj.type);
   const updateFlowPool = useFlowStore((state) => state.updateFlowPool);
   const setChatValueStore = useUtilityStore((state) => state.setChatValueStore);
@@ -99,7 +99,7 @@ export default function ChatView({
       return new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime();
     });
 
-    if (messages.length === 0 && !lockChat && chatInputNode && isTabHidden) {
+    if (messages.length === 0 && !isBuilding && chatInputNode && isTabHidden) {
       setChatValueStore(
         chatInputNode.data.node.template["input_value"].value ?? "",
       );
@@ -164,12 +164,10 @@ export default function ChatView({
     >
       <div ref={messagesRef} className="chat-message-div">
         {chatHistory &&
-          (lockChat || chatHistory?.length > 0 ? (
+          (isBuilding || chatHistory?.length > 0 ? (
             <>
               {chatHistory?.map((chat, index) => (
                 <MemoizedChatMessage
-                  setLockChat={setLockChat}
-                  lockChat={lockChat}
                   chat={chat}
                   lastMessage={chatHistory.length - 1 === index}
                   key={`${chat.id}-${index}`}
@@ -224,7 +222,6 @@ export default function ChatView({
       <div className="m-auto w-full max-w-[768px] md:w-5/6">
         <ChatInput
           noInput={!inputTypes.includes("ChatInput")}
-          lockChat={lockChat}
           sendMessage={({ repeat, files }) => {
             sendMessage({ repeat, files });
             track("Playground Message Sent");

src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx

@@ -26,7 +26,6 @@ import UploadFileButton from "./components/upload-file-button";
 import useAutoResizeTextArea from "./hooks/use-auto-resize-text-area";
 import useFocusOnUnlock from "./hooks/use-focus-unlock";
 export default function ChatInput({
-  lockChat,
   sendMessage,
   inputRef,
   noInput,
@@ -39,10 +38,10 @@ export default function ChatInput({
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const { validateFileSize } = useFileSizeValidator(setErrorData);
   const stopBuilding = useFlowStore((state) => state.stopBuilding);
-
+  const isBuilding = useFlowStore((state) => state.isBuilding);
   const chatValue = useUtilityStore((state) => state.chatValueStore);
 
-  useFocusOnUnlock(lockChat, inputRef);
+  useFocusOnUnlock(isBuilding, inputRef);
   useAutoResizeTextArea(chatValue, inputRef);
 
   const { mutate } = usePostUploadFile();
@@ -134,7 +133,7 @@ export default function ChatInput({
     return () => {
       document.removeEventListener("paste", handleFileChange);
     };
-  }, [handleFileChange, currentFlowId, lockChat]);
+  }, [handleFileChange, currentFlowId, isBuilding]);
 
   const send = () => {
     sendMessage({
@@ -147,7 +146,7 @@ export default function ChatInput({
   const checkSendingOk = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
     return (
       event.key === "Enter" &&
-      !lockChat &&
+      !isBuilding &&
       !event.shiftKey &&
       !event.nativeEvent.isComposing
     );
@@ -168,7 +167,7 @@ export default function ChatInput({
     return (
       <div className="flex h-full w-full flex-col items-center justify-center">
         <div className="flex w-full flex-col items-center justify-center gap-3 rounded-md border border-input bg-muted p-2 py-4">
-          {!lockChat ? (
+          {!isBuilding ? (
             <Button
               data-testid="button-send"
               className="font-semibold"
@@ -214,9 +213,9 @@ export default function ChatInput({
     <div className="flex w-full flex-col-reverse">
       <div className="flex w-full flex-col rounded-md border border-input p-4 hover:border-muted-foreground focus:border-[1.75px] has-[:focus]:border-primary">
         <TextAreaWrapper
+          isBuilding={isBuilding}
           checkSendingOk={checkSendingOk}
           send={send}
-          lockChat={lockChat}
           noInput={noInput}
           chatValue={chatValue}
           CHAT_INPUT_PLACEHOLDER={CHAT_INPUT_PLACEHOLDER}
@@ -239,9 +238,9 @@ export default function ChatInput({
           ))}
         </div>
         <div className="flex w-full items-end justify-between">
-          <div className={lockChat ? "cursor-not-allowed" : ""}>
+          <div className={isBuilding ? "cursor-not-allowed" : ""}>
             <UploadFileButton
-              lockChat={lockChat}
+              isBuilding={isBuilding}
               fileInputRef={fileInputRef}
               handleFileChange={handleFileChange}
               handleButtonClick={handleButtonClick}
@@ -250,7 +249,6 @@ export default function ChatInput({
           <div className="">
             <ButtonSendWrapper
               send={send}
-              lockChat={lockChat}
               noInput={noInput}
               chatValue={chatValue}
               files={files}

src/frontend/src/modals/IOModal/components/chatView/chatInput/components/button-send-wrapper.tsx

@@ -16,26 +16,23 @@ const BUTTON_STATES = {
 
 type ButtonSendWrapperProps = {
   send: () => void;
-  lockChat: boolean;
   noInput: boolean;
   chatValue: string;
   files: FilePreviewType[];
 };
 
 const ButtonSendWrapper = ({
   send,
-  lockChat,
   noInput,
   chatValue,
   files,
 }: ButtonSendWrapperProps) => {
   const stopBuilding = useFlowStore((state) => state.stopBuilding);
 
   const isBuilding = useFlowStore((state) => state.isBuilding);
-  const showStopButton = lockChat || files.some((file) => file.loading);
-  const showPlayButton = !lockChat && noInput;
+  const showStopButton = isBuilding || files.some((file) => file.loading);
   const showSendButton =
-    !(lockChat || files.some((file) => file.loading)) && !noInput;
+    !(isBuilding || files.some((file) => file.loading)) && !noInput;
 
   const getButtonState = () => {
     if (showStopButton) return BUTTON_STATES.SHOW_STOP;
@@ -58,7 +55,6 @@ const ButtonSendWrapper = ({
   return (
     <Button
       className={buttonClasses}
-      disabled={lockChat && !isBuilding}
       onClick={handleClick}
       unstyled
       data-testid={showStopButton ? "button-stop" : "button-send"}

src/frontend/src/modals/IOModal/components/chatView/chatInput/components/text-area-wrapper.tsx

@@ -6,7 +6,7 @@ import { classNames } from "../../../../../../utils/utils";
 const TextAreaWrapper = ({
   checkSendingOk,
   send,
-  lockChat,
+  isBuilding,
   noInput,
   chatValue,
   CHAT_INPUT_PLACEHOLDER,
@@ -36,10 +36,10 @@ const TextAreaWrapper = ({
     "form-input block w-full border-0 custom-scroll focus:border-ring rounded-none shadow-none focus:ring-0 p-0 sm:text-sm !bg-transparent";
 
   useEffect(() => {
-    if (!lockChat && !noInput) {
+    if (!isBuilding && !noInput) {
       inputRef.current?.focus();
     }
-  }, [lockChat, noInput]);
+  }, [isBuilding, noInput]);
 
   return (
     <Textarea
@@ -51,7 +51,7 @@ const TextAreaWrapper = ({
       }}
       rows={1}
       ref={inputRef}
-      disabled={lockChat || noInput}
+      disabled={isBuilding || noInput}
       style={{
         resize: "none",
         bottom: `${inputRef?.current?.scrollHeight}px`,

src/frontend/src/modals/IOModal/components/chatView/chatInput/components/upload-file-button.tsx

@@ -6,7 +6,7 @@ const UploadFileButton = ({
   fileInputRef,
   handleFileChange,
   handleButtonClick,
-  lockChat,
+  isBuilding,
 }) => {
   return (
     <ShadTooltip
@@ -16,16 +16,16 @@ const UploadFileButton = ({
     >
       <div>
         <input
-          disabled={lockChat}
+          disabled={isBuilding}
           type="file"
           ref={fileInputRef}
           style={{ display: "none" }}
           onChange={handleFileChange}
         />
         <Button
-          disabled={lockChat}
+          disabled={isBuilding}
           className={`flex h-[32px] w-[32px] items-center justify-center rounded-md bg-muted font-bold transition-all ${
-            lockChat
+            isBuilding
               ? "cursor-not-allowed"
               : "text-muted-foreground hover:text-primary"
           }`}

src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-focus-unlock.ts

@@ -1,14 +1,14 @@
 import { useEffect } from "react";
 
 const useFocusOnUnlock = (
-  lockChat: boolean,
+  isBuilding: boolean,
   inputRef: React.RefObject<HTMLInputElement>,
 ) => {
   useEffect(() => {
-    if (!lockChat && inputRef.current) {
+    if (!isBuilding && inputRef.current) {
       inputRef.current.focus();
     }
-  }, [lockChat, inputRef]);
+  }, [isBuilding, inputRef]);
 
   return inputRef;
 };

src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-upload.ts

@@ -1,3 +1,4 @@
+import useFlowStore from "@/stores/flowStore";
 import { AxiosResponse } from "axios";
 import { useEffect } from "react";
 import ShortUniqueId from "short-unique-id";
@@ -16,12 +17,12 @@ const useUpload = (
   ) => Promise<AxiosResponse<UploadFileTypeAPI>>,
   currentFlowId: string,
   setFiles: any,
-  lockChat: boolean,
 ) => {
   const setErrorData = useAlertStore((state) => state.setErrorData);
+  const isBuilding = useFlowStore((state) => state.isBuilding);
   useEffect(() => {
     const handlePaste = (event: ClipboardEvent): void => {
-      if (lockChat) {
+      if (isBuilding) {
         return;
       }
       const items = event.clipboardData?.items;
@@ -57,7 +58,7 @@ const useUpload = (
     return () => {
       document.removeEventListener("paste", handlePaste);
     };
-  }, [uploadFile, currentFlowId, lockChat]);
+  }, [uploadFile, currentFlowId, isBuilding]);
 
   return null;
 };

src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx

@@ -6,6 +6,7 @@ import { ENABLE_DATASTAX_LANGFLOW } from "@/customization/feature-flags";
 import useFlowsManagerStore from "@/stores/flowsManagerStore";
 import useFlowStore from "@/stores/flowStore";
 import { useUtilityStore } from "@/stores/utilityStore";
+import { ChatMessageType } from "@/types/chat";
 import Convert from "ansi-to-html";
 import { useEffect, useRef, useState } from "react";
 import Robot from "../../../../../assets/robot.png";
@@ -27,10 +28,8 @@ import { convertFiles } from "./helpers/convert-files";
 
 export default function ChatMessage({
   chat,
-  lockChat,
   lastMessage,
   updateChat,
-  setLockChat,
   closeChat,
 }: chatMessagePropsType): JSX.Element {
   const convert = new Convert({ newline: true });
@@ -49,22 +48,20 @@ export default function ChatMessage({
   const chatMessageRef = useRef(chatMessage);
   const [editMessage, setEditMessage] = useState(false);
   const [showError, setShowError] = useState(false);
+  const isBuilding = useFlowStore((state) => state.isBuilding);
 
   useEffect(() => {
     const chatMessageString = chat.message ? chat.message.toString() : "";
     setChatMessage(chatMessageString);
-  }, [chat]);
+    chatMessageRef.current = chatMessage;
+  }, [chat, isBuilding]);
 
   const playgroundScrollBehaves = useUtilityStore(
     (state) => state.playgroundScrollBehaves,
   );
   const setPlaygroundScrollBehaves = useUtilityStore(
     (state) => state.setPlaygroundScrollBehaves,
   );
-  // Sync ref with state
-  useEffect(() => {
-    chatMessageRef.current = chatMessage;
-  }, [chatMessage]);
 
   // The idea now is that chat.stream_url MAY be a URL if we should stream the output of the chat
   // probably the message is empty when we have a stream_url
@@ -103,21 +100,17 @@ export default function ChatMessage({
 
   useEffect(() => {
     if (streamUrl && !isStreaming) {
-      setLockChat(true);
       streamChunks(streamUrl)
         .then(() => {
-          setLockChat(false);
           if (updateChat) {
             updateChat(chat, chatMessageRef.current);
           }
         })
         .catch((error) => {
           console.error(error);
-          setLockChat(false);
         });
     }
   }, [streamUrl, chatMessage]);
-
   useEffect(() => {
     return () => {
       eventSource.current?.close();
@@ -320,7 +313,7 @@ export default function ChatMessage({
                 contentBlocks={chat.content_blocks}
                 isLoading={
                   chatMessage === "" &&
-                  lockChat &&
+                  isBuilding &&
                   chat.properties?.state === "partial"
                 }
                 state={chat.properties?.state}
@@ -360,7 +353,7 @@ export default function ChatMessage({
                         }
                         className="flex w-full flex-col"
                       >
-                        {chatMessage === "" && lockChat ? (
+                        {chatMessage === "" && isBuilding && lastMessage ? (
                           <IconComponent
                             name="MoreHorizontal"
                             className="h-8 w-8 animate-pulse"