Rework Workflow Save (#3145)

skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx

@@ -1,4 +1,3 @@
-import { getClient } from "@/api/AxiosClient";
 import { Button } from "@/components/ui/button";
 import {
   Dialog,
@@ -9,16 +8,18 @@ import {
   DialogTitle,
 } from "@/components/ui/dialog";
 import { toast } from "@/components/ui/use-toast";
-import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { useOnChange } from "@/hooks/useOnChange";
 import { useShouldNotifyWhenClosingTab } from "@/hooks/useShouldNotifyWhenClosingTab";
 import { DeleteNodeCallbackContext } from "@/store/DeleteNodeCallbackContext";
 import { useDebugStore } from "@/store/useDebugStore";
-import { useWorkflowHasChangesStore } from "@/store/WorkflowHasChangesStore";
+import {
+  useWorkflowHasChangesStore,
+  useWorkflowSave,
+  type WorkflowSaveData,
+} from "@/store/WorkflowHasChangesStore";
 import { useWorkflowPanelStore } from "@/store/WorkflowPanelStore";
 import { useWorkflowTitleStore } from "@/store/WorkflowTitleStore";
 import { ReloadIcon } from "@radix-ui/react-icons";
-import { useMutation, useQueryClient } from "@tanstack/react-query";
 import {
   Background,
   BackgroundVariant,
@@ -34,11 +35,9 @@ import {
   useReactFlow,
 } from "@xyflow/react";
 import "@xyflow/react/dist/style.css";
-import { AxiosError } from "axios";
 import { nanoid } from "nanoid";
-import { useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
-import { useBlocker, useParams } from "react-router-dom";
+import { useBlocker } from "react-router-dom";
-import { stringify as convertToYAML } from "yaml";
 import {
   AWSSecretParameter,
   debuggableWorkflowBlockTypes,
@@ -46,18 +45,15 @@ import {
   WorkflowEditorParameterTypes,
   WorkflowParameterTypes,
   WorkflowParameterValueType,
-  WorkflowSettings,
 } from "../types/workflowTypes";
 import {
   BitwardenCreditCardDataParameterYAML,
   BitwardenLoginCredentialParameterYAML,
   BitwardenSensitiveInformationParameterYAML,
-  BlockYAML,
   ContextParameterYAML,
   CredentialParameterYAML,
   OnePasswordCredentialParameterYAML,
   ParameterYAML,
-  WorkflowCreateYAMLRequest,
   WorkflowParameterYAML,
 } from "../types/workflowYamlTypes";
 import { WorkflowHeader } from "./WorkflowHeader";
@@ -265,9 +261,6 @@ function FlowRenderer({
 }: Props) {
   const reactFlowInstance = useReactFlow();
   const debugStore = useDebugStore();
-  const { workflowPermanentId } = useParams();
-  const credentialGetter = useCredentialGetter();
-  const queryClient = useQueryClient();
   const { workflowPanelState, setWorkflowPanelState, closeWorkflowPanel } =
     useWorkflowPanelStore();
   const { title, initializeTitle } = useWorkflowTitleStore();
@@ -289,107 +282,16 @@ function FlowRenderer({
     initializeTitle(initialTitle);
   }, [initialTitle, initializeTitle]);
 
-  const { hasChanges, setHasChanges } = useWorkflowHasChangesStore();
+  const workflowChangesStore = useWorkflowHasChangesStore();
-  useShouldNotifyWhenClosingTab(hasChanges);
+  const setGetSaveDataRef = useRef(workflowChangesStore.setGetSaveData);
+  setGetSaveDataRef.current = workflowChangesStore.setGetSaveData;
+  const saveWorkflow = useWorkflowSave();
+  useShouldNotifyWhenClosingTab(workflowChangesStore.hasChanges);
   const blocker = useBlocker(({ currentLocation, nextLocation }) => {
-    return hasChanges && nextLocation.pathname !== currentLocation.pathname;
+    return (
-  });
+      workflowChangesStore.hasChanges &&
-
+      nextLocation.pathname !== currentLocation.pathname
-  const saveWorkflowMutation = useMutation({
+    );
-    mutationFn: async (data: {
-      parameters: Array<ParameterYAML>;
-      blocks: Array<BlockYAML>;
-      title: string;
-      settings: WorkflowSettings;
-    }) => {
-      if (!workflowPermanentId) {
-        return;
-      }
-      const client = await getClient(credentialGetter);
-      const extraHttpHeaders: Record<string, string> = {};
-      if (data.settings.extraHttpHeaders) {
-        try {
-          const parsedHeaders = JSON.parse(data.settings.extraHttpHeaders);
-          if (
-            parsedHeaders &&
-            typeof parsedHeaders === "object" &&
-            !Array.isArray(parsedHeaders)
-          ) {
-            for (const [key, value] of Object.entries(parsedHeaders)) {
-              if (key && typeof key === "string") {
-                if (key in extraHttpHeaders) {
-                  toast({
-                    title: "Error",
-                    description: `Duplicate key '${key}' in extra http headers`,
-                    variant: "destructive",
-                  });
-                  continue;
-                }
-                extraHttpHeaders[key] = String(value);
-              }
-            }
-          }
-        } catch (error) {
-          toast({
-            title: "Error",
-            description: "Invalid JSON format in extra http headers",
-            variant: "destructive",
-          });
-          return;
-        }
-      }
-
-      const requestBody: WorkflowCreateYAMLRequest = {
-        title: data.title,
-        description: workflow.description,
-        proxy_location: data.settings.proxyLocation,
-        webhook_callback_url: data.settings.webhookCallbackUrl,
-        persist_browser_session: data.settings.persistBrowserSession,
-        model: data.settings.model,
-        max_screenshot_scrolls: data.settings.maxScreenshotScrolls,
-        totp_verification_url: workflow.totp_verification_url,
-        extra_http_headers: extraHttpHeaders,
-        use_cache: data.settings.useScriptCache,
-        cache_key: data.settings.scriptCacheKey,
-        workflow_definition: {
-          parameters: data.parameters,
-          blocks: data.blocks,
-        },
-        is_saved_task: workflow.is_saved_task,
-      };
-      const yaml = convertToYAML(requestBody);
-      return client.put<string, WorkflowApiResponse>(
-        `/workflows/${workflowPermanentId}`,
-        yaml,
-        {
-          headers: {
-            "Content-Type": "text/plain",
-          },
-        },
-      );
-    },
-    onSuccess: () => {
-      toast({
-        title: "Changes saved",
-        description: "Your changes have been saved",
-        variant: "success",
-      });
-      queryClient.invalidateQueries({
-        queryKey: ["workflow", workflowPermanentId],
-      });
-      queryClient.invalidateQueries({
-        queryKey: ["workflows"],
-      });
-      setHasChanges(false);
-    },
-    onError: (error: AxiosError) => {
-      const detail = (error.response?.data as { detail?: string })?.detail;
-      toast({
-        title: "Error",
-        description: detail ? detail : error.message,
-        variant: "destructive",
-      });
-    },
   });
 
   function doLayout(nodes: Array<AppNode>, edges: Array<Edge>) {
@@ -424,7 +326,7 @@ function FlowRenderer({
     setDebuggableBlockCount(debuggable.length);
   }, [nodes, edges]);
 
-  async function handleSave() {
+  const constructSaveData = useCallback((): WorkflowSaveData => {
     const blocks = getWorkflowBlocks(nodes, edges);
     const settings = getWorkflowSettings(nodes);
     const parametersInYAMLConvertibleJSON = convertToParametersYAML(parameters);
@@ -447,7 +349,7 @@ function FlowRenderer({
       overallParameters,
     );
 
-    return saveWorkflowMutation.mutateAsync({
+    return {
       parameters: [
         ...echoParameters,
         ...parametersInYAMLConvertibleJSON,
@@ -456,7 +358,16 @@ function FlowRenderer({
       blocks,
       title,
       settings,
-    });
+      workflow,
+    };
+  }, [nodes, edges, parameters, title, workflow]);
+
+  useEffect(() => {
+    setGetSaveDataRef.current(constructSaveData);
+  }, [constructSaveData]);
+
+  async function handleSave() {
+    return await saveWorkflow.mutateAsync();
   }
 
   function addNode({
@@ -537,7 +448,7 @@ function FlowRenderer({
       ...nodes.slice(previousNodeIndex + 1),
     ];
 
-    setHasChanges(true);
+    workflowChangesStore.setHasChanges(true);
     doLayout(newNodesAfter, [...editedEdges, ...newEdges]);
   }
 
@@ -627,7 +538,8 @@ function FlowRenderer({
       }
       return node;
     });
-    setHasChanges(true);
+    workflowChangesStore.setHasChanges(true);
+
     doLayout(newNodesWithUpdatedParameters, newEdges);
   }
 
@@ -717,9 +629,9 @@ function FlowRenderer({
                   blocker.proceed?.();
                 });
               }}
-              disabled={saveWorkflowMutation.isPending}
+              disabled={workflowChangesStore.saveIsPending}
             >
-              {saveWorkflowMutation.isPending && (
+              {workflowChangesStore.saveIsPending && (
                 <ReloadIcon className="mr-2 h-4 w-4 animate-spin" />
               )}
               Save changes
@@ -763,7 +675,7 @@ function FlowRenderer({
                   );
                 })
               ) {
-                setHasChanges(true);
+                workflowChangesStore.setHasChanges(true);
               }
               onNodesChange(changes);
             }}
@@ -807,7 +719,7 @@ function FlowRenderer({
             <Panel position="top-center" className={cn("h-20")}>
               <WorkflowHeader
                 debuggableBlockCount={debuggableBlockCount}
-                saving={saveWorkflowMutation.isPending}
+                saving={workflowChangesStore.saveIsPending}
                 parametersPanelOpen={
                   workflowPanelState.active &&
                   workflowPanelState.content === "parameters"

skyvern-frontend/src/routes/workflows/editor/WorkflowDebugger.tsx

@@ -62,9 +62,7 @@ function WorkflowDebugger() {
     return state.setCollapsed;
   });
 
-  const setHasChanges = useWorkflowHasChangesStore(
+  const workflowChangesStore = useWorkflowHasChangesStore();
-    (state) => state.setHasChanges,
-  );
 
   const handleOnCycle = () => {
     setOpenDialogue(true);
@@ -72,7 +70,7 @@ function WorkflowDebugger() {
 
   useMountEffect(() => {
     setCollapsed(true);
-    setHasChanges(false);
+    workflowChangesStore.setHasChanges(false);
 
     if (workflowPermanentId) {
       queryClient.removeQueries({

skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx

@@ -16,9 +16,7 @@ function WorkflowEditor() {
   const setCollapsed = useSidebarStore((state) => {
     return state.setCollapsed;
   });
-  const setHasChanges = useWorkflowHasChangesStore(
+  const workflowChangesStore = useWorkflowHasChangesStore();
-    (state) => state.setHasChanges,
-  );
 
   const { data: workflow, isLoading } = useWorkflowQuery({
     workflowPermanentId,
@@ -29,7 +27,7 @@ function WorkflowEditor() {
 
   useMountEffect(() => {
     setCollapsed(true);
-    setHasChanges(false);
+    workflowChangesStore.setHasChanges(false);
   });
 
   if (isLoading || isGlobalWorkflowsLoading) {

skyvern-frontend/src/routes/workflows/editor/WorkflowHeader.tsx

@@ -40,7 +40,7 @@ function WorkflowHeader({
   saving,
 }: Props) {
   const { title, setTitle } = useWorkflowTitleStore();
-  const { setHasChanges } = useWorkflowHasChangesStore();
+  const workflowChangesStore = useWorkflowHasChangesStore();
   const { blockLabel: urlBlockLabel, workflowPermanentId } = useParams();
   const { data: globalWorkflows } = useGlobalWorkflowsQuery();
   const navigate = useNavigate();
@@ -68,7 +68,7 @@ function WorkflowHeader({
           editable={true}
           onChange={(newTitle) => {
             setTitle(newTitle);
-            setHasChanges(true);
+            workflowChangesStore.setHasChanges(true);
           }}
           value={title}
           titleClassName="text-3xl"

skyvern-frontend/src/routes/workflows/editor/nodes/components/NodeHeader.tsx

@@ -21,6 +21,7 @@ import {
 } from "@/routes/workflows/types/workflowTypes";
 import { getInitialValues } from "@/routes/workflows/utils";
 import { useDebugStore } from "@/store/useDebugStore";
+import { useWorkflowSave } from "@/store/WorkflowHasChangesStore";
 import {
   useWorkflowSettingsStore,
   type WorkflowSettingsState,
@@ -155,6 +156,7 @@ function NodeHeader({
   const { data: debugSession } = useDebugSessionQuery({
     workflowPermanentId,
   });
+  const saveWorkflow = useWorkflowSave();
 
   useEffect(() => {
     if (!workflowRun || !workflowPermanentId || !workflowRunId) {
@@ -190,6 +192,8 @@ function NodeHeader({
 
   const runBlock = useMutation({
     mutationFn: async () => {
+      await saveWorkflow.mutateAsync();
+
       if (!workflowPermanentId) {
         console.error("There is no workflowPermanentId");
         toast({

skyvern-frontend/src/store/WorkflowHasChangesStore.ts

@@ -1,13 +1,179 @@
+import { AxiosError } from "axios";
+import { useEffect } from "react";
 import { create } from "zustand";
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import { stringify as convertToYAML } from "yaml";
 
-type WorkflowHasChangesStore = {
+import { getClient } from "@/api/AxiosClient";
-  hasChanges: boolean;
+import { toast } from "@/components/ui/use-toast";
-  setHasChanges: (hasChanges: boolean) => void;
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
+import type {
+  BlockYAML,
+  ParameterYAML,
+} from "@/routes/workflows/types/workflowYamlTypes";
+import type {
+  WorkflowApiResponse,
+  WorkflowSettings,
+} from "@/routes/workflows/types/workflowTypes";
+import { WorkflowCreateYAMLRequest } from "@/routes/workflows/types/workflowYamlTypes";
+type SaveData = {
+  parameters: Array<ParameterYAML>;
+  blocks: Array<BlockYAML>;
+  title: string;
+  settings: WorkflowSettings;
+  workflow: WorkflowApiResponse;
 };
 
-const useWorkflowHasChangesStore = create<WorkflowHasChangesStore>((set) => ({
+type WorkflowHasChangesStore = {
-  hasChanges: false,
+  getSaveData: () => SaveData | null;
-  setHasChanges: (hasChanges) => set({ hasChanges }),
+  hasChanges: boolean;
-}));
+  saveIsPending: boolean;
+  setGetSaveData: (getSaveData: () => SaveData) => void;
+  setHasChanges: (hasChanges: boolean) => void;
+  setSaveIsPending: (isPending: boolean) => void;
+};
 
-export { useWorkflowHasChangesStore };
+const useWorkflowHasChangesStore = create<WorkflowHasChangesStore>((set) => {
+  return {
+    hasChanges: false,
+    saveIsPending: false,
+    getSaveData: () => null,
+    setGetSaveData: (getSaveData: () => SaveData) => {
+      set({ getSaveData });
+    },
+    setHasChanges: (hasChanges: boolean) => {
+      set({ hasChanges });
+    },
+    setSaveIsPending: (isPending: boolean) => {
+      set({ saveIsPending: isPending });
+    },
+  };
+});
+
+const useWorkflowSave = () => {
+  const credentialGetter = useCredentialGetter();
+  const queryClient = useQueryClient();
+  const { getSaveData, setHasChanges, setSaveIsPending } =
+    useWorkflowHasChangesStore();
+
+  const saveWorkflowMutation = useMutation({
+    mutationFn: async () => {
+      const saveData = getSaveData();
+
+      if (!saveData) {
+        setHasChanges(false);
+        return;
+      }
+
+      const client = await getClient(credentialGetter);
+      const extraHttpHeaders: Record<string, string> = {};
+
+      if (saveData.settings.extraHttpHeaders) {
+        try {
+          const parsedHeaders = JSON.parse(saveData.settings.extraHttpHeaders);
+          if (
+            parsedHeaders &&
+            typeof parsedHeaders === "object" &&
+            !Array.isArray(parsedHeaders)
+          ) {
+            for (const [key, value] of Object.entries(parsedHeaders)) {
+              if (key && typeof key === "string") {
+                if (key in extraHttpHeaders) {
+                  toast({
+                    title: "Error",
+                    description: `Duplicate key '${key}' in extra http headers`,
+                    variant: "destructive",
+                  });
+                  continue;
+                }
+                extraHttpHeaders[key] = String(value);
+              }
+            }
+          }
+        } catch (error) {
+          toast({
+            title: "Error",
+            description: "Invalid JSON format in extra http headers",
+            variant: "destructive",
+          });
+          return;
+        }
+      }
+
+      const requestBody: WorkflowCreateYAMLRequest = {
+        title: saveData.title,
+        description: saveData.workflow.description,
+        proxy_location: saveData.settings.proxyLocation,
+        webhook_callback_url: saveData.settings.webhookCallbackUrl,
+        persist_browser_session: saveData.settings.persistBrowserSession,
+        model: saveData.settings.model,
+        max_screenshot_scrolls: saveData.settings.maxScreenshotScrolls,
+        totp_verification_url: saveData.workflow.totp_verification_url,
+        extra_http_headers: extraHttpHeaders,
+        use_cache: saveData.settings.useScriptCache,
+        cache_key: saveData.settings.scriptCacheKey,
+        workflow_definition: {
+          parameters: saveData.parameters,
+          blocks: saveData.blocks,
+        },
+        is_saved_task: saveData.workflow.is_saved_task,
+      };
+
+      const yaml = convertToYAML(requestBody);
+
+      return client.put<string, WorkflowApiResponse>(
+        `/workflows/${saveData.workflow.workflow_permanent_id}`,
+        yaml,
+        {
+          headers: {
+            "Content-Type": "text/plain",
+          },
+        },
+      );
+    },
+    onSuccess: () => {
+      const saveData = getSaveData();
+
+      if (!saveData) {
+        return;
+      }
+
+      toast({
+        title: "Changes saved",
+        description: "Your changes have been saved",
+        variant: "success",
+      });
+
+      queryClient.invalidateQueries({
+        queryKey: ["workflow", saveData.workflow.workflow_permanent_id],
+      });
+
+      queryClient.invalidateQueries({
+        queryKey: ["workflows"],
+      });
+
+      setHasChanges(false);
+    },
+    onError: (error: AxiosError) => {
+      const detail = (error.response?.data as { detail?: string })?.detail;
+
+      toast({
+        title: "Error",
+        description: detail ? detail : error.message,
+        variant: "destructive",
+      });
+    },
+  });
+
+  useEffect(() => {
+    setSaveIsPending(saveWorkflowMutation.isPending);
+  }, [saveWorkflowMutation.isPending, setSaveIsPending]);
+
+  return saveWorkflowMutation;
+};
+
+export {
+  useWorkflowSave,
+  useWorkflowHasChangesStore,
+  type SaveData as WorkflowSaveData,
+};