diff --git a/src/frontend/src/components/tableNodeCellRender/index.tsx b/src/frontend/src/components/tableNodeCellRender/index.tsx
index c93d6ce09..0db35ee57 100644
--- a/src/frontend/src/components/tableNodeCellRender/index.tsx
+++ b/src/frontend/src/components/tableNodeCellRender/index.tsx
@@ -1,4 +1,4 @@
-import { CustomCellRendererProps } from "ag-grid-react";
+import { CustomCellEditorProps, CustomCellRendererProps } from "ag-grid-react";
 import { classNames, cn, isTimeStampString } from "../../utils/utils";
 import ArrayReader from "../arrayReaderComponent";
 import DateReader from "../dateReaderComponent";
@@ -30,21 +30,27 @@ import { useState } from "react";
 import useFlowStore from "../../stores/flowStore";
 
 export default function TableNodeCellRender({
-  node: { data: templateData },
+  node: { data },
   value: {
-    value: templateValue,
+    value,
     nodeClass,
-    handleOnNewValue,
+    handleOnNewValue: handleOnNewValueNode,
     handleOnChangeDb,
   },
 }: CustomCellRendererProps) {
-  console.log(
-    templateData,
-    templateValue,
-    nodeClass,
-    handleOnNewValue,
-    handleOnChangeDb,
-  );
+  const handleOnNewValue = (newValue: any, name: string) => {
+    handleOnNewValueNode(newValue, name);
+    setTemplateData((old) => {
+      let newData = cloneDeep(old);
+      newData.value = newValue;
+      return newData;
+    });
+    setTemplateValue(newValue);
+  };
+
+  const [templateValue, setTemplateValue] = useState(value);
+  const [templateData, setTemplateData] = useState(data);
+
   const [errorDuplicateKey, setErrorDuplicateKey] = useState(false);
   const edges = useFlowStore((state) => state.edges);
 
@@ -71,79 +77,65 @@ export default function TableNodeCellRender({
     switch (templateData.type) {
       case "str":
         if (!templateData.options) {
-          return (
-            <div className="w-[300px]">
-              {templateData?.list ? (
-                <InputListComponent
-                  componentName={templateData.key ?? undefined}
-                  editNode={true}
-                  disabled={disabled}
-                  value={
-                    !templateValue || templateValue === ""
-                      ? [""]
-                      : templateValue
-                  }
-                  onChange={(value: string[]) => {
-                    handleOnNewValue(value, templateData.key);
-                  }}
-                />
-              ) : templateData.multiline ? (
-                <TextAreaComponent
-                  id={"textarea-edit-" + templateData.name}
-                  data-testid={"textarea-edit-" + templateData.name}
-                  disabled={disabled}
-                  editNode={true}
-                  value={templateData.value ?? ""}
-                  onChange={(value: string | string[]) => {
-                    handleOnNewValue(value, templateData.key);
-                  }}
-                />
-              ) : (
-                <InputGlobalComponent
-                  disabled={disabled}
-                  editNode={true}
-                  onChange={(value) =>
-                    handleOnNewValue(value, templateData.key)
-                  }
-                  setDb={(value) => {
-                    handleOnChangeDb(value, templateData.key);
-                  }}
-                  name={templateData.key}
-                  data={templateData}
-                />
-              )}
-            </div>
+          return templateData?.list ? (
+            <InputListComponent
+              componentName={templateData.key ?? undefined}
+              editNode={true}
+              disabled={disabled}
+              value={
+                !templateValue || templateValue === "" ? [""] : templateValue
+              }
+              onChange={(value: string[]) => {
+                handleOnNewValue(value, templateData.key);
+              }}
+            />
+          ) : templateData.multiline ? (
+            <TextAreaComponent
+              id={"textarea-edit-" + templateData.name}
+              data-testid={"textarea-edit-" + templateData.name}
+              disabled={disabled}
+              editNode={true}
+              value={templateValue ?? ""}
+              onChange={(value: string | string[]) => {
+                handleOnNewValue(value, templateData.key);
+              }}
+            />
+          ) : (
+            <InputGlobalComponent
+              disabled={disabled}
+              editNode={true}
+              onChange={(value) => handleOnNewValue(value, templateData.key)}
+              setDb={(value) => {
+                handleOnChangeDb(value, templateData.key);
+              }}
+              name={templateData.key}
+              data={templateData}
+            />
           );
         } else {
           return (
-            <div className="w-[300px]">
-              <Dropdown
-                editNode={true}
-                options={templateData.options}
-                onSelect={(value) => handleOnNewValue(value, templateData.key)}
-                value={templateData.value ?? "Choose an option"}
-                id={"dropdown-edit-" + templateData.name}
-              ></Dropdown>
-            </div>
+            <Dropdown
+              editNode={true}
+              options={templateData.options}
+              onSelect={(value) => handleOnNewValue(value, templateData.key)}
+              value={templateValue ?? "Choose an option"}
+              id={"dropdown-edit-" + templateData.name}
+            ></Dropdown>
           );
         }
 
       case "NestedDict":
         return (
-          <div className="  w-full">
-            <DictComponent
-              disabled={disabled}
-              editNode={true}
-              value={templateValue.toString() === "{}" ? {} : templateValue}
-              onChange={(newValue) => {
-                templateValue = newValue;
-                handleOnNewValue(newValue, templateData.key);
-              }}
-              id="editnode-div-dict-input"
-            />
-          </div>
+          <DictComponent
+            disabled={disabled}
+            editNode={true}
+            value={templateValue.toString() === "{}" ? {} : templateValue}
+            onChange={(newValue) => {
+              handleOnNewValue(newValue, templateData.key);
+            }}
+            id="editnode-div-dict-input"
+          />
         );
-        break;
 
       case "dict":
         return (
@@ -157,14 +149,13 @@ export default function TableNodeCellRender({
               disabled={disabled}
               editNode={true}
               value={
-                templateData.value?.length === 0 || !templateData.value
+                templateValue?.length === 0 || !templateValue
                   ? [{ "": "" }]
                   : convertObjToArray(templateValue, templateData.type)
               }
               duplicateKey={errorDuplicateKey}
               onChange={(newValue) => {
                 const valueToNumbers = convertValuesToNumbers(newValue);
-                templateValue = valueToNumbers;
                 setErrorDuplicateKey(hasDuplicateKeys(valueToNumbers));
                 handleOnNewValue(valueToNumbers, templateData.key);
               }}
@@ -172,128 +163,110 @@ export default function TableNodeCellRender({
             />
           </div>
         );
-        break;
 
       case "bool":
         return (
-          <div className="ml-auto">
-            {" "}
-            <ToggleShadComponent
-              id={"toggle-edit-" + templateData.name}
-              disabled={disabled}
-              enabled={templateData.value}
-              setEnabled={(isEnabled) => {
-                handleOnNewValue(isEnabled, templateData.key);
-              }}
-              size="small"
-              editNode={true}
-            />
-          </div>
+          <ToggleShadComponent
+            id={"toggle-edit-" + templateData.name}
+            disabled={disabled}
+            enabled={templateValue}
+            setEnabled={(isEnabled) => {
+              handleOnNewValue(isEnabled, templateData.key);
+            }}
+            size="small"
+            editNode={true}
+          />
         );
 
       case "float":
         return (
-          <div className="w-[300px]">
-            <FloatComponent
-              disabled={disabled}
-              editNode={true}
-              rangeSpec={templateData.rangeSpec}
-              value={templateData.value ?? ""}
-              onChange={(value) => {
-                handleOnNewValue(value, templateData.key);
-              }}
-            />
-          </div>
+          <FloatComponent
+            disabled={disabled}
+            editNode={true}
+            rangeSpec={templateData.rangeSpec}
+            value={templateValue ?? ""}
+            onChange={(value) => {
+              handleOnNewValue(value, templateData.key);
+            }}
+          />
         );
       case "int":
         return (
-          <div className="w-[300px]">
-            <IntComponent
-              rangeSpec={templateData.rangeSpec}
-              id={"edit-int-input-" + templateData.name}
-              disabled={disabled}
-              editNode={true}
-              value={templateData.value ?? ""}
-              onChange={(value) => {
-                handleOnNewValue(value, templateData.key);
-              }}
-            />
-          </div>
+          <IntComponent
+            rangeSpec={templateData.rangeSpec}
+            id={"edit-int-input-" + templateData.name}
+            disabled={disabled}
+            editNode={true}
+            value={templateValue ?? ""}
+            onChange={(value) => {
+              handleOnNewValue(value, templateData.key);
+            }}
+          />
         );
-        break;
 
       case "file":
         return (
-          <div className="w-[300px]">
-            <InputFileComponent
-              editNode={true}
-              disabled={disabled}
-              value={templateData.value ?? ""}
-              onChange={(value: string | string[]) => {
-                handleOnNewValue(value, templateData.key);
-              }}
-              fileTypes={templateData.fileTypes}
-              onFileChange={(filePath: string) => {
-                templateData.file_path = filePath;
-              }}
-            ></InputFileComponent>
-          </div>
+          <InputFileComponent
+            editNode={true}
+            disabled={disabled}
+            value={templateValue ?? ""}
+            onChange={(value: string | string[]) => {
+              handleOnNewValue(value, templateData.key);
+            }}
+            fileTypes={templateData.fileTypes}
+            onFileChange={(filePath: string) => {
+              templateData.file_path = filePath;
+            }}
+          />
         );
 
       case "prompt":
         return (
-          <div className="w-[300px]">
-            <PromptAreaComponent
-              readonly={nodeClass.flow ? true : false}
-              field_name={templateData.key}
-              editNode={true}
-              disabled={disabled}
-              nodeClass={nodeClass}
-              setNodeClass={(value) => {
-                nodeClass = value;
-              }}
-              value={templateValue ?? ""}
-              onChange={(value: string | string[]) => {
-                handleOnNewValue(value, templateData.key);
-              }}
-              id={"prompt-area-edit-" + templateData.name}
-              data-testid={"modal-prompt-input-" + templateData.name}
-            />
-          </div>
+          <PromptAreaComponent
+            readonly={nodeClass.flow ? true : false}
+            field_name={templateData.key}
+            editNode={true}
+            disabled={disabled}
+            nodeClass={nodeClass}
+            setNodeClass={(value) => {
+              nodeClass = value;
+            }}
+            value={templateValue ?? ""}
+            onChange={(value: string | string[]) => {
+              handleOnNewValue(value, templateData.key);
+            }}
+            id={"prompt-area-edit-" + templateData.name}
+            data-testid={"modal-prompt-input-" + templateData.name}
+          />
         );
-        break;
 
       case "code":
         return (
-          <div className="w-[300px]">
-            <CodeAreaComponent
-              readonly={nodeClass.flow && templateData.dynamic ? true : false}
-              dynamic={templateData.dynamic ?? false}
-              setNodeClass={(value) => {
-                nodeClass = value;
-              }}
-              nodeClass={nodeClass}
-              disabled={disabled}
-              editNode={true}
-              value={templateData.value ?? ""}
-              onChange={(value: string | string[]) => {
-                handleOnNewValue(value, templateData.key);
-              }}
-              id={"code-area-edit" + templateData.name}
-            />
-          </div>
+          <CodeAreaComponent
+            readonly={nodeClass.flow && templateData.dynamic ? true : false}
+            dynamic={templateData.dynamic ?? false}
+            setNodeClass={(value) => {
+              nodeClass = value;
+            }}
+            nodeClass={nodeClass}
+            disabled={disabled}
+            editNode={true}
+            value={templateValue ?? ""}
+            onChange={(value: string | string[]) => {
+              handleOnNewValue(value, templateData.key);
+            }}
+            id={"code-area-edit" + templateData.name}
+          />
         );
-        break;
       case "Any":
         return <>-</>;
-        break;
       default:
         return String(templateValue);
     }
   }
 
   return (
-    <div className="group flex h-full w-full items-center align-middle">
+    <div className="group flex h-full w-[300px] items-center justify-center">
       {getCellType()}
     </div>
   );
diff --git a/src/frontend/src/components/toggleShadComponent/index.tsx b/src/frontend/src/components/toggleShadComponent/index.tsx
index ef1dc0d33..e76518f03 100644
--- a/src/frontend/src/components/toggleShadComponent/index.tsx
+++ b/src/frontend/src/components/toggleShadComponent/index.tsx
@@ -29,20 +29,18 @@ export default function ToggleShadComponent({
   }
 
   return (
-    <div className={disabled ? "pointer-events-none cursor-not-allowed " : ""}>
-      <Switch
-        id={id}
-        data-testid={id}
-        style={{
-          transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
-        }}
-        disabled={disabled}
-        className=""
-        checked={enabled}
-        onCheckedChange={(isEnabled: boolean) => {
-          setEnabled(isEnabled);
-        }}
-      ></Switch>
-    </div>
+    <Switch
+      id={id}
+      data-testid={id}
+      style={{
+        transform: `scaleX(${scaleX}) scaleY(${scaleY})`,
+      }}
+      disabled={disabled}
+      className=""
+      checked={enabled}
+      onCheckedChange={(isEnabled: boolean) => {
+        setEnabled(isEnabled);
+      }}
+    ></Switch>
   );
 }
diff --git a/src/frontend/src/modals/editNodeModal/index.tsx b/src/frontend/src/modals/editNodeModal/index.tsx
index 771129d94..90aa35b22 100644
--- a/src/frontend/src/modals/editNodeModal/index.tsx
+++ b/src/frontend/src/modals/editNodeModal/index.tsx
@@ -1,5 +1,5 @@
 import { cloneDeep } from "lodash";
-import { forwardRef, useEffect, useState } from "react";
+import { forwardRef, useEffect, useRef, useState } from "react";
 import IconComponent from "../../components/genericIconComponent";
 import { Badge } from "../../components/ui/badge";
 import {
@@ -35,36 +35,28 @@ const EditNodeModal = forwardRef(
 
     const dataFromStore = nodes.find((node) => node.id === node.id)?.data;
 
-    const [myData, setMyData] = useState(dataFromStore ?? data);
+    const myData = useRef(dataFromStore ?? data);
 
     const setNode = useFlowStore((state) => state.setNode);
 
     function changeAdvanced(n) {
-      setMyData((old) => {
-        let newData = cloneDeep(old);
-        newData.node!.template[n].advanced =
-          !newData.node!.template[n].advanced;
-        return newData;
-      });
+      myData.current.node!.template[n].advanced =
+        !myData.current.node!.template[n].advanced;
     }
 
     const handleOnNewValue = (newValue: any, name) => {
-      setMyData((old) => {
-        let newData = cloneDeep(old);
-        newData.node!.template[name].value = newValue;
-        return newData;
-      });
+      myData.current.node!.template[name].value = newValue;
     };
 
     useEffect(() => {
       if (open) {
-        setMyData(data); // reset data to what it is on node when opening modal
+        myData.current = data; // reset data to what it is on node when opening modal
       }
     }, [open]);
 
-    const rowData = Object.keys(myData.node!.template)
+    const rowData = Object.keys(myData.current.node!.template)
       .filter((key: string) => {
-        const templateParam = myData.node!.template[
+        const templateParam = myData.current.node!.template[
           key
         ] as TemplateVariableType;
         return (
@@ -78,7 +70,7 @@ const EditNodeModal = forwardRef(
         );
       })
       .map((key: string) => {
-        const templateParam = myData.node!.template[
+        const templateParam = myData.current.node!.template[
           key
         ] as TemplateVariableType;
         return {
@@ -115,21 +107,16 @@ const EditNodeModal = forwardRef(
         field: "value",
         cellRenderer: TableNodeCellRender,
         valueGetter: (params: ValueGetterParams) => {
-          console.log("params", params);
           return {
             value: params.data.value,
-            nodeClass: myData.node,
+            nodeClass: myData.current.node,
             handleOnNewValue: handleOnNewValue,
             handleOnChangeDb: (value, key) => {
-              setMyData((oldData) => {
-                let newData = cloneDeep(oldData);
-                newData.node!.template[key].load_from_db = value;
-                return newData;
-              });
+              myData.current.node!.template[key].load_from_db = value;
             },
           };
         },
-        minWidth: 300,
+        minWidth: 330,
         flex: 1,
         resizable: false,
       },
@@ -159,14 +146,14 @@ const EditNodeModal = forwardRef(
         open={open}
         setOpen={setOpen}
         onChangeOpenModal={(open) => {
-          setMyData(data);
+          myData.current = data;
         }}
         onSubmit={() => {
           setNode(data.id, (old) => ({
             ...old,
             data: {
               ...old.data,
-              node: myData.node,
+              node: myData.current.node,
             },
           }));
           setOpen(false);
@@ -175,9 +162,9 @@ const EditNodeModal = forwardRef(
         <BaseModal.Trigger>
           <></>
         </BaseModal.Trigger>
-        <BaseModal.Header description={myData.node?.description!}>
-          <span className="pr-2">{myData.type}</span>
-          <Badge variant="secondary">ID: {myData.id}</Badge>
+        <BaseModal.Header description={myData.current.node?.description!}>
+          <span className="pr-2">{myData.current.type}</span>
+          <Badge variant="secondary">ID: {myData.current.id}</Badge>
         </BaseModal.Header>
         <BaseModal.Content>
           <div className="flex pb-2">