From 6eeaa583fff123fb9b4db6f2438ff0de4b674c78 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 14 Nov 2024 11:30:11 -0800
Subject: [PATCH] Add cancel workflow functionality (#1193)

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
---
 .../src/routes/workflows/WorkflowRun.tsx      | 105 ++++++++++++++++--
 1 file changed, 93 insertions(+), 12 deletions(-)

diff --git a/skyvern-frontend/src/routes/workflows/WorkflowRun.tsx b/skyvern-frontend/src/routes/workflows/WorkflowRun.tsx
index 01a73be4..be390690 100644
--- a/skyvern-frontend/src/routes/workflows/WorkflowRun.tsx
+++ b/skyvern-frontend/src/routes/workflows/WorkflowRun.tsx
@@ -43,9 +43,11 @@ import {
   Pencil2Icon,
   PlayIcon,
   ReaderIcon,
+  ReloadIcon,
 } from "@radix-ui/react-icons";
 import {
   keepPreviousData,
+  useMutation,
   useQuery,
   useQueryClient,
 } from "@tanstack/react-query";
@@ -59,9 +61,23 @@ import {
 } from "react-router-dom";
 import { TaskActions } from "../tasks/list/TaskActions";
 import { TaskListSkeletonRows } from "../tasks/list/TaskListSkeletonRows";
-import { statusIsNotFinalized, statusIsRunningOrQueued } from "../tasks/types";
+import {
+  statusIsFinalized,
+  statusIsNotFinalized,
+  statusIsRunningOrQueued,
+} from "../tasks/types";
 import { CodeEditor } from "./components/CodeEditor";
 import { useWorkflowQuery } from "./hooks/useWorkflowQuery";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
 
 type StreamMessage = {
   task_id: string;
@@ -144,6 +160,35 @@ function WorkflowRun() {
       workflowRun?.status === Status.Running ? "always" : false,
   });
 
+  const cancelWorkflowMutation = useMutation({
+    mutationFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client
+        .post(`/workflows/runs/${workflowRunId}/cancel`)
+        .then((response) => response.data);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["workflowRun", workflowRunId],
+      });
+      queryClient.invalidateQueries({
+        queryKey: ["workflowRun", workflowPermanentId, workflowRunId],
+      });
+      toast({
+        variant: "success",
+        title: "Workflow Canceled",
+        description: "The workflow has been successfully canceled.",
+      });
+    },
+    onError: (error) => {
+      toast({
+        variant: "destructive",
+        title: "Error",
+        description: error.message,
+      });
+    },
+  });
+
   const currentRunningTask = workflowTasks?.find(
     (task) => task.status === Status.Running,
   );
@@ -151,6 +196,8 @@ function WorkflowRun() {
   const workflowRunIsRunningOrQueued =
     workflowRun && statusIsRunningOrQueued(workflowRun);
 
+  const workflowRunIsFinalized = workflowRun && statusIsFinalized(workflowRun);
+
   useEffect(() => {
     if (!workflowRunIsRunningOrQueued) {
       return;
@@ -347,17 +394,51 @@ function WorkflowRun() {
               Edit
             </Link>
           </Button>
-          <Button asChild>
-            <Link
-              to={`/workflows/${workflowPermanentId}/run`}
-              state={{
-                data: parameters,
-              }}
-            >
-              <PlayIcon className="mr-2 h-4 w-4" />
-              Rerun
-            </Link>
-          </Button>
+          {workflowRunIsRunningOrQueued && (
+            <Dialog>
+              <DialogTrigger asChild>
+                <Button variant="destructive">Cancel</Button>
+              </DialogTrigger>
+              <DialogContent>
+                <DialogHeader>
+                  <DialogTitle>Are you sure?</DialogTitle>
+                  <DialogDescription>
+                    Are you sure you want to cancel this workflow run?
+                  </DialogDescription>
+                </DialogHeader>
+                <DialogFooter>
+                  <DialogClose asChild>
+                    <Button variant="secondary">Back</Button>
+                  </DialogClose>
+                  <Button
+                    variant="destructive"
+                    onClick={() => {
+                      cancelWorkflowMutation.mutate();
+                    }}
+                    disabled={cancelWorkflowMutation.isPending}
+                  >
+                    {cancelWorkflowMutation.isPending && (
+                      <ReloadIcon className="mr-2 h-4 w-4 animate-spin" />
+                    )}
+                    Cancel Workflow Run
+                  </Button>
+                </DialogFooter>
+              </DialogContent>
+            </Dialog>
+          )}
+          {workflowRunIsFinalized && (
+            <Button asChild>
+              <Link
+                to={`/workflows/${workflowPermanentId}/run`}
+                state={{
+                  data: parameters,
+                }}
+              >
+                <PlayIcon className="mr-2 h-4 w-4" />
+                Rerun
+              </Link>
+            </Button>
+          )}
         </div>
       </header>
       {workflowRun && statusIsNotFinalized(workflowRun) && (