From 882c35ef6b936ff250affe29c880ae2f20dbc3ec Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com>
Date: Mon, 2 Sep 2024 10:36:51 -0300
Subject: [PATCH] fix: stop building with abort controller (#3634)

* Added buildController and stopBuilding controllers

* Used stop building in header Stop button

* Set the build controller in API performstreamingrequest

* added setbuildcontroller to types
---
 .../headerComponent/components/menuBar/index.tsx           | 3 ++-
 src/frontend/src/controllers/API/api.tsx                   | 1 +
 src/frontend/src/stores/flowStore.ts                       | 7 +++++++
 src/frontend/src/types/zustand/flow/index.ts               | 3 +++
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
index 48860766ce6..fe5d7ee358a 100644
--- a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
+++ b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
@@ -50,6 +50,7 @@ export const MenuBar = ({}: {}): JSX.Element => {
   const updatedAt = currentSavedFlow?.updated_at;
   const onFlowPage = useFlowStore((state) => state.onFlowPage);
   const setCurrentFlow = useFlowsManagerStore((state) => state.setCurrentFlow);
+  const stopBuilding = useFlowStore((state) => state.stopBuilding);
 
   const changesNotSaved =
     customStringify(currentFlow) !== customStringify(currentSavedFlow) &&
@@ -310,7 +311,7 @@ export const MenuBar = ({}: {}): JSX.Element => {
               disabled={!isBuilding}
               onClick={(_) => {
                 if (isBuilding) {
-                  window.stop();
+                  stopBuilding();
                 }
               }}
               className={
diff --git a/src/frontend/src/controllers/API/api.tsx b/src/frontend/src/controllers/API/api.tsx
index 43059e4939c..e58e36bf500 100644
--- a/src/frontend/src/controllers/API/api.tsx
+++ b/src/frontend/src/controllers/API/api.tsx
@@ -228,6 +228,7 @@ async function performStreamingRequest({
     headers["Authorization"] = `Bearer ${accessToken}`;
   }
   const controller = new AbortController();
+  useFlowStore.getState().setBuildController(controller);
   const params = {
     method: method,
     headers: headers,
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 00c77eb294d..ef9a637a498 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -79,6 +79,9 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
   nodes: [],
   edges: [],
   isBuilding: false,
+  stopBuilding: () => {
+    get().buildController.abort();
+  },
   isPending: true,
   setHasIO: (hasIO) => {
     set({ hasIO });
@@ -766,6 +769,10 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       },
     });
   },
+  buildController: new AbortController(),
+  setBuildController: (controller) => {
+    set({ buildController: controller });
+  },
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index c27e056e24b..ac4705d8ebf 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -180,4 +180,7 @@ export type FlowStoreType = {
     edges?: Edge[];
     viewport?: Viewport;
   }) => void;
+  stopBuilding: () => void;
+  buildController: AbortController;
+  setBuildController: (controller: AbortController) => void;
 };