diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index c62abf6b9..0e4812c90 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -36,6 +36,7 @@ export default function GenericNode({
   const flowPool = useFlowStore((state) => state.flowPool);
   const buildFlow = useFlowStore((state) => state.buildFlow);
   const setNode = useFlowStore((state) => state.setNode);
+  const getBuildStatus = useFlowStore((state) => state.getBuildStatus);
   const name = nodeIconsLucide[data.type] ? data.type : types[data.type];
   const [inputName, setInputName] = useState(false);
   const [nodeName, setNodeName] = useState(data.node!.display_name);
@@ -262,7 +263,7 @@ export default function GenericNode({
         className={getNodeBorderClassName(
           selected,
           showNode,
-          data?.build_status,
+          getBuildStatus(data.id),
           validationStatus
         )}
       >
@@ -473,7 +474,10 @@ export default function GenericNode({
                 variant="outline"
                 className={"h-9 px-1.5"}
                 onClick={() => {
-                  if (data?.build_status === BuildStatus.BUILDING || isBuilding)
+                  if (
+                    getBuildStatus(data.id) === BuildStatus.BUILDING ||
+                    isBuilding
+                  )
                     return;
                   buildFlow(data.id);
                 }}
@@ -481,7 +485,7 @@ export default function GenericNode({
                 <div>
                   <Tooltip
                     title={
-                      data?.build_status === BuildStatus.BUILDING ? (
+                      getBuildStatus(data.id) === BuildStatus.BUILDING ? (
                         <span>Building...</span>
                       ) : !validationStatus ? (
                         <span className="flex">
@@ -507,7 +511,7 @@ export default function GenericNode({
                   >
                     <div className="generic-node-status-position flex items-center justify-center">
                       {renderIconPlayOrPauseComponents(
-                        data?.build_status,
+                        getBuildStatus(data.id),
                         validationStatus
                       )}
                     </div>
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 549b0a90c..dac285ddf 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -425,25 +425,20 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       viewport: get().reactFlowInstance?.getViewport()!,
     };
   },
-  updateBuildStatus: (nodeIdList: string[], status: BuildStatus) => {
-    nodeIdList.forEach((id) => {
-      const nodeToUpdate = get().nodes.find((node) => node.id === id);
-      if (nodeToUpdate) {
-        nodeToUpdate.data.build_status = status;
-        get().setNodes(get().nodes);
-      }
-    });
-  },
   updateVerticesBuild: (vertices: string[]) => {
     set({ verticesBuild: vertices });
   },
   verticesBuild: [],
   revertAllVerticesToBuild: () => {
     // set all vertices to TO_BUILD
-    const verticesIds = get()
-      .nodes.filter((node) => node.data.build_status === BuildStatus.BUILDING)
-      .map((node) => node.id);
+    let verticesIds: string[] = [];
+    Object.keys(get().buildStatus).forEach((id) => {
+      if (get().buildStatus[id] !== BuildStatus.TO_BUILD) {
+        verticesIds.push(id);
+      }
+    });
     get().updateBuildStatus(verticesIds, BuildStatus.TO_BUILD);
+    // additional code logic here
   },
   removeFromVerticesBuild: (vertices: string[]) => {
     set({
@@ -452,6 +447,25 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       ),
     });
   },
+  updateBuildStatus: (nodeIdList: string[], status: BuildStatus) => {
+    let newBuildStatus = cloneDeep(get().buildStatus);
+    // check if nodeIdList is an array
+    if (!Array.isArray(nodeIdList)) {
+      nodeIdList = [nodeIdList];
+    }
+    nodeIdList.forEach((id) => {
+      newBuildStatus[id] = status;
+    });
+    set({ buildStatus: newBuildStatus });
+  },
+  buildStatus: {},
+  getBuildStatus: (nodeId: string) => {
+    // if the node is not in the buildStatus object, set it to TO_BUILD
+    if (!get().buildStatus[nodeId]) {
+      get().buildStatus[nodeId] = BuildStatus.TO_BUILD;
+    }
+    return get().buildStatus[nodeId];
+  },
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index ae34aa651..59b8a5ca5 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -1,5 +1,4 @@
 import { Edge, Node, Viewport } from "reactflow";
-import { BuildStatus } from "../../constants/enums";
 import { FlowType } from "../flow";
 //kind and class are just representative names to represent the actual structure of the object received by the API
 export type APIDataType = { [key: string]: APIKindType };
@@ -37,7 +36,6 @@ export type APIClassType = {
     | CustomFieldsType
     | boolean
     | undefined;
-  build_status?: BuildStatus;
 };
 
 export type TemplateVariableType = {
diff --git a/src/frontend/src/types/flow/index.ts b/src/frontend/src/types/flow/index.ts
index 127e1a704..8575b331e 100644
--- a/src/frontend/src/types/flow/index.ts
+++ b/src/frontend/src/types/flow/index.ts
@@ -1,5 +1,4 @@
 import { ReactFlowJsonObject, XYPosition } from "reactflow";
-import { BuildStatus } from "../../constants/enums";
 import { APIClassType } from "../api/index";
 
 export type FlowType = {
@@ -29,7 +28,6 @@ export type NodeDataType = {
   node?: APIClassType;
   id: string;
   output_types?: string[];
-  build_status?: BuildStatus;
 };
 // FlowStyleType is the type of the style object that is used to style the
 // Flow card with an emoji and a color.
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 7f214b69d..2fc987de6 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -86,9 +86,11 @@ export type FlowStoreType = {
   unselectAll: () => void;
   buildFlow: (nodeId?: string) => Promise<void>;
   getFlow: () => { nodes: Node[]; edges: Edge[]; viewport: Viewport };
-  updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;
   updateVerticesBuild: (vertices: string[]) => void;
   removeFromVerticesBuild: (vertices: string[]) => void;
   revertAllVerticesToBuild: () => void;
   verticesBuild: string[];
+  updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;
+  buildStatus: { [key: string]: BuildStatus };
+  getBuildStatus: (nodeId: string) => BuildStatus;
 };