diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 9bdd5c9a0..73eea602c 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -444,25 +444,34 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     function handleBuildUpdate(
       vertexBuildData: VertexBuildTypeAPI,
       status: BuildStatus,
-      buildId: string
+      runId: string
     ) {
       if (vertexBuildData && vertexBuildData.inactive_vertices) {
         get().removeFromVerticesBuild(vertexBuildData.inactive_vertices);
       }
-      get().verticesBuild &&
+      if (vertexBuildData.next_vertices_ids) {
+        // next_vertices_ids is a list of vertices that are going to be built next
+        // verticesLayers is a list of list of vertices ids, where each list is a layer of vertices
+        // we want to add a new layer (next_vertices_ids) to the list of layers (verticesLayers)
+        // and the values of next_vertices_ids to the list of vertices ids (verticesIds)
+        const newLayers = [
+          ...get().verticesBuild!.verticesLayers,
+          vertexBuildData.next_vertices_ids,
+        ];
+        const newIds = [
+          ...get().verticesBuild!.verticesIds,
+          ...vertexBuildData.next_vertices_ids,
+        ];
         get().updateVerticesBuild({
-          verticesIds: [
-            ...get().verticesBuild!.verticesIds,
-            vertexBuildData.id,
-          ],
-          verticesLayers: [
-            ...get().verticesBuild!.verticesLayers,
-            vertexBuildData.next_vertices_ids,
-          ],
-          runId: vertexBuildData.run_id,
+          verticesIds: newIds,
+          verticesLayers: newLayers,
+          runId: runId,
         });
+        get().updateBuildStatus(newIds, BuildStatus.TO_BUILD);
+      }
+
       get().addDataToFlowPool(
-        { ...vertexBuildData, buildId },
+        { ...vertexBuildData, buildId: runId },
         vertexBuildData.id
       );
 
@@ -514,6 +523,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       runId: string;
     } | null
   ) => {
+    console.log("updateVerticesBuild", vertices);
     set({ verticesBuild: vertices });
   },
   verticesBuild: null,
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 06f30fb2c..e10c01835 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -134,7 +134,7 @@ export type Component = {
 };
 
 export type VerticesOrderTypeAPI = {
-  ids: Array<Array<string>>;
+  ids: Array<string>;
   run_id: string;
 };
 
diff --git a/src/frontend/src/utils/buildUtils.ts b/src/frontend/src/utils/buildUtils.ts
index ce3a542ab..8bfced23d 100644
--- a/src/frontend/src/utils/buildUtils.ts
+++ b/src/frontend/src/utils/buildUtils.ts
@@ -31,6 +31,8 @@ function getInactiveVertexData(vertexId: string): VertexBuildTypeAPI {
     id: vertexId,
     data: inactiveData,
     params: "Inactive",
+    run_id: "",
+    next_vertices_ids: [],
     inactive_vertices: null,
     valid: false,
     timestamp: new Date().toISOString(),
@@ -61,27 +63,26 @@ export async function updateVerticesOrder(
       useFlowStore.getState().setIsBuilding(false);
       throw new Error("Invalid nodes");
     }
-    let verticesOrder: Array<Array<string>> = orderResponse.data.ids;
-    let verticesLayers: Array<Array<string>> = [];
+    let verticesLayers: Array<Array<string>> = [orderResponse.data.ids];
     const runId = orderResponse.data.run_id;
-    if (nodeId) {
-      for (let i = 0; i < verticesOrder.length; i += 1) {
-        const innerArray = verticesOrder[i];
-        const idIndex = innerArray.indexOf(nodeId);
-        if (idIndex !== -1) {
-          // If there's a nodeId, we want to run just that component and not the entire layer
-          // because a layer contains dependencies for the next layer
-          // and we are stopping at the layer that contains the nodeId
-          verticesLayers.push([innerArray[idIndex]]);
-          break; // Stop searching after finding the first occurrence
-        }
-        // If the targetId is not found, include the entire inner array
-        verticesLayers.push(innerArray);
-      }
-    } else {
-      verticesLayers = verticesOrder;
-    }
-    const verticesIds = verticesOrder.flat();
+    // if (nodeId) {
+    //   for (let i = 0; i < verticesOrder.length; i += 1) {
+    //     const innerArray = verticesOrder[i];
+    //     const idIndex = innerArray.indexOf(nodeId);
+    //     if (idIndex !== -1) {
+    //       // If there's a nodeId, we want to run just that component and not the entire layer
+    //       // because a layer contains dependencies for the next layer
+    //       // and we are stopping at the layer that contains the nodeId
+    //       verticesLayers.push([innerArray[idIndex]]);
+    //       break; // Stop searching after finding the first occurrence
+    //     }
+    //     // If the targetId is not found, include the entire inner array
+    //     verticesLayers.push(innerArray);
+    //   }
+    // } else {
+    //   verticesLayers = verticesOrder;
+    // }
+    const verticesIds = orderResponse.data.ids;
     useFlowStore.getState().updateVerticesBuild({
       verticesLayers,
       verticesIds,
@@ -128,14 +129,36 @@ export async function buildVertices({
   let currentLayerIndex = 0; // Start with the first layer
   // Set each vertex state to building
   const buildResults: Array<boolean> = [];
-  console.log(verticesLayers);
-  while (currentLayerIndex < verticesLayers.length) {
-    const currentLayer = verticesLayers[currentLayerIndex];
+
+  // Build each layer
+  while (
+    currentLayerIndex <
+    (useFlowStore.getState().verticesBuild?.verticesLayers! || []).length
+  ) {
+    // Get the current layer
+    const currentLayer =
+      useFlowStore.getState().verticesBuild?.verticesLayers![currentLayerIndex];
+    // If there are no more layers, we are done
+    if (!currentLayer) {
+      if (onBuildComplete) {
+        const allNodesValid = buildResults.every((result) => result);
+        onBuildComplete(allNodesValid);
+        useFlowStore.getState().setIsBuilding(false);
+      }
+      return;
+    }
+    // If there is a callback for the start of the build, call it
     if (onBuildStart) onBuildStart(currentLayer);
+    // Build each vertex in the current layer
     await Promise.all(
       currentLayer.map(async (vertexId) => {
         // Check if id is in the list of inactive nodes
-        if (!verticesIds.includes(vertexId) && onBuildUpdate) {
+        if (
+          !useFlowStore
+            .getState()
+            .verticesBuild?.verticesIds.includes(vertexId) &&
+          onBuildUpdate
+        ) {
           // If it is, skip building and set the state to inactive
           onBuildUpdate(
             getInactiveVertexData(vertexId),
@@ -145,6 +168,7 @@ export async function buildVertices({
           buildResults.push(false);
           return;
         }
+        // Build the vertex
         await buildVertex({
           flowId,
           id: vertexId,
@@ -164,6 +188,9 @@ export async function buildVertices({
         }
       })
     );
+    // Once the current layer is built, move to the next layer
+    currentLayerIndex += 1;
+    console.log(useFlowStore.getState().verticesBuild?.verticesLayers);
 
     if (stop) {
       break;