diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 90167bc36..abb625aee 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -31,9 +31,11 @@ import { FlowsContext } from "../../../../contexts/flowsContext";
 import { typesContext } from "../../../../contexts/typesContext";
 import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 import { postCustomComponentUpdate } from "../../../../controllers/API";
+import { APIClassType } from "../../../../types/api";
 import { ParameterComponentType } from "../../../../types/components";
 import { NodeDataType } from "../../../../types/flow";
 import {
+  cleanEdges,
   convertObjToArray,
   convertValuesToNumbers,
   hasDuplicateKeys,
@@ -69,7 +71,8 @@ export default function ParameterComponent({
   const { setErrorData } = useContext(alertContext);
   const updateNodeInternals = useUpdateNodeInternals();
   const [position, setPosition] = useState(0);
-  const { setTabsState, tabId, flows, tabsState } = useContext(FlowsContext);
+  const { setTabsState, tabId, flows, tabsState, updateFlow } =
+    useContext(FlowsContext);
 
   const flow = flows.find((flow) => flow.id === tabId)?.data?.nodes ?? null;
 
@@ -139,6 +142,40 @@ export default function ParameterComponent({
     renderTooltips();
   };
 
+  const handleNodeClass = (newNodeClass: APIClassType, code?: string): void => {
+    if (data.node!.template[name].value !== newNodeClass.template[name].value) {
+      takeSnapshot();
+    }
+    data.node = newNodeClass;
+    data.node.template[name].value = code;
+    updateNodeInternals(data.id);
+    // Set state to pending
+    //@ts-ignore
+    if (data.node!.template[name].value !== code) {
+      const tabs = cloneDeep(tabsState);
+      tabs[tabId].isPending = false;
+      tabs[tabId].formKeysData = tabsState[tabId].formKeysData;
+      setTabsState({
+        ...tabs,
+      });
+    }
+    renderTooltips();
+    let flow = flows.find((flow) => flow.id === tabId);
+    if (reactFlowInstance && flow && flow.data) {
+      cleanEdges({
+        flow: {
+          edges: flow.data!.edges,
+          nodes: flow.data!.nodes,
+        },
+        updateEdge: (edge) => {
+          reactFlowInstance.setEdges(edge);
+          updateNodeInternals(data.id);
+        },
+      });
+      updateFlow(flow);
+    }
+  };
+
   const [errorDuplicateKey, setErrorDuplicateKey] = useState(false);
 
   useEffect(() => {
@@ -443,9 +480,7 @@ export default function ParameterComponent({
                   : false
               }
               dynamic={data.node?.template[name].dynamic ?? false}
-              setNodeClass={(nodeClass) => {
-                data.node = nodeClass;
-              }}
+              setNodeClass={handleNodeClass}
               nodeClass={data.node}
               disabled={disabled}
               value={data.node?.template[name].value ?? ""}
@@ -479,15 +514,11 @@ export default function ParameterComponent({
             <PromptAreaComponent
               readonly={data.node?.flow ? true : false}
               field_name={name}
-              setNodeClass={(nodeClass) => {
-                data.node = nodeClass;
-              }}
+              setNodeClass={handleNodeClass}
               nodeClass={data.node}
               disabled={disabled}
               value={data.node?.template[name].value ?? ""}
-              onChange={(e) => {
-                handleOnNewValue(e);
-              }}
+              onChange={handleOnNewValue}
               id={"prompt-input-" + index}
               data-testid={"prompt-input-" + index}
             />
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 526c10cfa..9e6b7024b 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -14,11 +14,7 @@ import { undoRedoContext } from "../../contexts/undoRedoContext";
 import NodeToolbarComponent from "../../pages/FlowPage/components/nodeToolbarComponent";
 import { validationStatusType } from "../../types/components";
 import { NodeDataType } from "../../types/flow";
-import {
-  cleanEdges,
-  handleKeyDown,
-  scapedJSONStringfy,
-} from "../../utils/reactflowUtils";
+import { handleKeyDown, scapedJSONStringfy } from "../../utils/reactflowUtils";
 import { nodeColors, nodeIconsLucide } from "../../utils/styleUtils";
 import { classNames, getFieldTitle } from "../../utils/utils";
 import ParameterComponent from "./components/parameterComponent";
@@ -85,27 +81,10 @@ export default function GenericNode({
 
   useEffect(() => {
     countHandles();
-  }, []);
+  }, [data, data.node]);
 
   // State for outline color
   const { sseData, isBuilding } = useSSE();
-  useEffect(() => {
-    let flow = flows.find((flow) => flow.id === tabId);
-    if (reactFlowInstance && flow && flow.data) {
-      cleanEdges({
-        flow: {
-          edges: flow.data!.edges,
-          nodes: flow.data!.nodes,
-        },
-        updateEdge: (edge) => {
-          reactFlowInstance.setEdges(edge);
-          updateNodeInternals(data.id);
-        },
-      });
-      updateFlow(flow);
-    }
-    countHandles();
-  }, [modalContextOpen]);
 
   useEffect(() => {
     setNodeDescription(data.node!.description);
diff --git a/src/frontend/src/components/promptComponent/index.tsx b/src/frontend/src/components/promptComponent/index.tsx
index f5878af12..5a2ced971 100644
--- a/src/frontend/src/components/promptComponent/index.tsx
+++ b/src/frontend/src/components/promptComponent/index.tsx
@@ -31,9 +31,7 @@ export default function PromptAreaComponent({
         value={value}
         buttonText="Check & Save"
         modalTitle="Edit Prompt"
-        setValue={(value: string) => {
-          onChange(value);
-        }}
+        setValue={onChange}
         nodeClass={nodeClass}
         setNodeClass={setNodeClass}
       >
diff --git a/src/frontend/src/modals/EditNodeModal/index.tsx b/src/frontend/src/modals/EditNodeModal/index.tsx
index 17a6bb70e..f6283870e 100644
--- a/src/frontend/src/modals/EditNodeModal/index.tsx
+++ b/src/frontend/src/modals/EditNodeModal/index.tsx
@@ -150,9 +150,9 @@ const EditNodeModal = forwardRef(
                         .map((templateParam, index) => {
                           let id = {
                             inputTypes:
-                              data.node!.template[templateParam].input_types,
-                            type: data.node!.template[templateParam].type,
-                            id: data.id,
+                              myData.node!.template[templateParam].input_types,
+                            type: myData.node!.template[templateParam].type,
+                            id: myData.id,
                             fieldName: templateParam,
                           };
                           let disabled =
@@ -160,11 +160,11 @@ const EditNodeModal = forwardRef(
                               (edge) =>
                                 edge.targetHandle ===
                                 scapedJSONStringfy(
-                                  data.node!.template[templateParam].proxy
+                                  myData.node!.template[templateParam].proxy
                                     ? {
                                         ...id,
                                         proxy:
-                                          data.node?.template[templateParam]
+                                          myData.node?.template[templateParam]
                                             .proxy,
                                       }
                                     : id
diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
index dd561aa97..f488704c8 100644
--- a/src/frontend/src/modals/codeAreaModal/index.tsx
+++ b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -88,8 +88,7 @@ export default function CodeAreaModal({
       .then((apiReturn) => {
         const { data } = apiReturn;
         if (data) {
-          setNodeClass(data);
-          setValue(code);
+          setNodeClass(data, code);
           setError({ detail: { error: undefined, traceback: undefined } });
           setOpen(false);
         }
diff --git a/src/frontend/src/modals/genericModal/index.tsx b/src/frontend/src/modals/genericModal/index.tsx
index ce63ca004..355d0afcd 100644
--- a/src/frontend/src/modals/genericModal/index.tsx
+++ b/src/frontend/src/modals/genericModal/index.tsx
@@ -114,7 +114,6 @@ export default function GenericModal({
         // if field_name is an empty string, then we need to set it
         // to the first key of the custom_fields object
         if (field_name === "") {
-          console.log(apiReturn.data?.frontend_node?.custom_fields);
           field_name = Array.isArray(
             apiReturn.data?.frontend_node?.custom_fields?.[""]
           )
@@ -134,9 +133,8 @@ export default function GenericModal({
               JSON.stringify(apiReturn.data?.frontend_node) !==
               JSON.stringify({})
             ) {
-              setNodeClass!(apiReturn.data?.frontend_node);
+              setNodeClass!(apiReturn.data?.frontend_node, inputValue);
               setModalOpen(closeModal);
-              setValue(inputValue);
               setIsEdit(false);
               setSuccessData({
                 title: "Prompt is ready",
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index fbd68f136..4ff343e6c 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -95,7 +95,7 @@ export type TextAreaComponentType = {
 export type PromptAreaComponentType = {
   field_name?: string;
   nodeClass?: APIClassType;
-  setNodeClass?: (value: APIClassType) => void;
+  setNodeClass?: (value: APIClassType, code?: string) => void;
   disabled: boolean;
   onChange: (value: string[] | string) => void;
   value: string;
@@ -110,7 +110,7 @@ export type CodeAreaComponentType = {
   value: string;
   editNode?: boolean;
   nodeClass?: APIClassType;
-  setNodeClass?: (value: APIClassType) => void;
+  setNodeClass?: (value: APIClassType, code?: string) => void;
   dynamic?: boolean;
   id?: string;
   readonly?: boolean;
@@ -500,7 +500,7 @@ export type codeAreaModalPropsType = {
   setValue: (value: string) => void;
   value: string;
   nodeClass: APIClassType | undefined;
-  setNodeClass: (Class: APIClassType) => void | undefined;
+  setNodeClass: (Class: APIClassType, code?: string) => void | undefined;
   children: ReactNode;
   dynamic?: boolean;
   readonly?: boolean;
@@ -526,7 +526,7 @@ export type genericModalPropsType = {
   modalTitle: string;
   type: number;
   nodeClass?: APIClassType;
-  setNodeClass?: (Class: APIClassType) => void;
+  setNodeClass?: (Class: APIClassType, code?: string) => void;
   children: ReactNode;
   id?: string;
   readonly?: boolean;