From 47c1a6f0e879cce7645b5ad9061dab9ba33b6c2d Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Wed, 26 Jun 2024 15:57:24 -0700
Subject: [PATCH] Task cancelling frontend (#522)

---
 skyvern-frontend/src/api/types.ts             |   1 +
 .../src/components/StatusBadge.tsx            |   3 +-
 .../src/routes/tasks/detail/TaskDetails.tsx   | 101 ++++++++++++++++--
 .../src/routes/tasks/list/TaskHistory.tsx     |   2 +
 4 files changed, 98 insertions(+), 9 deletions(-)

diff --git a/skyvern-frontend/src/api/types.ts b/skyvern-frontend/src/api/types.ts
index c9d55b6e..0b697e63 100644
--- a/skyvern-frontend/src/api/types.ts
+++ b/skyvern-frontend/src/api/types.ts
@@ -19,6 +19,7 @@ export const Status = {
   Completed: "completed",
   Queued: "queued",
   TimedOut: "timed_out",
+  Canceled: "canceled",
 } as const;
 
 export type Status = (typeof Status)[keyof typeof Status];
diff --git a/skyvern-frontend/src/components/StatusBadge.tsx b/skyvern-frontend/src/components/StatusBadge.tsx
index b92eba88..4e97e3b8 100644
--- a/skyvern-frontend/src/components/StatusBadge.tsx
+++ b/skyvern-frontend/src/components/StatusBadge.tsx
@@ -12,7 +12,8 @@ function StatusBadge({ status }: Props) {
   } else if (
     status === "failed" ||
     status === "terminated" ||
-    status === "timed_out"
+    status === "timed_out" ||
+    status === "canceled"
   ) {
     variant = "destructive";
   } else if (status === "running") {
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
index bbb75e7d..68d215f2 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
@@ -1,17 +1,36 @@
 import { getClient } from "@/api/AxiosClient";
 import { Status, TaskApiResponse } from "@/api/types";
 import { StatusBadge } from "@/components/StatusBadge";
+import { Button } from "@/components/ui/button";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
 import { Label } from "@/components/ui/label";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Textarea } from "@/components/ui/textarea";
+import { toast } from "@/components/ui/use-toast";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { cn } from "@/util/utils";
-import { keepPreviousData, useQuery } from "@tanstack/react-query";
+import { ReloadIcon } from "@radix-ui/react-icons";
+import {
+  keepPreviousData,
+  useMutation,
+  useQuery,
+  useQueryClient,
+} from "@tanstack/react-query";
 import { NavLink, Outlet, useParams } from "react-router-dom";
 
 function TaskDetails() {
   const { taskId } = useParams();
   const credentialGetter = useCredentialGetter();
+  const queryClient = useQueryClient();
 
   const {
     data: task,
@@ -36,6 +55,35 @@ function TaskDetails() {
     placeholderData: keepPreviousData,
   });
 
+  const cancelTaskMutation = useMutation({
+    mutationFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client
+        .post(`/tasks/${taskId}/cancel`)
+        .then((response) => response.data);
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["task", taskId],
+      });
+      queryClient.invalidateQueries({
+        queryKey: ["tasks"],
+      });
+      toast({
+        variant: "success",
+        title: "Task Canceled",
+        description: "The task has been successfully canceled.",
+      });
+    },
+    onError: (error) => {
+      toast({
+        variant: "destructive",
+        title: "Error",
+        description: error.message,
+      });
+    },
+  });
+
   if (taskIsError) {
     return <div>Error: {taskError?.message}</div>;
   }
@@ -53,6 +101,9 @@ function TaskDetails() {
     </div>
   ) : null;
 
+  const taskIsRunningOrQueued =
+    task?.status === Status.Running || task?.status === Status.Queued;
+
   const showFailureReason =
     task?.status === Status.Failed ||
     task?.status === Status.Terminated ||
@@ -70,13 +121,47 @@ function TaskDetails() {
 
   return (
     <div className="flex flex-col gap-8">
-      <div className="flex items-center gap-4">
-        <span className="text-lg">{taskId}</span>
-        {taskIsLoading ? (
-          <Skeleton className="w-28 h-8" />
-        ) : task ? (
-          <StatusBadge status={task?.status} />
-        ) : null}
+      <div className="flex justify-between items-center">
+        <div className="flex items-center gap-4">
+          <span className="text-lg">{taskId}</span>
+          {taskIsLoading ? (
+            <Skeleton className="w-28 h-8" />
+          ) : task ? (
+            <StatusBadge status={task?.status} />
+          ) : null}
+        </div>
+        {taskIsRunningOrQueued && (
+          <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 task?
+                </DialogDescription>
+              </DialogHeader>
+              <DialogFooter>
+                <DialogClose asChild>
+                  <Button variant="secondary">Back</Button>
+                </DialogClose>
+                <Button
+                  variant="destructive"
+                  onClick={() => {
+                    cancelTaskMutation.mutate();
+                  }}
+                  disabled={cancelTaskMutation.isPending}
+                >
+                  {cancelTaskMutation.isPending && (
+                    <ReloadIcon className="mr-2 h-4 w-4 animate-spin" />
+                  )}
+                  Cancel Task
+                </Button>
+              </DialogFooter>
+            </DialogContent>
+          </Dialog>
+        )}
       </div>
       {taskIsLoading ? (
         <div className="flex items-center gap-2">
diff --git a/skyvern-frontend/src/routes/tasks/list/TaskHistory.tsx b/skyvern-frontend/src/routes/tasks/list/TaskHistory.tsx
index 4fa1c9b0..50fc376d 100644
--- a/skyvern-frontend/src/routes/tasks/list/TaskHistory.tsx
+++ b/skyvern-frontend/src/routes/tasks/list/TaskHistory.tsx
@@ -48,6 +48,8 @@ function TaskHistory() {
       params.append("task_status", "failed");
       params.append("task_status", "terminated");
       params.append("task_status", "timed_out");
+      params.append("task_status", "canceled");
+
       return client
         .get("/tasks", {
           params,