From e1e9d7baef3c5350a6474212d9f2fc8616d056a2 Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Tue, 11 Feb 2025 13:39:56 -0300
Subject: [PATCH] refactor: migrate chat lock state to useFlowStore (#6166)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 🔧 (frontend): remove unused lockChat and setLockChat props from ChatViewWrapper and ChatView components
♻️ (frontend): refactor ChatMessage component to use useFlowStore for lockChat state management instead of passing it as a prop

* 🔧 (chat-message.tsx): Remove unused variables setLockChat and lockChat
♻️ (chat-message.tsx): Refactor code to use isBuilding state instead of setLockChat and lockChat variables
♻️ (new-modal.tsx): Refactor code to remove setLockChat function and references
♻️ (flowStore.ts): Refactor code to remove setLockChat function and lockChat variable
♻️ (components/index.ts): Refactor code to remove setLockChat function and lockChat variable
♻️ (flow/index.ts): Refactor code to remove setLockChat function and lockChat variable
♻️ (buildUtils.ts): Refactor code to remove setLockChat function and references

* ♻️ (chat-view.tsx): refactor variable name 'lockChat' to 'isBuilding' for better clarity and semantics in the code.

* 🔧 (chat-view.tsx, chat-input.tsx, button-send-wrapper.tsx, text-area-wrapper.tsx, upload-file-button.tsx, use-focus-unlock.ts, use-upload.ts, chat-message.tsx, index.ts): Remove lockChat variable and replace it with isBuilding variable to improve code readability and consistency.

* ♻️ (button-send-wrapper.tsx): remove unnecessary disabled prop from Button component to improve code readability and maintainability
---
 .../IOModal/components/chat-view-wrapper.tsx  |  4 ----
 .../IOModal/components/chatView/chat-view.tsx | 11 ++++------
 .../chatView/chatInput/chat-input.tsx         | 18 +++++++---------
 .../components/button-send-wrapper.tsx        |  8 ++-----
 .../components/text-area-wrapper.tsx          |  8 +++----
 .../components/upload-file-button.tsx         |  8 +++----
 .../chatInput/hooks/use-focus-unlock.ts       |  6 +++---
 .../chatView/chatInput/hooks/use-upload.ts    |  7 ++++---
 .../chatView/chatMessage/chat-message.tsx     | 19 ++++++-----------
 src/frontend/src/modals/IOModal/new-modal.tsx | 21 ++-----------------
 .../modals/IOModal/types/chat-view-wrapper.ts |  2 --
 src/frontend/src/stores/flowStore.ts          | 15 -------------
 src/frontend/src/types/components/index.ts    |  9 --------
 src/frontend/src/types/zustand/flow/index.ts  |  4 ----
 src/frontend/src/utils/buildUtils.ts          |  8 -------
 15 files changed, 37 insertions(+), 111 deletions(-)

diff --git a/src/frontend/src/modals/IOModal/components/chat-view-wrapper.tsx b/src/frontend/src/modals/IOModal/components/chat-view-wrapper.tsx
index 32f01029c..6bcf576fe 100644
--- a/src/frontend/src/modals/IOModal/components/chat-view-wrapper.tsx
+++ b/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
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chat-view.tsx b/src/frontend/src/modals/IOModal/components/chatView/chat-view.tsx
index a278d7587..34245259a 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chat-view.tsx
+++ b/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");
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx
index fbb32f10e..50a04151b 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/chat-input.tsx
+++ b/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}
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/button-send-wrapper.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/button-send-wrapper.tsx
index ddc3ceda2..b5bbdc020 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/button-send-wrapper.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/button-send-wrapper.tsx
@@ -16,7 +16,6 @@ const BUTTON_STATES = {
 
 type ButtonSendWrapperProps = {
   send: () => void;
-  lockChat: boolean;
   noInput: boolean;
   chatValue: string;
   files: FilePreviewType[];
@@ -24,7 +23,6 @@ type ButtonSendWrapperProps = {
 
 const ButtonSendWrapper = ({
   send,
-  lockChat,
   noInput,
   chatValue,
   files,
@@ -32,10 +30,9 @@ const ButtonSendWrapper = ({
   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"}
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/text-area-wrapper.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/text-area-wrapper.tsx
index 93fc6cbf0..9060aa0f4 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/text-area-wrapper.tsx
+++ b/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`,
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/upload-file-button.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/upload-file-button.tsx
index e0a76cb8a..7fcbbe379 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/components/upload-file-button.tsx
+++ b/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"
           }`}
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-focus-unlock.ts b/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-focus-unlock.ts
index 6e951d80f..09a9e1531 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-focus-unlock.ts
+++ b/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;
 };
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-upload.ts b/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-upload.ts
index b6b31647f..e0cc545d5 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatInput/hooks/use-upload.ts
+++ b/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;
 };
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx
index 8ad21cdbb..981313140 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx
+++ b/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,11 +48,13 @@ 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,
@@ -61,10 +62,6 @@ export default function ChatMessage({
   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"
diff --git a/src/frontend/src/modals/IOModal/new-modal.tsx b/src/frontend/src/modals/IOModal/new-modal.tsx
index 271e02de5..39170d886 100644
--- a/src/frontend/src/modals/IOModal/new-modal.tsx
+++ b/src/frontend/src/modals/IOModal/new-modal.tsx
@@ -113,8 +113,7 @@ export default function IOModal({
 
   const buildFlow = useFlowStore((state) => state.buildFlow);
   const setIsBuilding = useFlowStore((state) => state.setIsBuilding);
-  const lockChat = useFlowStore((state) => state.lockChat);
-  const setLockChat = useFlowStore((state) => state.setLockChat);
+
   const isBuilding = useFlowStore((state) => state.isBuilding);
   const messages = useMessagesStore((state) => state.messages);
   const [sessions, setSessions] = useState<string[]>(
@@ -147,8 +146,6 @@ export default function IOModal({
       files?: string[];
     }): Promise<void> => {
       if (isBuilding) return;
-      setIsBuilding(true);
-      setLockChat(true);
       setChatValue("");
       for (let i = 0; i < repeat; i++) {
         await buildFlow({
@@ -157,24 +154,12 @@ export default function IOModal({
           files: files,
           silent: true,
           session: sessionId,
-          setLockChat,
         }).catch((err) => {
           console.error(err);
-          setLockChat(false);
         });
       }
-      // refetch();
-      setLockChat(false);
     },
-    [
-      isBuilding,
-      setIsBuilding,
-      setLockChat,
-      chatValue,
-      chatInput?.id,
-      sessionId,
-      buildFlow,
-    ],
+    [isBuilding, setIsBuilding, chatValue, chatInput?.id, sessionId, buildFlow],
   );
 
   useEffect(() => {
@@ -331,8 +316,6 @@ export default function IOModal({
                 messagesFetched={messagesFetched}
                 sessionId={sessionId}
                 sendMessage={sendMessage}
-                lockChat={lockChat}
-                setLockChat={setLockChat}
                 canvasOpen={canvasOpen}
                 setOpen={setOpen}
               />
diff --git a/src/frontend/src/modals/IOModal/types/chat-view-wrapper.ts b/src/frontend/src/modals/IOModal/types/chat-view-wrapper.ts
index a269442ea..31b101a4a 100644
--- a/src/frontend/src/modals/IOModal/types/chat-view-wrapper.ts
+++ b/src/frontend/src/modals/IOModal/types/chat-view-wrapper.ts
@@ -14,8 +14,6 @@ export type ChatViewWrapperProps = {
   messagesFetched: boolean;
   sessionId: string;
   sendMessage: (options: { repeat: number; files?: string[] }) => Promise<void>;
-  lockChat: boolean;
-  setLockChat: (locked: boolean) => void;
   canvasOpen: boolean | undefined;
   setOpen: (open: boolean) => void;
 };
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index a246ec2c6..b869a348a 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -49,7 +49,6 @@ import useAlertStore from "./alertStore";
 import { useDarkStore } from "./darkStore";
 import useFlowsManagerStore from "./flowsManagerStore";
 import { useGlobalVariablesStore } from "./globalVariablesStore/globalVariables";
-import { useMessagesStore } from "./messagesStore";
 import { useTypesStore } from "./typesStore";
 
 // this is our useStore hook that we can use in our components to get parts of the store and call actions
@@ -101,11 +100,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     set({ componentsToUpdate: outdatedNodes });
   },
   onFlowPage: false,
-  lockChat: false,
-  setLockChat: (lockChat) => {
-    useMessagesStore.setState({ displayLoadingMessage: lockChat });
-    set({ lockChat });
-  },
   setOnFlowPage: (FlowPage) => set({ onFlowPage: FlowPage }),
   flowState: undefined,
   flowBuildStatus: {},
@@ -120,7 +114,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     );
     set({ isBuilding: false });
     get().revertBuiltStatusFromBuilding();
-    get().setLockChat(false);
     useAlertStore.getState().setErrorData({
       title: "Build stopped",
     });
@@ -595,7 +588,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     input_value,
     files,
     silent,
-    setLockChat,
     session,
   }: {
     startNodeId?: string;
@@ -603,11 +595,9 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     input_value?: string;
     files?: string[];
     silent?: boolean;
-    setLockChat?: (lock: boolean) => void;
     session?: string;
   }) => {
     get().setIsBuilding(true);
-    get().setLockChat(true);
     const currentFlow = useFlowsManagerStore.getState().currentFlow;
     const setSuccessData = useAlertStore.getState().setSuccessData;
     const setErrorData = useAlertStore.getState().setErrorData;
@@ -627,7 +617,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     }
     if (error) {
       get().setIsBuilding(false);
-      get().setLockChat(false);
       throw new Error("Invalid components");
     }
 
@@ -734,7 +723,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       flowId: currentFlow!.id,
       startNodeId,
       stopNodeId,
-      setLockChat,
       onGetOrderSuccess: () => {
         if (!silent) {
           setNoticeData({ title: "Running components" });
@@ -759,7 +747,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
           false,
         );
         get().setIsBuilding(false);
-        get().setLockChat(false);
       },
       onBuildUpdate: handleBuildUpdate,
       onBuildError: (title: string, list: string[], elementList) => {
@@ -779,7 +766,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         );
         setErrorData({ list, title });
         get().setIsBuilding(false);
-        get().setLockChat(false);
         get().buildController.abort();
       },
       onBuildStart: (elementList) => {
@@ -807,7 +793,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       logBuilds: get().onFlowPage,
     });
     get().setIsBuilding(false);
-    get().setLockChat(false);
     get().revertBuiltStatusFromBuilding();
   },
   getFlow: () => {
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 01950fb90..2e6d5e0a8 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -505,7 +505,6 @@ export type ChatInputType = {
   inputRef: {
     current: any;
   };
-  lockChat: boolean;
   noInput: boolean;
   sendMessage: ({
     repeat,
@@ -593,9 +592,7 @@ export type codeAreaModalPropsType = {
 
 export type chatMessagePropsType = {
   chat: ChatMessageType;
-  lockChat: boolean;
   lastMessage: boolean;
-  setLockChat: (lock: boolean) => void;
   updateChat: (
     chat: ChatMessageType,
     message: string,
@@ -772,8 +769,6 @@ export type chatViewProps = {
     repeat: number;
     files?: string[];
   }) => void;
-  lockChat: boolean;
-  setLockChat: (lock: boolean) => void;
   visibleSession?: string;
   focusChat?: string;
   closeChat?: () => void;
@@ -794,15 +789,11 @@ export type toolbarSelectItemProps = {
 };
 
 export type clearChatPropsType = {
-  lockChat: boolean;
-  setLockChat: (lock: boolean) => void;
   setChatHistory: (chatHistory: ChatMessageType) => void;
   method: string;
 };
 
 export type handleSelectPropsType = {
   event: string;
-  lockChat: boolean;
-  setLockChat: (lock: boolean) => void;
   setChatHistory: (chatHistory: ChatMessageType) => void;
 };
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 8cdaffa07..119c99e08 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -145,10 +145,8 @@ export type FlowStoreType = {
     input_value,
     files,
     silent,
-    setLockChat,
     session,
   }: {
-    setLockChat?: (lock: boolean) => void;
     startNodeId?: string;
     stopNodeId?: string;
     input_value?: string;
@@ -184,8 +182,6 @@ export type FlowStoreType = {
     buildId?: string,
   ) => void;
   getNodePosition: (nodeId: string) => { x: number; y: number };
-  setLockChat: (lock: boolean) => void;
-  lockChat: boolean;
   updateFreezeStatus: (nodeIds: string[], freeze: boolean) => void;
   currentFlow: FlowType | undefined;
   setCurrentFlow: (flow: FlowType | undefined) => void;
diff --git a/src/frontend/src/utils/buildUtils.ts b/src/frontend/src/utils/buildUtils.ts
index 3553f2f52..febb926c2 100644
--- a/src/frontend/src/utils/buildUtils.ts
+++ b/src/frontend/src/utils/buildUtils.ts
@@ -14,7 +14,6 @@ import { VertexLayerElementType } from "../types/zustand/flow";
 import { isStringArray, tryParseJson } from "./utils";
 
 type BuildVerticesParams = {
-  setLockChat?: (lock: boolean) => void;
   flowId: string; // Assuming FlowType is the type for your flow
   input_value?: any; // Replace any with the actual type if it's not any
   files?: string[];
@@ -66,7 +65,6 @@ function getInactiveVertexData(vertexId: string): VertexBuildTypeAPI {
 
 export async function updateVerticesOrder(
   flowId: string,
-  setLockChat?: (lock: boolean) => void,
   startNodeId?: string | null,
   stopNodeId?: string | null,
   nodes?: Node[],
@@ -94,7 +92,6 @@ export async function updateVerticesOrder(
         list: [error.response?.data?.detail ?? "Unknown Error"],
       });
       useFlowStore.getState().setIsBuilding(false);
-      setLockChat && setLockChat(false);
       throw new Error("Invalid components");
     }
     // orderResponse.data.ids,
@@ -154,7 +151,6 @@ export async function buildFlowVertices({
   nodes,
   edges,
   logBuilds,
-  setLockChat,
   session,
 }: BuildVerticesParams) {
   const inputs = {};
@@ -226,7 +222,6 @@ export async function buildFlowVertices({
             return true;
           } catch (e) {
             useFlowStore.getState().setIsBuilding(false);
-            setLockChat && setLockChat(false);
             return false;
           }
         }
@@ -381,7 +376,6 @@ export async function buildVertices({
   onValidateNodes,
   nodes,
   edges,
-  setLockChat,
 }: BuildVerticesParams) {
   // if startNodeId and stopNodeId are provided
   // something is wrong
@@ -390,7 +384,6 @@ export async function buildVertices({
   }
   let verticesOrderResponse = await updateVerticesOrder(
     flowId,
-    setLockChat,
     startNodeId,
     stopNodeId,
     nodes,
@@ -401,7 +394,6 @@ export async function buildVertices({
       onValidateNodes(verticesOrderResponse.verticesToRun);
     } catch (e) {
       useFlowStore.getState().setIsBuilding(false);
-      setLockChat && setLockChat(false);
       return;
     }
   }