Canceling runs

apps/webapp/app/components/run/RunOverview.tsx

@@ -1,11 +1,19 @@
 import { conform, useForm } from "@conform-to/react";
 import { parse } from "@conform-to/zod";
 import { BoltIcon, ForwardIcon } from "@heroicons/react/24/solid";
-import { Form, Outlet, useActionData, useNavigate } from "@remix-run/react";
+import {
+  Form,
+  Outlet,
+  useActionData,
+  useLocation,
+  useNavigate,
+  useNavigation,
+} from "@remix-run/react";
 import { JobRunStatus, RuntimeEnvironmentType } from "@trigger.dev/database";
 import { useMemo } from "react";
 import { usePathName } from "~/hooks/usePathName";
 import { Run } from "~/presenters/RunPresenter.server";
+import { cancelSchema } from "~/routes/resources.runs.$runId.cancel";
 import { schema } from "~/routes/resources.runs.$runId.rerun";
 import { formatDuration } from "~/utils";
 import { cn } from "~/utils/cn";
@@ -123,6 +131,7 @@ export function RunOverview({
                 status={basicStatus}
               />
             )}
+            {!hasFinished(run.status) && <CancelRun runId={run.id} />}
           </PageButtons>
         </PageTitleRow>
         <PageInfoRow>
@@ -419,3 +428,42 @@ function RerunPopover({
     </Popover>
   );
 }
+
+export function CancelRun({ runId }: { runId: string }) {
+  const lastSubmission = useActionData();
+  const location = useLocation();
+  const navigation = useNavigation();
+
+  const [form, { redirectUrl }] = useForm({
+    id: "cancel-run",
+    lastSubmission,
+    onValidate({ formData }) {
+      return parse(formData, { schema: cancelSchema });
+    },
+  });
+
+  const isLoading =
+    navigation.state === "submitting" && navigation.formData !== undefined;
+
+  return (
+    <Form
+      method="post"
+      action={`/resources/runs/${runId}/cancel`}
+      {...form.props}
+    >
+      <input
+        {...conform.input(redirectUrl, { type: "hidden" })}
+        defaultValue={location.pathname}
+      />
+      <Button
+        type="submit"
+        LeadingIcon={isLoading ? "spinner-white" : "stop"}
+        leadingIconClassName="text-white"
+        variant="danger/small"
+        disabled={isLoading}
+      >
+        {isLoading ? "Canceling" : "Cancel run"}
+      </Button>
+    </Form>
+  );
+}

apps/webapp/app/components/runs/RunStatuses.tsx

@@ -3,18 +3,21 @@ import {
   CheckCircleIcon,
   ClockIcon,
   ExclamationTriangleIcon,
+  StopIcon,
   WrenchIcon,
   XCircleIcon,
 } from "@heroicons/react/24/solid";
 import { cn } from "~/utils/cn";
 import { Spinner } from "../primitives/Spinner";
+import { HandRaisedIcon, NoSymbolIcon } from "@heroicons/react/20/solid";
 
 export function hasFinished(status: JobRunStatus): boolean {
   return (
     status === "SUCCESS" ||
     status === "FAILURE" ||
     status === "ABORTED" ||
-    status === "TIMED_OUT"
+    status === "TIMED_OUT" ||
+    status === "CANCELED"
   );
 }
 
@@ -89,6 +92,12 @@ export function RunStatusIcon({
       return (
         <Spinner className={cn(runStatusClassNameColor(status), className)} />
       );
+    case "CANCELED":
+      return (
+        <NoSymbolIcon
+          className={cn(runStatusClassNameColor(status), className)}
+        />
+      );
   }
 }
 
@@ -119,6 +128,8 @@ export function runBasicStatus(status: JobRunStatus): RunBasicStatus {
       return "FAILED";
     case "PREPROCESSING":
       return "PENDING";
+    case "CANCELED":
+      return "FAILED";
   }
 }
 
@@ -142,6 +153,8 @@ export function runStatusTitle(status: JobRunStatus): string {
       return "Aborted";
     case "PREPROCESSING":
       return "Preprocessing";
+    case "CANCELED":
+      return "Canceled";
   }
 }
 
@@ -165,5 +178,7 @@ export function runStatusClassNameColor(status: JobRunStatus): string {
       return "text-rose-500";
     case "PREPROCESSING":
       return "text-blue-500";
+    case "CANCELED":
+      return "text-slate-500";
   }
 }

apps/webapp/app/routes/api.v1.runs.$runId.tasks.ts

@@ -192,6 +192,18 @@ export class RunTaskService {
         return existingTask;
       }
 
+      const run = await tx.jobRun.findUnique({
+        where: {
+          id: runId,
+        },
+        select: {
+          status: true,
+        },
+      });
+
+      if (!run) throw new Error("Run not found");
+      if (run.status === "CANCELED") return;
+
       // If task.delayUntil is set and is in the future, we'll set the task's status to "WAITING", else set it to RUNNING
       const status =
         (taskBody.delayUntil && taskBody.delayUntil.getTime() > Date.now()) ||

apps/webapp/app/routes/resources.runs.$runId.cancel.ts

@@ -0,0 +1,37 @@
+import { parse } from "@conform-to/zod";
+import { ActionFunction, json } from "@remix-run/node";
+import { z } from "zod";
+import { redirectWithSuccessMessage } from "~/models/message.server";
+import { CancelRunService } from "~/services/runs/cancelRun.server";
+
+export const cancelSchema = z.object({
+  redirectUrl: z.string(),
+});
+
+const ParamSchema = z.object({
+  runId: z.string(),
+});
+
+export const action: ActionFunction = async ({ request, params }) => {
+  const { runId } = ParamSchema.parse(params);
+
+  const formData = await request.formData();
+  const submission = parse(formData, { schema: cancelSchema });
+
+  if (!submission.value) {
+    return json(submission);
+  }
+
+  try {
+    const cancelRunService = new CancelRunService();
+    await cancelRunService.call({ runId });
+
+    return redirectWithSuccessMessage(
+      submission.value.redirectUrl,
+      request,
+      `Canceled run. Any pending tasks will be canceled.`
+    );
+  } catch (error: any) {
+    return json({ errors: { body: error.message } }, { status: 400 });
+  }
+};

apps/webapp/app/services/runs/cancelRun.server.ts

@@ -0,0 +1,43 @@
+import { PrismaClient, prisma } from "~/db.server";
+import { workerQueue } from "../worker.server";
+
+export class CancelRunService {
+  #prismaClient: PrismaClient;
+
+  constructor(prismaClient: PrismaClient = prisma) {
+    this.#prismaClient = prismaClient;
+  }
+
+  public async call({ runId }: { runId: string }) {
+    try {
+      return await this.#prismaClient.$transaction(async (tx) => {
+        const run = await tx.jobRun.update({
+          select: {
+            queueId: true,
+          },
+          where: { id: runId },
+          data: {
+            status: "CANCELED",
+            queue: {
+              update: {
+                jobCount: {
+                  decrement: 1,
+                },
+              },
+            },
+          },
+        });
+
+        await workerQueue.enqueue(
+          "startQueuedRuns",
+          {
+            id: run.queueId,
+          },
+          { tx }
+        );
+      });
+    } catch (error) {
+      throw error;
+    }
+  }
+}

apps/webapp/app/services/runs/continueRun.server.ts

@@ -2,7 +2,7 @@ import { EXECUTE_JOB_RETRY_LIMIT } from "~/consts";
 import { $transaction, Prisma, PrismaClient, prisma } from "~/db.server";
 import { workerQueue } from "../worker.server";
 
-const RESUMABLE_STATUSES = ["FAILURE", "TIMED_OUT", "ABORTED"];
+const RESUMABLE_STATUSES = ["FAILURE", "TIMED_OUT", "ABORTED", "CANCELED"];
 
 export class ContinueRunService {
   #prismaClient: PrismaClient;

packages/database/prisma/migrations/20230711175328_added_canceled_run_status/migration.sql

@@ -0,0 +1,2 @@
+-- AlterEnum
+ALTER TYPE "JobRunStatus" ADD VALUE 'CANCELED';

packages/database/prisma/schema.prisma

@@ -709,6 +709,7 @@ enum JobRunStatus {
   FAILURE
   TIMED_OUT
   ABORTED
+  CANCELED
 }
 
 model JobRunExecution {