diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index c39249b22..9e60bb7d2 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -46,7 +46,6 @@ export default function GenericNode({
   );
   const [validationStatus, setValidationStatus] =
     useState<validationStatusType | null>(null);
-  const [showNode, setShowNode] = useState<boolean>(true);
   const [handles, setHandles] = useState<boolean[] | []>([]);
   let numberOfInputs: boolean[] = [];
 
@@ -108,12 +107,6 @@ export default function GenericNode({
     countHandles();
   }, [data]);
 
-  useEffect(() => {
-    setTimeout(() => {
-      updateNodeInternals(data.id);
-    }, 300);
-  }, [showNode]);
-
   useEffect(() => {
     setNodeDescription(data.node!.description);
   }, [data.node!.description]);
@@ -132,6 +125,9 @@ export default function GenericNode({
       setValidationStatus(null);
     }
   }, [sseData, data.id]);
+  
+  const showNode = data.showNode ?? false;
+
   return (
     <>
       <NodeToolbar>
@@ -139,7 +135,7 @@ export default function GenericNode({
           position={{ x: xPos, y: yPos }}
           data={data}
           deleteNode={deleteNode}
-          setShowNode={setShowNode}
+          setShowNode={(showNode: boolean) => {data.showNode = showNode}}
           numberOfHandles={handles}
           showNode={showNode}
         ></NodeToolbarComponent>
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 0b450fcd1..fb43309fd 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -23,6 +23,7 @@ import {
   updateFlowPosition,
 } from "../../../../utils/reactflowUtils";
 import { classNames } from "../../../../utils/utils";
+import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 
 export default function NodeToolbarComponent({
   data,
@@ -65,6 +66,7 @@ export default function NodeToolbarComponent({
   const isGroup = data.node?.flow ? true : false;
 
   const { paste, saveComponent, version, flows } = useContext(FlowsContext);
+  const { takeSnapshot } = useContext(undoRedoContext);
   const reactFlowInstance = useReactFlow();
   const [showModalAdvanced, setShowModalAdvanced] = useState(false);
   const [showconfirmShare, setShowconfirmShare] = useState(false);
@@ -83,6 +85,7 @@ export default function NodeToolbarComponent({
         setShowModalAdvanced(true);
         break;
       case "show":
+        takeSnapshot();
         setShowNode((prev) => !prev);
         updateNodeInternals(data.id);
         break;
@@ -98,6 +101,7 @@ export default function NodeToolbarComponent({
       case "disabled":
         break;
       case "ungroup":
+        takeSnapshot();
         updateFlowPosition(position, data.node?.flow!);
         expandGroupNode(data, reactFlowInstance, getNodeId);
         break;
diff --git a/src/frontend/src/types/flow/index.ts b/src/frontend/src/types/flow/index.ts
index e66c96304..9f1821a08 100644
--- a/src/frontend/src/types/flow/index.ts
+++ b/src/frontend/src/types/flow/index.ts
@@ -22,6 +22,7 @@ export type NodeType = {
 };
 
 export type NodeDataType = {
+  showNode?: boolean;
   type: string;
   node?: APIClassType;
   id: string;