diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 9c1b4bbf1..09af96c6f 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -1,6 +1,6 @@
 import { cloneDeep } from "lodash";
 import React, { ReactNode, useEffect, useRef, useState } from "react";
-import { Handle, Position } from "reactflow";
+import { Handle, Position, useUpdateNodeInternals } from "reactflow";
 import ShadTooltip from "../../../../components/ShadTooltipComponent";
 import CodeAreaComponent from "../../../../components/codeAreaComponent";
 import DictComponent from "../../../../components/dictComponent";
@@ -124,27 +124,32 @@ export default function ParameterComponent({
     renderTooltips();
   };
 
+  const updateNodeInternals = useUpdateNodeInternals();
+
   const handleNodeClass = (newNodeClass: APIClassType, code?: string): void => {
     if (!data.node) return;
     if (data.node!.template[name].value !== code) {
       takeSnapshot();
     }
-
+    
+    
     setNode(data.id, (oldNode) => {
       let newNode = cloneDeep(oldNode);
-
+      
       newNode.data = {
         ...newNode.data,
         node: newNodeClass,
         description: newNodeClass.description ?? data.node!.description,
         display_name: newNodeClass.display_name ?? data.node!.display_name,
       };
-
+      
       newNode.data.node.template[name].value = code;
-
+      
       return newNode;
     });
-
+    
+    updateNodeInternals(data.id);
+    
     renderTooltips();
   };
 
@@ -268,6 +273,9 @@ export default function ParameterComponent({
             <Handle
               type={left ? "target" : "source"}
               position={left ? Position.Left : Position.Right}
+              key={proxy
+                ? scapedJSONStringfy({ ...id, proxy })
+                : scapedJSONStringfy(id)}
               id={
                 proxy
                   ? scapedJSONStringfy({ ...id, proxy })
@@ -340,6 +348,9 @@ export default function ParameterComponent({
                 <Handle
                   type={left ? "target" : "source"}
                   position={left ? Position.Left : Position.Right}
+                  key={proxy
+                    ? scapedJSONStringfy({ ...id, proxy })
+                    : scapedJSONStringfy(id)}
                   id={
                     proxy
                       ? scapedJSONStringfy({ ...id, proxy })
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index 57716b5b5..634083b2c 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -183,7 +183,7 @@ export default function Page({
       takeSnapshot();
       onConnect(params);
     },
-    [setEdges, takeSnapshot, addEdge]
+    [takeSnapshot, onConnect]
   );
 
   const onNodeDragStart: NodeDragHandler = useCallback(() => {