diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 624189f11..8183b3dab 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -10,6 +10,7 @@ import {
   RUN_TIMESTAMP_PREFIX,
   STATUS_BUILD,
   STATUS_BUILDING,
+  TOOLTIP_OUTDATED_NODE,
 } from "../../constants/constants";
 import { BuildStatus } from "../../constants/enums";
 import { countHandlesFn } from "../helpers/count-handles";
@@ -34,6 +35,8 @@ import useValidationStatusString from "../hooks/use-validation-status-string";
 import getFieldTitle from "../utils/get-field-title";
 import sortFields from "../utils/sort-fields";
 import ParameterComponent from "./components/parameterComponent";
+import { postCustomComponent } from "../../controllers/API";
+import { cloneDeep } from "lodash";
 
 export default function GenericNode({
   data,
@@ -201,6 +204,37 @@ export default function GenericNode({
     setShowNode(data.showNode ?? true);
   }, [data.showNode]);
 
+  const handleUpdateCode = () => {
+    takeSnapshot();
+    // to update we must get the code from the templates in useTypesStore
+    const thisNodeTemplate = templates[data.type]?.template;
+    // if the template does not have a code key
+    // return
+    if (!thisNodeTemplate?.code) return;
+
+    const currentCode = thisNodeTemplate.code.value;
+    if (data.node) {
+      postCustomComponent(currentCode, data.node)
+        .then((apiReturn) => {
+          const { data } = apiReturn;
+          if (data && updateNodeCode) {
+            updateNodeCode(data, currentCode, "code");
+          }
+        })
+        .catch((err) => {
+          console.log(err);
+        });
+      setNode(data.id, (oldNode) => {
+        let newNode = cloneDeep(oldNode);
+        newNode.data = {
+          ...data,
+        };
+        newNode.data.node.template.code.value = currentCode;
+        return newNode;
+      });
+    }
+  };
+
   const memoizedNodeToolbarComponent = useMemo(() => {
     return (
       <NodeToolbar>
@@ -221,8 +255,6 @@ export default function GenericNode({
           showNode={showNode}
           openAdvancedModal={false}
           onCloseAdvancedModal={() => {}}
-          updateNodeCode={updateNodeCode}
-          isOutdated={isOutdated}
           selected={selected}
         />
       </NodeToolbar>
@@ -439,12 +471,18 @@ export default function GenericNode({
               <>
                 <div className="flex flex-shrink-0 items-center gap-1">
                   {isOutdated && (
-                    <Button variant="secondary" className={"h-9 px-1.5"}>
-                      <IconComponent
-                        name="AlertTriangle"
-                        className="h-5 w-5 text-status-yellow"
-                      />
-                    </Button>
+                    <ShadTooltip content={TOOLTIP_OUTDATED_NODE}>
+                      <Button
+                        onClick={handleUpdateCode}
+                        variant="secondary"
+                        className={"h-9 px-1.5"}
+                      >
+                        <IconComponent
+                          name="AlertTriangle"
+                          className="h-5 w-5 text-status-yellow"
+                        />
+                      </Button>
+                    </ShadTooltip>
                   )}
                   <ShadTooltip
                     content={
diff --git a/src/frontend/src/constants/constants.ts b/src/frontend/src/constants/constants.ts
index a113b3911..d84989994 100644
--- a/src/frontend/src/constants/constants.ts
+++ b/src/frontend/src/constants/constants.ts
@@ -661,6 +661,9 @@ export const OUTPUT_TYPES = new Set([
 export const CHAT_FIRST_INITIAL_TEXT =
   "Start a conversation and click the agent's thoughts";
 
+export const TOOLTIP_OUTDATED_NODE =
+  "Your component is outdated. Click to update (data may be lost)";
+
 export const CHAT_SECOND_INITIAL_TEXT = "to inspect the chaining process.";
 
 export const ZERO_NOTIFICATIONS = "No new notifications";
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index e3bef9114..8636907d8 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -40,10 +40,8 @@ export default function NodeToolbarComponent({
   showNode,
   name = "code",
   selected,
-  updateNodeCode,
   setShowState,
   onCloseAdvancedModal,
-  isOutdated,
 }: nodeToolbarPropsType): JSX.Element {
   const nodeLength = Object.keys(data.node!.template).filter(
     (templateField) =>
@@ -179,37 +177,6 @@ export default function NodeToolbarComponent({
             paneY: nodes.find((node) => node.id === data.id)?.position.y,
           },
         );
-        break;
-      case "update":
-        takeSnapshot();
-        // to update we must get the code from the templates in useTypesStore
-        const thisNodeTemplate = templates[data.type]?.template;
-        // if the template does not have a code key
-        // return
-        if (!thisNodeTemplate?.code) return;
-
-        const currentCode = thisNodeTemplate.code.value;
-        if (data.node) {
-          postCustomComponent(currentCode, data.node)
-            .then((apiReturn) => {
-              const { data } = apiReturn;
-              if (data && updateNodeCode) {
-                updateNodeCode(data, currentCode, "code");
-              }
-            })
-            .catch((err) => {
-              console.log(err);
-            });
-          setNode(data.id, (oldNode) => {
-            let newNode = cloneDeep(oldNode);
-            newNode.data = {
-              ...data,
-            };
-            newNode.data.node.template.code.value = currentCode;
-            return newNode;
-          });
-        }
-
         break;
     }
   };
@@ -508,19 +475,6 @@ export default function NodeToolbarComponent({
                   dataTestId="copy-button-modal"
                 />
               </SelectItem>
-              {isOutdated && (
-                <SelectItem value={"update"}>
-                  <ToolbarSelectItem
-                    keyboardKey="U"
-                    isMac={navigator.userAgent.toUpperCase().includes("MAC")}
-                    shift={false}
-                    value={"Update"}
-                    icon={"Code"}
-                    dataTestId="update-button-modal"
-                    ping={isOutdated}
-                  />
-                </SelectItem>
-              )}
               {hasStore && (
                 <SelectItem
                   value={"Share"}
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 7d9712aa9..ade4f7c13 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -538,13 +538,7 @@ export type nodeToolbarPropsType = {
   openAdvancedModal?: boolean;
   onCloseAdvancedModal?: (close: boolean) => void;
   selected: boolean;
-  updateNodeCode?: (
-    newNodeClass: APIClassType,
-    code: string,
-    name: string,
-  ) => void;
   setShowState: (show: boolean | SetStateAction<boolean>) => void;
-  isOutdated?: boolean;
 };
 
 export type parsedDataType = {