diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 018e99a7c..f66e1e638 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -260,6 +260,7 @@ export default function GenericNode({
           openAdvancedModal={false}
           onCloseAdvancedModal={() => {}}
           selected={selected}
+          updateNode={handleUpdateCode}
         />
       </NodeToolbar>
     );
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 295384fea..72c437a05 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -46,6 +46,7 @@ export default function NodeToolbarComponent({
   selected,
   setShowState,
   onCloseAdvancedModal,
+  updateNode,
 }: nodeToolbarPropsType): JSX.Element {
   const version = useDarkStore((state) => state.version);
   const [showModalAdvanced, setShowModalAdvanced] = useState(false);
@@ -307,6 +308,9 @@ export default function NodeToolbarComponent({
       case "delete":
         deleteNode(data.id);
         break;
+      case "update":
+        updateNode();
+        break;
       case "copy":
         const node = nodes.filter((node) => node.id === data.id);
         setLastCopiedSelection({ nodes: _.cloneDeep(node), edges: [] });
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 8d1fef5d1..3cae3f838 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -491,7 +491,7 @@ export type ChatInputType = {
   isDragging: boolean;
   files: FilePreviewType[];
   setFiles: (
-    files: FilePreviewType[] | ((prev: FilePreviewType[]) => FilePreviewType[])
+    files: FilePreviewType[] | ((prev: FilePreviewType[]) => FilePreviewType[]),
   ) => void;
   chatValue: string;
   inputRef: {
@@ -547,6 +547,7 @@ export type nodeToolbarPropsType = {
   onCloseAdvancedModal?: (close: boolean) => void;
   selected: boolean;
   setShowState: (show: boolean | SetStateAction<boolean>) => void;
+  updateNode: () => void;
 };
 
 export type parsedDataType = {
@@ -593,7 +594,7 @@ export type chatMessagePropsType = {
   updateChat: (
     chat: ChatMessageType,
     message: string,
-    stream_url?: string
+    stream_url?: string,
   ) => void;
 };
 
@@ -685,12 +686,12 @@ export type codeTabsPropsType = {
       value: string,
       node: NodeType,
       template: TemplateVariableType,
-      tweak: tweakType
+      tweak: tweakType,
     ) => string;
     buildTweakObject?: (
       tw: string,
       changes: string | string[] | boolean | number | Object[] | Object,
-      template: TemplateVariableType
+      template: TemplateVariableType,
     ) => Promise<string | void>;
   };
   activeTweaks?: boolean;