diff --git a/src/frontend/src/contexts/tabsContext.tsx b/src/frontend/src/contexts/tabsContext.tsx
index 283345ead..f1e5dd3d5 100644
--- a/src/frontend/src/contexts/tabsContext.tsx
+++ b/src/frontend/src/contexts/tabsContext.tsx
@@ -1,5 +1,5 @@
 import { AxiosError } from "axios";
-import _ from "lodash";
+import _, { cloneDeep } from "lodash";
 import {
   ReactNode,
   createContext,
@@ -458,6 +458,7 @@ export function TabsProvider({ children }: { children: ReactNode }) {
         id: source,
       });
       sourceHandleObject.id = source;
+
       edge.data.sourceHandle = sourceHandleObject;
       const targetHandleObject: targetHandleType = scapeJSONParse(
         edge.targetHandle!
@@ -482,6 +483,7 @@ export function TabsProvider({ children }: { children: ReactNode }) {
           sourceHandle,
           targetHandle,
           id,
+          data: cloneDeep(edge.data),
           style: { stroke: "#555" },
           className:
             targetHandleObject.type === "Text"
@@ -538,7 +540,8 @@ export function TabsProvider({ children }: { children: ReactNode }) {
     let data = flow?.data ? flow.data : null;
     if (data) {
       processFlowEdges(flow);
-      processFlowNodes(flow);
+      //prevent node update for now
+      // processFlowNodes(flow);
       //add animation to text type edges
       updateEdges(data.edges);
       // updateNodes(data.nodes, data.edges);
diff --git a/src/frontend/src/modals/EditNodeModal/index.tsx b/src/frontend/src/modals/EditNodeModal/index.tsx
index 04042104e..af70d6b0e 100644
--- a/src/frontend/src/modals/EditNodeModal/index.tsx
+++ b/src/frontend/src/modals/EditNodeModal/index.tsx
@@ -7,6 +7,7 @@ import {
   useRef,
   useState,
 } from "react";
+import { useUpdateNodeInternals } from "reactflow";
 import ShadTooltip from "../../components/ShadTooltipComponent";
 import CodeAreaComponent from "../../components/codeAreaComponent";
 import DictComponent from "../../components/dictComponent";
@@ -64,6 +65,7 @@ const EditNodeModal = forwardRef(
     ref
   ) => {
     const [modalOpen, setModalOpen] = useState(open ?? false);
+    const updateNodeInternals = useUpdateNodeInternals();
 
     const myData = useRef(data);
 
@@ -83,11 +85,14 @@ const EditNodeModal = forwardRef(
     const handleOnNewValue = (newValue: any, name) => {
       myData.current.node!.template[name].value = newValue;
       setDataValue(newValue);
+      updateNodeInternals(data.id);
     };
 
     useEffect(() => {
-      myData.current = data; // reset data to what it is on node when opening modal
-      onClose!(modalOpen);
+      if (modalOpen) {
+        myData.current = data; // reset data to what it is on node when opening modal
+        onClose!(modalOpen);
+      }
     }, [modalOpen]);
 
     const [errorDuplicateKey, setErrorDuplicateKey] = useState(false);
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index 20fb826e1..0512859c7 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -310,7 +310,6 @@ export default function Page({
           event.dataTransfer.getData("nodedata")
         );
 
-        // If data type is not "chatInput" or if there are no "chatInputNode" nodes present in the ReactFlow instance, create a new node
         // Calculate the position where the node should be created
         const position = reactFlowInstance!.project({
           x: event.clientX - reactflowBounds!.left,
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index cf62f2694..eecb0ccbb 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -1009,7 +1009,8 @@ export function expandGroupNode(
       gNodes[nodeIndex].data.node!.template[field].show = show;
       gNodes[nodeIndex].data.node!.template[field].advanced = advanced;
       gNodes[nodeIndex].data.node!.template[field].display_name = display_name;
-      gNodes[nodeIndex].selected = false;
+      // keep the nodes selected after ungrouping
+      // gNodes[nodeIndex].selected = false;
       if (proxy) {
         gNodes[nodeIndex].data.node!.template[field].proxy = proxy;
       } else {