diff --git a/src/frontend/src/components/codeTabsComponent/index.tsx b/src/frontend/src/components/codeTabsComponent/index.tsx
index 0229434f4..b1bc7ef34 100644
--- a/src/frontend/src/components/codeTabsComponent/index.tsx
+++ b/src/frontend/src/components/codeTabsComponent/index.tsx
@@ -35,7 +35,6 @@ import {
   convertObjToArray,
   convertValuesToNumbers,
   hasDuplicateKeys,
-  unselectAllNodes,
 } from "../../utils/reactflowUtils";
 import { classNames } from "../../utils/utils";
 import DictComponent from "../dictComponent";
@@ -54,6 +53,7 @@ export default function CodeTabsComponent({
   const [data, setData] = useState(flow ? flow["data"]!["nodes"] : null);
   const [openAccordion, setOpenAccordion] = useState<string[]>([]);
   const dark = useDarkStore((state) => state.dark);
+  const unselectAll = useFlowStore((state) => state.unselectAll);
 
   const setNodes = useFlowStore((state) => state.setNodes);
 
@@ -67,12 +67,7 @@ export default function CodeTabsComponent({
 
   useEffect(() => {
     if (tweaks && data) {
-      unselectAllNodes({
-        data,
-        updateNodes: (nodes) => {
-          setNodes(nodes);
-        },
-      });
+      unselectAll();
     }
   }, []);
 
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 25be5d192..85c7ec60f 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -296,6 +296,17 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         get().reactFlowInstance?.getViewport() ?? { x: 0, y: 0, zoom: 1 }
       );
   },
+  unselectAll: () => {
+    let newNodes = cloneDeep(get().nodes);
+    newNodes.forEach((node) => {
+      node.selected = false;
+      let newEdges = cleanEdges(newNodes, get().edges);
+      set({
+        nodes: newNodes,
+        edges: newEdges,
+      });
+    });
+  },
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 273a1958f..1555b82c2 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -54,4 +54,5 @@ export type FlowStoreType = {
   setFilterEdge: (newState) => void;
   getFilterEdge: any[];
   onConnect: (connection: Connection) => void;
+  unselectAll: () => void;
 };