From 4fcd7027ad42f3d1f0647dc4e65d6ecc7d7a4650 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <lucas.edu.oli@hotmail.com>
Date: Wed, 19 Jun 2024 15:12:50 -0300
Subject: [PATCH] implemented Refresh button on node toolbar dropdown that
 appears when the user edits and the component is outdated.

---
 .../src/CustomNodes/GenericNode/index.tsx     |  9 +++++---
 .../hooks/use-check-code-validity.tsx         | 16 ++++++--------
 .../hooks/use-update-node-code.tsx            |  2 ++
 .../components/nodeToolbarComponent/index.tsx | 21 ++++++++++++-------
 src/frontend/src/types/components/index.ts    |  2 +-
 src/frontend/src/utils/styleUtils.ts          |  2 ++
 6 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index ce0547eb6..922c33dc6 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -75,6 +75,7 @@ export default function GenericNode({
     data.node?.description!,
   );
   const [isOutdated, setIsOutdated] = useState(false);
+  const [isUserEdited, setIsUserEdited] = useState(false);
   const buildStatus = useFlowStore(
     (state) => state.flowBuildStatus[data.id]?.status,
   );
@@ -96,6 +97,7 @@ export default function GenericNode({
     data.node!,
     setNode,
     setIsOutdated,
+    setIsUserEdited,
     updateNodeInternals,
   );
 
@@ -170,7 +172,7 @@ export default function GenericNode({
     deleteNode(data.id);
   }
 
-  useCheckCodeValidity(data, templates, setIsOutdated, types);
+  useCheckCodeValidity(data, templates, setIsOutdated, setIsUserEdited, types);
   useUpdateValidationStatus(data?.id, flowPool, setValidationStatus);
   useValidationStatusString(validationStatus, setValidationString);
 
@@ -329,8 +331,8 @@ export default function GenericNode({
           showNode={showNode}
           openAdvancedModal={false}
           onCloseAdvancedModal={() => {}}
-          selected={selected}
           updateNode={handleUpdateCode}
+          isOutdated={isOutdated && isUserEdited}
         />
       </NodeToolbar>
     );
@@ -344,6 +346,7 @@ export default function GenericNode({
     showNode,
     updateNodeCode,
     isOutdated,
+    isUserEdited,
     selected,
     shortcuts,
     //    openWDoubleCLick,
@@ -435,7 +438,7 @@ export default function GenericNode({
                           {data.node?.display_name}
                         </div>
                       </ShadTooltip>
-                      {isOutdated && (
+                      {isOutdated && !isUserEdited && (
                         <ShadTooltip content={TOOLTIP_OUTDATED_NODE}>
                           <Button
                             onClick={handleUpdateCode}
diff --git a/src/frontend/src/CustomNodes/hooks/use-check-code-validity.tsx b/src/frontend/src/CustomNodes/hooks/use-check-code-validity.tsx
index ec4d586f6..dd9027069 100644
--- a/src/frontend/src/CustomNodes/hooks/use-check-code-validity.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-check-code-validity.tsx
@@ -6,6 +6,7 @@ const useCheckCodeValidity = (
   data: NodeDataType,
   templates: { [key: string]: any },
   setIsOutdated: (value: boolean) => void,
+  setIsUserEdited: (value: boolean) => void,
   types,
 ) => {
   useEffect(() => {
@@ -24,17 +25,12 @@ const useCheckCodeValidity = (
     const currentCode = thisNodeTemplate.code?.value;
     const thisNodesCode = data.node!.template?.code?.value;
     const componentsToIgnore = ["CustomComponent", "Prompt"];
-    if (
-      currentCode !== thisNodesCode &&
-      !componentsToIgnore.includes(data.type) &&
-      !(data.node?.edited ?? false)
-    ) {
-      setIsOutdated(true);
-    } else {
-      setIsOutdated(false);
-    }
+    setIsOutdated(
+      currentCode !== thisNodesCode && !componentsToIgnore.includes(data.type),
+    );
+    setIsUserEdited(data.node?.edited ?? false);
     // template.code can be undefined
-  }, [data.node?.template?.code?.value, templates, setIsOutdated]);
+  }, [data.node?.template?.code?.value, templates, setIsOutdated, setIsUserEdited]);
 };
 
 export default useCheckCodeValidity;
diff --git a/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx b/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
index f1593597f..d6c972ddf 100644
--- a/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
@@ -7,6 +7,7 @@ const useUpdateNodeCode = (
   dataNode: APIClassType, // Define YourNodeType according to your data structure
   setNode: (id: string, callback: (oldNode) => any) => void,
   setIsOutdated: (value: boolean) => void,
+  setIsUserEdited: (value: boolean) => void,
   updateNodeInternals: (id: string) => void,
 ) => {
   const updateNodeCode = useCallback(
@@ -24,6 +25,7 @@ const useUpdateNodeCode = (
 
         newNode.data.node.template[name].value = code;
         setIsOutdated(false);
+        setIsUserEdited(false);
 
         return newNode;
       });
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 89755f198..5e58825e8 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -41,10 +41,10 @@ export default function NodeToolbarComponent({
   numberOfOutputHandles,
   showNode,
   name = "code",
-  selected,
   setShowState,
   onCloseAdvancedModal,
   updateNode,
+  isOutdated,
 }: nodeToolbarPropsType): JSX.Element {
   const version = useDarkStore((state) => state.version);
   const [showModalAdvanced, setShowModalAdvanced] = useState(false);
@@ -90,13 +90,6 @@ export default function NodeToolbarComponent({
     return;
   }
 
-  function handleUpdateWShortcut(e: KeyboardEvent) {
-    e.preventDefault();
-    if (hasApiKey || hasStore) {
-      handleSelectChange("update");
-    }
-  }
-
   function handleGroupWShortcut(e: KeyboardEvent) {
     e.preventDefault();
     if (isGroup) {
@@ -609,6 +602,18 @@ export default function NodeToolbarComponent({
                   dataTestId="copy-button-modal"
                 />
               </SelectItem>
+              {isOutdated && (
+                <SelectItem value={"update"}>
+                  <ToolbarSelectItem
+                    shortcut={
+                      shortcuts.find((obj) => obj.name === "Update")?.shortcut!
+                    }
+                    value={"Restore"}
+                    icon={"RefreshCcwDot"}
+                    dataTestId="update-button-modal"
+                  />
+                </SelectItem>
+              )}
               {hasStore && (
                 <SelectItem
                   value={"Share"}
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index f9f985653..15090b4a0 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -562,7 +562,7 @@ export type nodeToolbarPropsType = {
   name?: string;
   openAdvancedModal?: boolean;
   onCloseAdvancedModal?: (close: boolean) => void;
-  selected: boolean;
+  isOutdated: boolean;
   setShowState: (show: boolean | SetStateAction<boolean>) => void;
   updateNode: () => void;
 };
diff --git a/src/frontend/src/utils/styleUtils.ts b/src/frontend/src/utils/styleUtils.ts
index c47d40495..8d53ebb05 100644
--- a/src/frontend/src/utils/styleUtils.ts
+++ b/src/frontend/src/utils/styleUtils.ts
@@ -109,6 +109,7 @@ import {
   PocketKnife,
   Redo,
   RefreshCcw,
+  RefreshCcwDot,
   Repeat,
   RotateCcw,
   Save,
@@ -453,6 +454,7 @@ export const nodeIconsLucide: iconsType = {
   ChevronLeft,
   SlidersHorizontal,
   Palette,
+  RefreshCcwDot,
   FolderUp,
   Blocks,
   ChevronDown,