v3: task run status and canceling runs (#941)

* WIP task run status, revamped resuming task dependencies

* Don’t select a span when toggling collapsing

* Cancel runs and attempts, in prod and dev

apps/coordinator/package.json

@@ -4,7 +4,6 @@
   "version": "0.0.1",
   "description": "",
   "main": "dist/index.cjs",
-  "type": "module",
   "scripts": {
     "build": "npm run build:bundle",
     "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"",
@@ -31,4 +30,4 @@
     "tsx": "^4.7.0",
     "typescript": "^5.3.3"
   }
-}
+}

apps/coordinator/src/index.ts

@@ -301,6 +301,16 @@ class TaskCoordinator {
 
           taskSocket.emit("RESUME_AFTER_DURATION", message);
         },
+        REQUEST_ATTEMPT_CANCELLATION: async (message) => {
+          const taskSocket = await this.#getAttemptSocket(message.attemptId);
+
+          if (!taskSocket) {
+            logger.log("Socket for attempt not found", { attemptId: message.attemptId });
+            return;
+          }
+
+          taskSocket.emit("REQUEST_ATTEMPT_CANCELLATION", message);
+        },
       },
     });
 
@@ -385,11 +395,21 @@ class TaskCoordinator {
 
           if (!executionAck) {
             logger.error("no execution ack", { attemptId: socket.data.attemptId });
+
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+
             return;
           }
 
           if (!executionAck.success) {
             logger.error("execution unsuccessful", { attemptId: socket.data.attemptId });
+
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+
             return;
           }
 

apps/coordinator/tsconfig.json

@@ -9,10 +9,10 @@
     "strict": true,
     "skipLibCheck": true,
     "paths": {
-      "@trigger.dev/core/v3": ["../core/src/v3"],
-      "@trigger.dev/core/v3/*": ["../core/src/v3/*"],
-      "@trigger.dev/core-apps": ["../core-apps/src"],
-      "@trigger.dev/core-apps/*": ["../core-apps/src/*"]
+      "@trigger.dev/core/v3": ["../../packages/core/src/v3"],
+      "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"],
+      "@trigger.dev/core-apps": ["../../packages/core-apps/src"],
+      "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"]
     }
   }
 }

apps/docker-provider/package.json

@@ -4,7 +4,6 @@
   "version": "0.0.1",
   "description": "",
   "main": "dist/index.cjs",
-  "type": "module",
   "scripts": {
     "build": "npm run build:bundle",
     "build:bundle": "esbuild src/index.ts --bundle --outfile=dist/index.mjs --platform=node --format=esm --target=esnext --banner:js=\"const require = createRequire(import.meta.url);\"",
@@ -29,4 +28,4 @@
     "tsx": "^4.7.0",
     "typescript": "^5.3.3"
   }
-}
+}

apps/docker-provider/tsconfig.json

@@ -7,10 +7,10 @@
     "strict": true,
     "skipLibCheck": true,
     "paths": {
-      "@trigger.dev/core/v3": ["../core/src/v3"],
-      "@trigger.dev/core/v3/*": ["../core/src/v3/*"],
-      "@trigger.dev/core-apps": ["../core-apps/src"],
-      "@trigger.dev/core-apps/*": ["../core-apps/src/*"]
+      "@trigger.dev/core/v3": ["../../packages/core/src/v3"],
+      "@trigger.dev/core/v3/*": ["../../packages/core/src/v3/*"],
+      "@trigger.dev/core-apps": ["../../packages/core-apps/src"],
+      "@trigger.dev/core-apps/*": ["../../packages/core-apps/src/*"]
     }
   }
 }

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx

@@ -334,7 +334,6 @@ function TasksTreeView({
                       onClick={(e) => {
                         e.stopPropagation();
                         toggleExpandNode(node.id);
-                        selectNode(node.id);
                         scrollToNode(node.id);
                       }}
                     >

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route.tsx

@@ -52,7 +52,7 @@ export default function Page() {
         <div className={cn("grid h-full max-h-full grid-cols-1")}>
           <ResizablePanelGroup direction="horizontal" className="h-full max-h-full">
             <ResizablePanel order={1} minSize={20} defaultSize={30}>
-              <div className="flex flex-col px-3">
+              <div className="flex h-full max-h-full flex-col overflow-hidden px-3">
                 {tasks.length === 0 ? (
                   <NoTaskInstructions />
                 ) : (
@@ -81,7 +81,7 @@ function TaskSelector({ tasks }: { tasks: TaskListItem[] }) {
   const project = useProject();
 
   return (
-    <div className="flex flex-col divide-y divide-charcoal-800">
+    <div className="flex flex-col divide-y divide-charcoal-800 overflow-y-auto">
       {tasks.map((t) => (
         <NavLink key={t.id} to={v3TestTaskPath(organization, project, t)}>
           {({ isActive, isPending }) => (

apps/webapp/app/routes/api.v2.runs.$runParam.cancel.ts

@@ -0,0 +1,52 @@
+import type { ActionFunctionArgs } from "@remix-run/server-runtime";
+import { json } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { prisma } from "~/db.server";
+import { authenticateApiRequest } from "~/services/apiAuth.server";
+import { CancelTaskRunService } from "~/v3/services/cancelTaskRun.server";
+
+const ParamsSchema = z.object({
+  runParam: z.string(),
+});
+
+export async function action({ request, params }: ActionFunctionArgs) {
+  // Ensure this is a POST request
+  if (request.method.toUpperCase() !== "POST") {
+    return { status: 405, body: "Method Not Allowed" };
+  }
+
+  // Authenticate the request
+  const authenticationResult = await authenticateApiRequest(request);
+
+  if (!authenticationResult) {
+    return json({ error: "Invalid or Missing API Key" }, { status: 401 });
+  }
+
+  const parsed = ParamsSchema.safeParse(params);
+
+  if (!parsed.success) {
+    return json({ error: "Invalid or Missing runId" }, { status: 400 });
+  }
+
+  const { runParam } = parsed.data;
+
+  const taskRun = await prisma.taskRun.findUnique({
+    where: {
+      friendlyId: runParam,
+    },
+  });
+
+  if (!taskRun) {
+    return json({ error: "Run not found" }, { status: 404 });
+  }
+
+  const service = new CancelTaskRunService();
+
+  try {
+    await service.call(taskRun);
+  } catch (error) {
+    return json({ error: "Internal Server Error" }, { status: 500 });
+  }
+
+  return json({ message: "Run cancelled" }, { status: 200 });
+}

apps/webapp/app/services/worker.server.ts

@@ -30,6 +30,9 @@ import { DeliverWebhookRequestService } from "./sources/deliverWebhookRequest.se
 import { PerformTaskOperationService } from "./tasks/performTaskOperation.server";
 import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout.server";
 import { ResumeTaskService } from "./tasks/resumeTask.server";
+import { ResumeTaskRunDependenciesService } from "~/v3/services/resumeTaskRunDependencies.server";
+import { ResumeBatchRunService } from "~/v3/services/resumeBatchRun.server";
+import { ResumeTaskDependencyService } from "~/v3/services/resumeTaskDependency.server";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -107,6 +110,17 @@ const workerCatalog = {
   "v3.indexDeployment": z.object({
     id: z.string(),
   }),
+  "v3.resumeTaskRunDependencies": z.object({
+    attemptId: z.string(),
+  }),
+  "v3.resumeBatchRun": z.object({
+    batchRunId: z.string(),
+    sourceTaskAttemptId: z.string(),
+  }),
+  "v3.resumeTaskDependency": z.object({
+    dependencyId: z.string(),
+    sourceTaskAttemptId: z.string(),
+  }),
 };
 
 const executionWorkerCatalog = {
@@ -443,6 +457,33 @@ function getWorkerQueue() {
           return await service.call(payload.id);
         },
       },
+      "v3.resumeTaskRunDependencies": {
+        priority: 0,
+        maxAttempts: 5,
+        handler: async (payload, job) => {
+          const service = new ResumeTaskRunDependenciesService();
+
+          return await service.call(payload.attemptId);
+        },
+      },
+      "v3.resumeBatchRun": {
+        priority: 0,
+        maxAttempts: 5,
+        handler: async (payload, job) => {
+          const service = new ResumeBatchRunService();
+
+          return await service.call(payload.batchRunId, payload.sourceTaskAttemptId);
+        },
+      },
+      "v3.resumeTaskDependency": {
+        priority: 0,
+        maxAttempts: 5,
+        handler: async (payload, job) => {
+          const service = new ResumeTaskDependencyService();
+
+          return await service.call(payload.dependencyId, payload.sourceTaskAttemptId);
+        },
+      },
     },
   });
 }

apps/webapp/app/v3/authenticatedSocketConnection.server.ts

@@ -10,6 +10,7 @@ import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { DevQueueConsumer } from "./marqs/devQueueConsumer.server";
 import type { WebSocket, MessageEvent, CloseEvent, ErrorEvent } from "ws";
+import { env } from "~/env.server";
 
 export class AuthenticatedSocketConnection {
   public id: string;
@@ -26,6 +27,10 @@ export class AuthenticatedSocketConnection {
       schema: serverWebsocketMessages,
       sender: async (message) => {
         return new Promise((resolve, reject) => {
+          if (!ws.OPEN) {
+            return reject(new Error("Websocket is not open"));
+          }
+
           ws.send(JSON.stringify(message), {}, (err) => {
             if (err) {
               reject(err);
@@ -84,6 +89,8 @@ export class AuthenticatedSocketConnection {
   }
 
   async #handleClose(ev: CloseEvent) {
+    logger.debug("[AuthenticatedSocketConnection] Websocket closed", { ev });
+
     await this._consumer.stop();
 
     this.onClose.post(ev);

apps/webapp/app/v3/eventRepository.server.ts

@@ -44,6 +44,7 @@ export type TraceAttributes = Partial<
     CreatableEvent,
     | "attemptId"
     | "isError"
+    | "isCancelled"
     | "runId"
     | "runIsTest"
     | "output"

apps/webapp/app/v3/marqs/devPubSub.server.ts

@@ -0,0 +1,38 @@
+import { z } from "zod";
+import { singleton } from "~/utils/singleton";
+import { ZodPubSub, ZodSubscriber } from "../utils/zodPubSub.server";
+import { env } from "~/env.server";
+
+const messageCatalog = {
+  CANCEL_ATTEMPT: z.object({
+    version: z.literal("v1").default("v1"),
+    backgroundWorkerId: z.string(),
+    attemptId: z.string(),
+    taskRunId: z.string(),
+  }),
+};
+
+export type DevSubscriber = ZodSubscriber<typeof messageCatalog>;
+
+export const devPubSub = singleton("devPubSub", initializeDevPubSub);
+
+function initializeDevPubSub() {
+  return new ZodPubSub({
+    redis: {
+      port: env.REDIS_PORT,
+      host: env.REDIS_HOST,
+      username: env.REDIS_USERNAME,
+      password: env.REDIS_PASSWORD,
+      enableAutoPipelining: true,
+      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+    },
+    schema: {
+      CANCEL_ATTEMPT: z.object({
+        version: z.literal("v1").default("v1"),
+        backgroundWorkerId: z.string(),
+        attemptId: z.string(),
+        taskRunId: z.string(),
+      }),
+    },
+  });
+}

apps/webapp/app/v3/marqs/devQueueConsumer.server.ts

@@ -17,6 +17,7 @@ import { marqs } from "../marqs.server";
 import { CancelAttemptService } from "../services/cancelAttempt.server";
 import { CompleteAttemptService } from "../services/completeAttempt.server";
 import { attributesFromAuthenticatedEnv } from "../tracer.server";
+import { DevSubscriber, devPubSub } from "./devPubSub.server";
 
 const tracer = trace.getTracer("devQueueConsumer");
 
@@ -36,9 +37,11 @@ export type DevQueueConsumerOptions = {
 
 export class DevQueueConsumer {
   private _backgroundWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
+  private _backgroundWorkerSubscriber: Map<string, DevSubscriber> = new Map();
   private _deprecatedWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
   private _enabled = false;
-  private _options: Required<DevQueueConsumerOptions>;
+  private _maximumItemsPerTrace: number;
+  private _traceTimeoutSeconds: number;
   private _perTraceCountdown: number | undefined;
   private _lastNewTrace: Date | undefined;
   private _currentSpanContext: Context | undefined;
@@ -51,12 +54,10 @@ export class DevQueueConsumer {
   constructor(
     public env: AuthenticatedEnvironment,
     private _sender: ZodMessageSender<typeof serverWebsocketMessages>,
-    options: DevQueueConsumerOptions = {}
+    private _options: DevQueueConsumerOptions = {}
   ) {
-    this._options = {
-      maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace
-      traceTimeoutSeconds: options.traceTimeoutSeconds ?? 60, // 60 seconds
-    };
+    this._traceTimeoutSeconds = _options.traceTimeoutSeconds ?? 60;
+    this._maximumItemsPerTrace = _options.maximumItemsPerTrace ?? 1_000;
   }
 
   // This method is called when a background worker is deprecated and will no longer be used unless a run is locked to it
@@ -87,6 +88,21 @@ export class DevQueueConsumer {
 
     logger.debug("Registered background worker", { backgroundWorker: backgroundWorker.id });
 
+    const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`);
+
+    subscriber.on("CANCEL_ATTEMPT", async (message) => {
+      await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
+        backgroundWorkerId: backgroundWorker.friendlyId,
+        data: {
+          type: "CANCEL_ATTEMPT",
+          taskAttemptId: message.attemptId,
+          taskRunId: message.taskRunId,
+        },
+      });
+    });
+
+    this._backgroundWorkerSubscriber.set(backgroundWorker.id, subscriber);
+
     // Start reading from the queue if we haven't already
     this.#enable();
   }
@@ -133,6 +149,16 @@ export class DevQueueConsumer {
 
     // We need to cancel all the in progress task run attempts and ack the messages so they will stop processing
     await this.#cancelInProgressAttempts(reason);
+
+    // We need to unsubscribe from the background worker channels
+    for (const [id, subscriber] of this._backgroundWorkerSubscriber) {
+      logger.debug("Unsubscribing from background worker channel", { id });
+
+      await subscriber.stopListening();
+      this._backgroundWorkerSubscriber.delete(id);
+
+      logger.debug("Unsubscribed from background worker channel", { id });
+    }
   }
 
   async #cancelInProgressAttempts(reason: string) {
@@ -144,6 +170,10 @@ export class DevQueueConsumer {
 
     this._inProgressAttempts.clear();
 
+    logger.debug("Cancelling in progress attempts", {
+      attempts: Array.from(inProgressAttempts.keys()),
+    });
+
     for (const [attemptId, messageId] of inProgressAttempts) {
       await this.#cancelInProgressAttempt(attemptId, messageId, service, cancelledAt, reason);
     }
@@ -156,6 +186,8 @@ export class DevQueueConsumer {
     cancelledAt: Date,
     reason: string
   ) {
+    logger.debug("Cancelling in progress attempt", { attemptId, messageId });
+
     try {
       await cancelAttemptService.call(attemptId, messageId, cancelledAt, reason, this.env);
     } catch (e) {
@@ -189,7 +221,7 @@ export class DevQueueConsumer {
     // Check if the trace has expired
     if (
       this._perTraceCountdown === 0 ||
-      Date.now() - this._lastNewTrace!.getTime() > this._options.traceTimeoutSeconds * 1000 ||
+      Date.now() - this._lastNewTrace!.getTime() > this._traceTimeoutSeconds * 1000 ||
       this._currentSpanContext === undefined ||
       this._endSpanInNextIteration
     ) {
@@ -309,6 +341,7 @@ export class DevQueueConsumer {
       data: {
         lockedAt: new Date(),
         lockedById: backgroundTask.id,
+        status: "EXECUTING",
       },
       include: {
         attempts: {
@@ -365,6 +398,7 @@ export class DevQueueConsumer {
         backgroundWorkerTaskId: backgroundTask.id,
         status: "EXECUTING" as const,
         queueId: queue.id,
+        runtimeEnvironmentId: this.env.id,
       },
     });
 
@@ -441,6 +475,11 @@ export class DevQueueConsumer {
         },
       });
 
+      logger.debug("Saving the in progress attempt", {
+        taskRunAttempt: taskRunAttempt.id,
+        messageId: message.messageId,
+      });
+
       this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId);
     } catch (e) {
       if (e instanceof Error) {

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -31,9 +31,11 @@ const MessageBody = z.discriminatedUnion("type", [
   z.object({
     type: z.literal("RESUME"),
     completedAttemptIds: z.string().array(),
+    resumableAttemptId: z.string(),
   }),
   z.object({
     type: z.literal("RESUME_AFTER_DURATION"),
+    resumableAttemptId: z.string(),
   }),
 ]);
 
@@ -42,6 +44,8 @@ type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTas
 export type SharedQueueConsumerOptions = {
   maximumItemsPerTrace?: number;
   traceTimeoutSeconds?: number;
+  nextTickInterval?: number;
+  interval?: number;
 };
 
 export class SharedQueueConsumer {
@@ -66,6 +70,8 @@ export class SharedQueueConsumer {
     this._options = {
       maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace
       traceTimeoutSeconds: options.traceTimeoutSeconds ?? 60, // 60 seconds
+      nextTickInterval: options.nextTickInterval ?? 1000, // 1 second
+      interval: options.interval ?? 100, // 100ms
     };
   }
 
@@ -233,7 +239,7 @@ export class SharedQueueConsumer {
     const message = await marqs?.dequeueMessageInSharedQueue();
 
     if (!message) {
-      setTimeout(() => this.#doWork(), 1000);
+      setTimeout(() => this.#doWork(), this._options.nextTickInterval);
       return;
     }
 
@@ -257,7 +263,7 @@ export class SharedQueueConsumer {
         envId,
       });
       await marqs?.acknowledgeMessage(message.messageId);
-      setTimeout(() => this.#doWork(), 100);
+      setTimeout(() => this.#doWork(), this._options.interval);
       return;
     }
 
@@ -272,7 +278,7 @@ export class SharedQueueConsumer {
 
       await marqs?.acknowledgeMessage(message.messageId);
 
-      setTimeout(() => this.#doWork(), 100);
+      setTimeout(() => this.#doWork(), this._options.interval);
       return;
     }
 
@@ -290,7 +296,19 @@ export class SharedQueueConsumer {
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
+          return;
+        }
+
+        if (existingTaskRun.status !== "PENDING") {
+          logger.debug("Task run is not pending, aborting", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+            taskRun: existingTaskRun.id,
+            status: existingTaskRun.status,
+          });
+          await marqs?.acknowledgeMessage(message.messageId);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -321,7 +339,7 @@ export class SharedQueueConsumer {
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -332,7 +350,7 @@ export class SharedQueueConsumer {
             deployment: deployment.id,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -351,7 +369,7 @@ export class SharedQueueConsumer {
 
           await marqs?.acknowledgeMessage(message.messageId);
 
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -383,7 +401,7 @@ export class SharedQueueConsumer {
 
           await marqs?.acknowledgeMessage(message.messageId);
 
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -398,7 +416,7 @@ export class SharedQueueConsumer {
 
         if (!queue) {
           await marqs?.nackMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 1000);
+          setTimeout(() => this.#doWork(), this._options.nextTickInterval);
           return;
         }
 
@@ -417,6 +435,7 @@ export class SharedQueueConsumer {
             backgroundWorkerTaskId: backgroundTask.id,
             status: "PENDING" as const,
             queueId: queue.id,
+            runtimeEnvironmentId: environment.id,
           },
         });
 
@@ -462,7 +481,7 @@ export class SharedQueueConsumer {
           // Finally we need to nack the message so it can be retried
           await marqs?.nackMessage(message.messageId);
         } finally {
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
         }
         break;
       }
@@ -474,41 +493,47 @@ export class SharedQueueConsumer {
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
-        const resumableRun = await prisma.taskRun.findFirst({
+        const resumableRun = await prisma.taskRun.findUnique({
           where: {
             id: message.messageId,
           },
+        });
+
+        if (!resumableRun) {
+          logger.error("Resumable run not found", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+          });
+          await marqs?.acknowledgeMessage(message.messageId);
+          setTimeout(() => this.#doWork(), this._options.interval);
+          return;
+        }
+
+        const resumableAttempt = await prisma.taskRunAttempt.findUnique({
+          where: {
+            id: messageBody.data.resumableAttemptId,
+          },
           include: {
-            attempts: {
+            checkpoints: {
+              take: 1,
               orderBy: {
                 createdAt: "desc",
               },
-              take: 1,
-              include: {
-                checkpoints: {
-                  take: 1,
-                  orderBy: {
-                    createdAt: "desc",
-                  },
-                },
-              },
             },
           },
         });
 
-        const resumableAttempt = resumableRun?.attempts[0];
-
         if (!resumableAttempt) {
           logger.error("Resumable attempt not found", {
             queueMessage: message.data,
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -523,7 +548,7 @@ export class SharedQueueConsumer {
 
         if (!queue) {
           await marqs?.nackMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 1000);
+          setTimeout(() => this.#doWork(), this._options.nextTickInterval);
           return;
         }
 
@@ -543,7 +568,7 @@ export class SharedQueueConsumer {
               resumableAttemptId: resumableAttempt.id,
             });
             await marqs?.acknowledgeMessage(message.messageId);
-            setTimeout(() => this.#doWork(), 100);
+            setTimeout(() => this.#doWork(), this._options.interval);
             return;
           }
 
@@ -553,6 +578,13 @@ export class SharedQueueConsumer {
             },
             data: {
               status: "EXECUTING",
+              taskRun: {
+                update: {
+                  data: {
+                    status: "EXECUTING",
+                  },
+                },
+              },
             },
           });
 
@@ -565,7 +597,7 @@ export class SharedQueueConsumer {
             reason: latestCheckpoint.reason ?? undefined,
           });
 
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -593,7 +625,7 @@ export class SharedQueueConsumer {
               messageId: message.messageId,
             });
             await marqs?.acknowledgeMessage(message.messageId);
-            setTimeout(() => this.#doWork(), 100);
+            setTimeout(() => this.#doWork(), this._options.interval);
             return;
           }
 
@@ -601,7 +633,7 @@ export class SharedQueueConsumer {
 
           if (!completion) {
             await marqs?.acknowledgeMessage(message.messageId);
-            setTimeout(() => this.#doWork(), 100);
+            setTimeout(() => this.#doWork(), this._options.interval);
             return;
           }
 
@@ -614,7 +646,7 @@ export class SharedQueueConsumer {
 
           if (!executionPayload) {
             await marqs?.acknowledgeMessage(message.messageId);
-            setTimeout(() => this.#doWork(), 100);
+            setTimeout(() => this.#doWork(), this._options.interval);
             return;
           }
 
@@ -641,43 +673,34 @@ export class SharedQueueConsumer {
           // Finally we need to nack the message so it can be retried
           await marqs?.nackMessage(message.messageId);
         } finally {
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
         }
         break;
       }
       // Resume after duration-based wait
       case "RESUME_AFTER_DURATION": {
-        const resumableRun = await prisma.taskRun.findFirst({
+        const resumableAttempt = await prisma.taskRunAttempt.findUnique({
           where: {
-            id: message.messageId,
+            id: messageBody.data.resumableAttemptId,
           },
           include: {
-            attempts: {
+            checkpoints: {
+              take: 1,
               orderBy: {
                 createdAt: "desc",
               },
-              take: 1,
-              include: {
-                checkpoints: {
-                  take: 1,
-                  orderBy: {
-                    createdAt: "desc",
-                  },
-                },
-              },
             },
+            taskRun: true,
           },
         });
 
-        const resumableAttempt = resumableRun?.attempts[0];
-
         if (!resumableAttempt) {
           logger.error("Resumable attempt not found", {
             queueMessage: message.data,
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -687,7 +710,7 @@ export class SharedQueueConsumer {
             messageId: message.messageId,
           });
           await marqs?.acknowledgeMessage(message.messageId);
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
           return;
         }
 
@@ -702,7 +725,7 @@ export class SharedQueueConsumer {
               resumableAttemptId: resumableAttempt.id,
             });
             await marqs?.acknowledgeMessage(message.messageId);
-            setTimeout(() => this.#doWork(), 100);
+            setTimeout(() => this.#doWork(), this._options.interval);
             return;
           }
 
@@ -736,7 +759,7 @@ export class SharedQueueConsumer {
           // Finally we need to nack the message so it can be retried
           await marqs?.nackMessage(message.messageId);
         } finally {
-          setTimeout(() => this.#doWork(), 100);
+          setTimeout(() => this.#doWork(), this._options.interval);
         }
         break;
       }
@@ -811,14 +834,14 @@ class SharedQueueTasks {
       include: {
         backgroundWorker: true,
         backgroundWorkerTask: true,
+        runtimeEnvironment: {
+          include: {
+            organization: true,
+            project: true,
+          },
+        },
         taskRun: {
           include: {
-            runtimeEnvironment: {
-              include: {
-                organization: true,
-                project: true,
-              },
-            },
             tags: true,
             batchItem: {
               include: {
@@ -836,6 +859,30 @@ class SharedQueueTasks {
       return;
     }
 
+    if (attempt.status === "CANCELED") {
+      return;
+    }
+
+    if (attempt.status === "FAILED") {
+      return;
+    }
+
+    if (attempt.status === "COMPLETED") {
+      return;
+    }
+
+    if (attempt.taskRun.status === "CANCELED") {
+      return;
+    }
+
+    if (attempt.taskRun.status === "COMPLETED_SUCCESSFULLY") {
+      return;
+    }
+
+    if (attempt.taskRun.status === "COMPLETED_WITH_ERRORS") {
+      return;
+    }
+
     if (setToExecuting) {
       await prisma.taskRunAttempt.update({
         where: {
@@ -843,6 +890,13 @@ class SharedQueueTasks {
         },
         data: {
           status: "EXECUTING",
+          taskRun: {
+            update: {
+              data: {
+                status: "EXECUTING",
+              },
+            },
+          },
         },
       });
     }
@@ -877,20 +931,20 @@ class SharedQueueTasks {
         name: queue.name,
       },
       environment: {
-        id: taskRun.runtimeEnvironment.id,
-        slug: taskRun.runtimeEnvironment.slug,
-        type: taskRun.runtimeEnvironment.type,
+        id: attempt.runtimeEnvironment.id,
+        slug: attempt.runtimeEnvironment.slug,
+        type: attempt.runtimeEnvironment.type,
       },
       organization: {
-        id: taskRun.runtimeEnvironment.organization.id,
-        slug: taskRun.runtimeEnvironment.organization.slug,
-        name: taskRun.runtimeEnvironment.organization.title,
+        id: attempt.runtimeEnvironment.organization.id,
+        slug: attempt.runtimeEnvironment.organization.slug,
+        name: attempt.runtimeEnvironment.organization.title,
       },
       project: {
-        id: taskRun.runtimeEnvironment.project.id,
-        ref: taskRun.runtimeEnvironment.project.externalRef,
-        slug: taskRun.runtimeEnvironment.project.slug,
-        name: taskRun.runtimeEnvironment.project.name,
+        id: attempt.runtimeEnvironment.project.id,
+        ref: attempt.runtimeEnvironment.project.externalRef,
+        slug: attempt.runtimeEnvironment.project.slug,
+        name: attempt.runtimeEnvironment.project.name,
       },
       batch: taskRun.batchItem?.batchTaskRun
         ? { id: taskRun.batchItem.batchTaskRun.friendlyId }
@@ -904,8 +958,8 @@ class SharedQueueTasks {
 
     const environmentRepository = new EnvironmentVariablesRepository();
     const variables = await environmentRepository.getEnvironmentVariables(
-      attempt.taskRun.runtimeEnvironment.projectId,
-      attempt.taskRun.runtimeEnvironmentId
+      attempt.runtimeEnvironment.projectId,
+      attempt.runtimeEnvironmentId
     );
 
     const payload: ProdTaskRunExecutionPayload = {

apps/webapp/app/v3/services/cancelAttempt.server.ts

@@ -5,6 +5,7 @@ import { BaseService } from "./baseService.server";
 import { logger } from "~/services/logger.server";
 
 import { PrismaClientOrTransaction, prisma } from "~/db.server";
+import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
 
 export class CancelAttemptService extends BaseService {
   public async call(
@@ -49,6 +50,14 @@ export class CancelAttemptService extends BaseService {
         },
         data: {
           status: "CANCELED",
+          completedAt: cancelledAt,
+          taskRun: {
+            update: {
+              data: {
+                status: "INTERRUPTED",
+              },
+            },
+          },
         },
       });
 
@@ -65,6 +74,10 @@ export class CancelAttemptService extends BaseService {
           return eventRepository.cancelEvent(event, cancelledAt, reason);
         })
       );
+
+      if (environment?.type !== "DEVELOPMENT") {
+        await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma);
+      }
     });
   }
 }
@@ -78,14 +91,10 @@ async function getAuthenticatedEnvironmentFromAttempt(
       friendlyId,
     },
     include: {
-      taskRun: {
+      runtimeEnvironment: {
         include: {
-          runtimeEnvironment: {
-            include: {
-              organization: true,
-              project: true,
-            },
-          },
+          organization: true,
+          project: true,
         },
       },
     },
@@ -95,5 +104,5 @@ async function getAuthenticatedEnvironmentFromAttempt(
     return;
   }
 
-  return taskRunAttempt?.taskRun.runtimeEnvironment;
+  return taskRunAttempt?.runtimeEnvironment;
 }

apps/webapp/app/v3/services/cancelTaskRun.server.ts

@@ -0,0 +1,129 @@
+import { TaskRun, TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database";
+import { eventRepository } from "../eventRepository.server";
+import { marqs } from "../marqs.server";
+import { devPubSub } from "../marqs/devPubSub.server";
+import { BaseService } from "./baseService.server";
+import { socketIo } from "../handleSocketIo.server";
+import { assertUnreachable } from "../utils/asserts.server";
+import { CancelAttemptService } from "./cancelAttempt.server";
+import { logger } from "~/services/logger.server";
+
+const CANCELLABLE_STATUSES: Array<TaskRunStatus> = [
+  "PENDING",
+  "EXECUTING",
+  "PAUSED",
+  "WAITING_TO_RESUME",
+  "PAUSED",
+  "RETRYING_AFTER_FAILURE",
+];
+
+const CANCELLABLE_ATTEMPT_STATUSES: Array<TaskRunAttemptStatus> = [
+  "EXECUTING",
+  "PAUSED",
+  "PENDING",
+];
+
+export class CancelTaskRunService extends BaseService {
+  public async call(taskRun: TaskRun) {
+    // Make sure the task run is in a cancellable state
+    if (!CANCELLABLE_STATUSES.includes(taskRun.status)) {
+      return;
+    }
+
+    // Remove the task run from the queue if it's there for some reason
+    await marqs?.acknowledgeMessage(taskRun.id);
+
+    // Set the task run status to cancelled
+    const cancelledTaskRun = await this._prisma.taskRun.update({
+      where: {
+        id: taskRun.id,
+      },
+      data: {
+        status: "CANCELED",
+      },
+      include: {
+        attempts: {
+          where: {
+            status: {
+              in: CANCELLABLE_ATTEMPT_STATUSES,
+            },
+          },
+          include: {
+            backgroundWorker: true,
+            runtimeEnvironment: true,
+          },
+        },
+        dependency: true,
+        runtimeEnvironment: true,
+      },
+    });
+
+    const inProgressEvents = await eventRepository.queryIncompleteEvents({
+      runId: taskRun.friendlyId,
+    });
+
+    logger.debug("Cancelling in-progress events", {
+      inProgressEvents: inProgressEvents.map((event) => event.id),
+    });
+
+    await Promise.all(
+      inProgressEvents.map((event) => {
+        return eventRepository.cancelEvent(event, new Date(), "Task run was cancelled by user");
+      })
+    );
+
+    // Cancel any in progress attempts
+    for (const attempt of cancelledTaskRun.attempts) {
+      if (attempt.runtimeEnvironment.type === "DEVELOPMENT") {
+        // Signal the task run attempt to stop
+        await devPubSub.publish(
+          `backgroundWorker:${attempt.backgroundWorkerId}:${attempt.id}`,
+          "CANCEL_ATTEMPT",
+          {
+            attemptId: attempt.friendlyId,
+            backgroundWorkerId: attempt.backgroundWorker.friendlyId,
+            taskRunId: cancelledTaskRun.friendlyId,
+          }
+        );
+      } else {
+        switch (attempt.status) {
+          case "EXECUTING": {
+            // We need to send a cancel message to the coordinator
+            socketIo.coordinatorNamespace.emit("REQUEST_ATTEMPT_CANCELLATION", {
+              version: "v1",
+              attemptId: attempt.id,
+            });
+
+            break;
+          }
+          case "PENDING":
+          case "PAUSED": {
+            logger.debug("Cancelling pending or paused attempt", {
+              attempt,
+            });
+
+            const service = new CancelAttemptService();
+
+            await service.call(
+              attempt.friendlyId,
+              taskRun.id,
+              new Date(),
+              "Task run was cancelled by user"
+            );
+
+            break;
+          }
+          case "CANCELED":
+          case "COMPLETED":
+          case "FAILED": {
+            // Do nothing
+            break;
+          }
+          default: {
+            assertUnreachable(attempt.status);
+          }
+        }
+      }
+    }
+  }
+}

apps/webapp/app/v3/services/completeAttempt.server.ts

@@ -1,72 +1,131 @@
+import { Attributes } from "@opentelemetry/api";
 import {
   RetryOptions,
   TaskRunContext,
   TaskRunExecution,
   TaskRunExecutionResult,
+  TaskRunFailedExecutionResult,
+  TaskRunSuccessfulExecutionResult,
   defaultRetryOptions,
   flattenAttributes,
 } from "@trigger.dev/core/v3";
+import { PrismaClientOrTransaction } from "~/db.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
+import { safeJsonParse } from "~/utils/json";
 import { eventRepository } from "../eventRepository.server";
 import { marqs } from "../marqs.server";
 import { BaseService } from "./baseService.server";
-import { Attributes } from "@opentelemetry/api";
-import { logger } from "~/services/logger.server";
+import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
+import { CancelAttemptService } from "./cancelAttempt.server";
+
+type FoundAttempt = Awaited<ReturnType<typeof findAttempt>>;
 
 export class CompleteAttemptService extends BaseService {
   public async call(
     completion: TaskRunExecutionResult,
     execution: TaskRunExecution,
     env?: AuthenticatedEnvironment
-  ): Promise<"ACKNOWLEDGED" | "RETRIED" | "FAILED"> {
-    const taskRunAttempt = completion.ok
-      ? await this._prisma.taskRunAttempt.update({
-          where: { friendlyId: completion.id },
-          data: {
-            status: "COMPLETED",
-            completedAt: new Date(),
-            output: completion.output,
-            outputType: completion.outputType,
-          },
-          include: {
-            taskRun: {
-              include: {
-                batchItem: true,
-                dependency: {
-                  include: {
-                    dependentAttempt: true,
-                    dependentBatchRun: true,
-                  },
-                },
-              },
-            },
-            backgroundWorkerTask: true,
-          },
-        })
-      : await this._prisma.taskRunAttempt.update({
-          where: { friendlyId: completion.id },
-          data: {
-            status: "FAILED",
-            completedAt: new Date(),
-            error: completion.error,
-          },
-          include: {
-            taskRun: {
-              include: {
-                batchItem: true,
-                dependency: {
-                  include: {
-                    dependentAttempt: true,
-                    dependentBatchRun: true,
-                  },
-                },
-              },
-            },
-            backgroundWorkerTask: true,
-          },
-        });
+  ) {
+    const taskRunAttempt = await findAttempt(this._prisma, completion.id);
 
-    if (!completion.ok && completion.retry !== undefined) {
+    if (!taskRunAttempt) {
+      logger.error("[CompleteAttemptService] Task run attempt not found", { id: completion.id });
+
+      // Update the task run to be failed
+      await this._prisma.taskRun.update({
+        where: {
+          friendlyId: execution.run.id,
+        },
+        data: {
+          status: "SYSTEM_FAILURE",
+        },
+      });
+
+      return "FAILED";
+    }
+
+    if (completion.ok) {
+      return await this.#completeAttemptSuccessfully(completion, taskRunAttempt, env);
+    } else {
+      return await this.#completeAttemptFailed(completion, execution, taskRunAttempt, env);
+    }
+  }
+
+  async #completeAttemptSuccessfully(
+    completion: TaskRunSuccessfulExecutionResult,
+    taskRunAttempt: NonNullable<FoundAttempt>,
+    env?: AuthenticatedEnvironment
+  ) {
+    await this._prisma.taskRunAttempt.update({
+      where: { friendlyId: completion.id },
+      data: {
+        status: "COMPLETED",
+        completedAt: new Date(),
+        output: completion.output,
+        outputType: completion.outputType,
+        taskRun: {
+          update: {
+            data: {
+              status: "COMPLETED_SUCCESSFULLY",
+            },
+          },
+        },
+      },
+    });
+
+    logger.debug("Completed attempt successfully, ACKing message", taskRunAttempt);
+
+    await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId);
+
+    // Now we need to "complete" the task run event/span
+    await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
+      endTime: new Date(),
+      attributes: {
+        isError: false,
+        output: completion.output ? (safeJsonParse(completion.output) as Attributes) : undefined,
+      },
+    });
+
+    if (!env || env.type !== "DEVELOPMENT") {
+      await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma);
+    }
+
+    return "ACKNOWLEDGED";
+  }
+
+  async #completeAttemptFailed(
+    completion: TaskRunFailedExecutionResult,
+    execution: TaskRunExecution,
+    taskRunAttempt: NonNullable<FoundAttempt>,
+    env?: AuthenticatedEnvironment
+  ) {
+    if (
+      completion.error.type === "INTERNAL_ERROR" &&
+      completion.error.code === "TASK_RUN_CANCELLED"
+    ) {
+      // We need to cancel the task run instead of fail it
+      const cancelService = new CancelAttemptService();
+
+      return await cancelService.call(
+        taskRunAttempt.friendlyId,
+        taskRunAttempt.taskRunId,
+        new Date(),
+        "Cancelled by user",
+        env
+      );
+    }
+
+    await this._prisma.taskRunAttempt.update({
+      where: { friendlyId: completion.id },
+      data: {
+        status: "FAILED",
+        completedAt: new Date(),
+        error: completion.error,
+      },
+    });
+
+    if (completion.retry !== undefined) {
       const retryConfig = taskRunAttempt.backgroundWorkerTask.retryConfig
         ? {
             ...defaultRetryOptions,
@@ -111,6 +170,15 @@ export class CompleteAttemptService extends BaseService {
 
       logger.debug("Retrying", { taskRun: taskRunAttempt.taskRun.friendlyId });
 
+      await this._prisma.taskRun.update({
+        where: {
+          id: taskRunAttempt.taskRunId,
+        },
+        data: {
+          status: "RETRYING_AFTER_FAILURE",
+        },
+      });
+
       if (environment.type === "DEVELOPMENT") {
         // This is already an EXECUTE message so we can just NACK
         await marqs?.nackMessage(taskRunAttempt.taskRunId, completion.retry.timestamp);
@@ -127,179 +195,31 @@ export class CompleteAttemptService extends BaseService {
       }
 
       return "RETRIED";
-    }
-    // Attempt succeeded or this was the last retry
-    else {
+    } else {
+      // No more retries, we need to fail the task run
       logger.debug("Completed attempt, ACKing message", taskRunAttempt);
 
       await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId);
 
       // Now we need to "complete" the task run event/span
-      if (completion.ok) {
-        await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
-          endTime: new Date(),
-          attributes: {
-            isError: false,
-            output: completion.output ? (JSON.parse(completion.output) as Attributes) : undefined,
-          },
-        });
-      } else {
-        await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
-          endTime: new Date(),
-          attributes: {
-            isError: true,
-          },
-        });
-      }
+      await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
+        endTime: new Date(),
+        attributes: {
+          isError: true,
+        },
+      });
 
-      const { batchItem, dependency } = taskRunAttempt.taskRun;
+      await this._prisma.taskRun.update({
+        where: {
+          id: taskRunAttempt.taskRunId,
+        },
+        data: {
+          status: "COMPLETED_WITH_ERRORS",
+        },
+      });
 
-      // This run is part of a batch so we should update its status
-      if (batchItem) {
-        logger.debug("Completing attempt with batch item", { batchItem });
-
-        await this._prisma.batchTaskRunItem.update({
-          where: {
-            id: batchItem.id,
-          },
-          data: {
-            status: completion.ok ? "COMPLETED" : "FAILED",
-          },
-        });
-
-        const finalizedBatchRun = await this._prisma.batchTaskRun.findFirst({
-          where: {
-            id: batchItem.batchTaskRunId,
-            dependentTaskAttemptId: {
-              not: null,
-            },
-            items: {
-              every: {
-                status: {
-                  not: "PENDING",
-                },
-              },
-            },
-          },
-          include: {
-            dependentTaskAttempt: {
-              include: {
-                taskRun: true,
-              },
-            },
-            items: {
-              include: {
-                taskRun: {
-                  include: {
-                    attempts: {
-                      orderBy: {
-                        completedAt: "desc",
-                      },
-                      take: 1,
-                      select: {
-                        id: true,
-                      },
-                    },
-                  },
-                },
-              },
-            },
-          },
-        });
-
-        // This batch has a dependent attempt and just finalized, we should resume that attempt
-        if (finalizedBatchRun && finalizedBatchRun.dependentTaskAttempt) {
-          const environment =
-            env ?? (await this.#getEnvironment(taskRunAttempt.taskRun.runtimeEnvironmentId));
-
-          if (!environment) {
-            logger.error("Environment not found", {
-              attemptId: taskRunAttempt.id,
-              envId: taskRunAttempt.taskRun.runtimeEnvironmentId,
-            });
-            return "FAILED";
-          }
-
-          if (environment.type === "DEVELOPMENT") {
-            return "ACKNOWLEDGED";
-          }
-
-          const dependentRun = finalizedBatchRun.dependentTaskAttempt.taskRun;
-
-          if (finalizedBatchRun.dependentTaskAttempt.status === "PAUSED") {
-            await marqs?.enqueueMessage(
-              environment,
-              dependentRun.queue,
-              dependentRun.id,
-              {
-                type: "RESUME",
-                completedAttemptIds: [taskRunAttempt.id],
-              },
-              dependentRun.concurrencyKey ?? undefined
-            );
-          } else {
-            await marqs?.replaceMessage(dependentRun.id, {
-              type: "RESUME",
-              completedAttemptIds: finalizedBatchRun.items.map(
-                (item) => item.taskRun.attempts[0]?.id
-              ),
-            });
-          }
-        }
-      }
-
-      if (dependency) {
-        logger.debug("Completing attempt with dependency", { dependency });
-
-        const environment =
-          env ?? (await this.#getEnvironment(taskRunAttempt.taskRun.runtimeEnvironmentId));
-
-        if (!environment) {
-          logger.error("Environment not found", {
-            attemptId: taskRunAttempt.id,
-            envId: taskRunAttempt.taskRun.runtimeEnvironmentId,
-          });
-          return "FAILED";
-        }
-
-        if (environment.type === "DEVELOPMENT") {
-          return "ACKNOWLEDGED";
-        }
-
-        if (dependency.dependentAttempt) {
-          const dependentRun = await this._prisma.taskRun.findFirst({
-            where: {
-              id: dependency.dependentAttempt.taskRunId,
-            },
-          });
-
-          if (!dependentRun) {
-            logger.error("Dependent task run does not exist", {
-              attemptId: taskRunAttempt.id,
-              envId: taskRunAttempt.taskRun.runtimeEnvironmentId,
-              taskRunId: dependency.taskRunId,
-            });
-            return "FAILED";
-          }
-
-          if (dependency.dependentAttempt.status === "PAUSED") {
-            await marqs?.enqueueMessage(
-              environment,
-              dependentRun.queue,
-              dependentRun.id,
-              {
-                type: "RESUME",
-                completedAttemptIds: [taskRunAttempt.id],
-              },
-              dependentRun.concurrencyKey ?? undefined
-            );
-          } else {
-            await marqs?.replaceMessage(dependentRun.id, {
-              type: "RESUME",
-              completedAttemptIds: [taskRunAttempt.id],
-            });
-          }
-        }
+      if (!env || env.type !== "DEVELOPMENT") {
+        await ResumeTaskRunDependenciesService.enqueue(taskRunAttempt.id, this._prisma);
       }
 
       return "ACKNOWLEDGED";
@@ -329,3 +249,13 @@ export class CompleteAttemptService extends BaseService {
     });
   }
 }
+
+async function findAttempt(prismaClient: PrismaClientOrTransaction, friendlyId: string) {
+  return prismaClient.taskRunAttempt.findUnique({
+    where: { friendlyId },
+    include: {
+      taskRun: true,
+      backgroundWorkerTask: true,
+    },
+  });
+}

apps/webapp/app/v3/services/createCheckpoint.server.ts

@@ -51,7 +51,7 @@ export class CreateCheckpointService {
       case "WAIT_FOR_DURATION": {
         await marqs?.replaceMessage(
           attempt.taskRunId,
-          { type: "RESUME_AFTER_DURATION" },
+          { type: "RESUME_AFTER_DURATION", resumableAttemptId: attempt.id },
           Date.now() + params.reason.ms
         );
         break;

apps/webapp/app/v3/services/resumeBatchRun.server.ts

@@ -0,0 +1,104 @@
+import { PrismaClientOrTransaction } from "~/db.server";
+import { workerQueue } from "~/services/worker.server";
+import { marqs } from "../marqs.server";
+import { BaseService } from "./baseService.server";
+
+export class ResumeBatchRunService extends BaseService {
+  public async call(batchRunId: string, sourceTaskAttemptId: string) {
+    const batchRun = await this._prisma.batchTaskRun.findFirst({
+      where: {
+        id: batchRunId,
+        dependentTaskAttemptId: {
+          not: null,
+        },
+        status: "PENDING",
+        items: {
+          every: {
+            taskRunAttemptId: {
+              not: null,
+            },
+          },
+        },
+      },
+      include: {
+        dependentTaskAttempt: {
+          include: {
+            runtimeEnvironment: {
+              include: {
+                project: true,
+                organization: true,
+              },
+            },
+            taskRun: true,
+          },
+        },
+        items: true,
+      },
+    });
+
+    if (!batchRun || !batchRun.dependentTaskAttempt) {
+      return;
+    }
+
+    await this._prisma.batchTaskRun.update({
+      where: {
+        id: batchRun.id,
+      },
+      data: {
+        status: "COMPLETED",
+      },
+    });
+
+    // We need to update the batchRun status so we don't resume it again
+
+    // This batch has a dependent attempt and just finalized, we should resume that attempt
+    const environment = batchRun.dependentTaskAttempt.runtimeEnvironment;
+
+    // If we are in development, we don't need to resume the dependent task (that will happen automatically)
+    if (environment.type === "DEVELOPMENT") {
+      return;
+    }
+
+    const dependentRun = batchRun.dependentTaskAttempt.taskRun;
+
+    if (batchRun.dependentTaskAttempt.status === "PAUSED") {
+      await marqs?.enqueueMessage(
+        environment,
+        dependentRun.queue,
+        dependentRun.id,
+        {
+          type: "RESUME",
+          completedAttemptIds: [sourceTaskAttemptId],
+          resumableAttemptId: batchRun.dependentTaskAttempt.id,
+        },
+        dependentRun.concurrencyKey ?? undefined
+      );
+    } else {
+      await marqs?.replaceMessage(dependentRun.id, {
+        type: "RESUME",
+        completedAttemptIds: batchRun.items.map((item) => item.taskRunAttemptId).filter(Boolean),
+        resumableAttemptId: batchRun.dependentTaskAttempt.id,
+      });
+    }
+  }
+
+  static async enqueue(
+    batchRunId: string,
+    sourceTaskAttemptId: string,
+    tx: PrismaClientOrTransaction,
+    runAt?: Date
+  ) {
+    return await workerQueue.enqueue(
+      "v3.resumeBatchRun",
+      {
+        batchRunId,
+        sourceTaskAttemptId,
+      },
+      {
+        tx,
+        runAt,
+        queueName: `resumeBatchRun-${batchRunId}`,
+      }
+    );
+  }
+}

apps/webapp/app/v3/services/resumeTaskDependency.server.ts

@@ -0,0 +1,78 @@
+import { PrismaClientOrTransaction } from "~/db.server";
+import { workerQueue } from "~/services/worker.server";
+import { marqs } from "../marqs.server";
+import { BaseService } from "./baseService.server";
+
+export class ResumeTaskDependencyService extends BaseService {
+  public async call(dependencyId: string, sourceTaskAttemptId: string) {
+    const dependency = await this._prisma.taskRunDependency.findUnique({
+      where: { id: dependencyId },
+      include: {
+        taskRun: {
+          include: {
+            runtimeEnvironment: {
+              include: {
+                project: true,
+                organization: true,
+              },
+            },
+          },
+        },
+        dependentAttempt: {
+          include: {
+            taskRun: true,
+          },
+        },
+      },
+    });
+
+    // Dependencies with a dependentBatchRun are handled already by the ResumeBatchRunService
+    if (!dependency || !dependency.dependentAttempt) {
+      return;
+    }
+
+    if (dependency.taskRun.runtimeEnvironment.type === "DEVELOPMENT") {
+      return;
+    }
+    const dependentRun = dependency.dependentAttempt.taskRun;
+
+    if (dependency.dependentAttempt.status === "PAUSED") {
+      await marqs?.enqueueMessage(
+        dependency.taskRun.runtimeEnvironment,
+        dependentRun.queue,
+        dependentRun.id,
+        {
+          type: "RESUME",
+          completedAttemptIds: [sourceTaskAttemptId],
+          resumableAttemptId: dependency.dependentAttempt.id,
+        },
+        dependentRun.concurrencyKey ?? undefined
+      );
+    } else {
+      await marqs?.replaceMessage(dependentRun.id, {
+        type: "RESUME",
+        completedAttemptIds: [sourceTaskAttemptId],
+        resumableAttemptId: dependency.dependentAttempt.id,
+      });
+    }
+  }
+
+  static async enqueue(
+    dependencyId: string,
+    sourceTaskAttemptId: string,
+    tx: PrismaClientOrTransaction,
+    runAt?: Date
+  ) {
+    return await workerQueue.enqueue(
+      "v3.resumeTaskDependency",
+      {
+        dependencyId,
+        sourceTaskAttemptId,
+      },
+      {
+        tx,
+        runAt,
+      }
+    );
+  }
+}

apps/webapp/app/v3/services/resumeTaskRunDependencies.server.ts

@@ -0,0 +1,86 @@
+import { BatchTaskRunItem, TaskRunAttempt, TaskRunDependency } from "@trigger.dev/database";
+import { $transaction, PrismaClientOrTransaction } from "~/db.server";
+import { workerQueue } from "~/services/worker.server";
+import { BaseService } from "./baseService.server";
+import { ResumeBatchRunService } from "./resumeBatchRun.server";
+import { ResumeTaskDependencyService } from "./resumeTaskDependency.server";
+
+export class ResumeTaskRunDependenciesService extends BaseService {
+  public async call(attemptId: string) {
+    const taskAttempt = await this._prisma.taskRunAttempt.findUnique({
+      where: { id: attemptId },
+      include: {
+        taskRun: {
+          include: {
+            runtimeEnvironment: true,
+            batchItem: true,
+            dependency: {
+              include: {
+                dependentAttempt: true,
+                dependentBatchRun: true,
+              },
+            },
+          },
+        },
+        backgroundWorkerTask: true,
+      },
+    });
+
+    if (!taskAttempt) {
+      return;
+    }
+
+    if (taskAttempt.taskRun.runtimeEnvironment.type === "DEVELOPMENT") {
+      return;
+    }
+
+    const { batchItem, dependency } = taskAttempt.taskRun;
+
+    if (!batchItem && !dependency) {
+      return;
+    }
+
+    if (batchItem) {
+      await this.#resumeBatchItem(batchItem, taskAttempt);
+      return;
+    }
+
+    if (dependency && dependency.dependentAttempt) {
+      await this.#resumeDependency(dependency, taskAttempt);
+    }
+  }
+
+  async #resumeBatchItem(batchItem: BatchTaskRunItem, taskAttempt: TaskRunAttempt) {
+    await $transaction(this._prisma, async (tx) => {
+      await tx.batchTaskRunItem.update({
+        where: {
+          id: batchItem.id,
+        },
+        data: {
+          status: "COMPLETED",
+          taskRunAttemptId: taskAttempt.id,
+        },
+      });
+
+      await ResumeBatchRunService.enqueue(batchItem.batchTaskRunId, taskAttempt.id, tx);
+    });
+  }
+
+  async #resumeDependency(dependency: TaskRunDependency, taskAttempt: TaskRunAttempt) {
+    await ResumeTaskDependencyService.enqueue(dependency.id, taskAttempt.id, this._prisma);
+  }
+
+  static async enqueue(attemptId: string, tx: PrismaClientOrTransaction, runAt?: Date) {
+    return await workerQueue.enqueue(
+      "v3.resumeTaskRunDependencies",
+      {
+        attemptId,
+      },
+      {
+        tx,
+        runAt,
+        jobKey: `resumeTaskRunDependencies:${attemptId}`,
+      }
+    );
+  }
+}

apps/webapp/app/v3/services/triggerTask.server.ts

@@ -100,6 +100,7 @@ export class TriggerTaskService extends BaseService {
 
             const taskRun = await tx.taskRun.create({
               data: {
+                status: "PENDING",
                 number: counter.lastNumber,
                 friendlyId: generateFriendlyId("run"),
                 runtimeEnvironmentId: environment.id,

apps/webapp/app/v3/sharedSocketConnection.ts

@@ -47,7 +47,10 @@ export class SharedSocketConnection {
       },
     });
 
-    this._sharedConsumer = new SharedQueueConsumer(this._sender);
+    this._sharedConsumer = new SharedQueueConsumer(this._sender, {
+      interval: 100,
+      nextTickInterval: 1000,
+    });
 
     socket.on("disconnect", this.#handleClose.bind(this));
     socket.on("error", this.#handleError.bind(this));

apps/webapp/app/v3/utils/asserts.server.ts

@@ -0,0 +1,3 @@
+export function assertUnreachable(x: never): never {
+  throw new Error("Didn't expect to get here");
+}

apps/webapp/app/v3/utils/zodPubSub.server.ts

@@ -0,0 +1,117 @@
+import { Logger } from "@trigger.dev/core-backend";
+import { ZodMessageCatalogSchema, ZodMessageHandler, ZodMessageSender } from "@trigger.dev/core/v3";
+import Redis, { RedisOptions } from "ioredis";
+import { z } from "zod";
+import { logger } from "~/services/logger.server";
+import { safeJsonParse } from "~/utils/json";
+
+export type ZodPubSubOptions<TMessageCatalog extends ZodMessageCatalogSchema> = {
+  redis: RedisOptions;
+  schema: TMessageCatalog;
+};
+
+export interface ZodSubscriber<TMessageCatalog extends ZodMessageCatalogSchema> {
+  on<K extends keyof TMessageCatalog>(
+    eventName: K,
+    listener: (payload: z.infer<TMessageCatalog[K]>) => Promise<void>
+  ): void;
+
+  stopListening(): Promise<void>;
+}
+
+class RedisZodSubscriber<TMessageCatalog extends ZodMessageCatalogSchema>
+  implements ZodSubscriber<TMessageCatalog>
+{
+  private _subscriber: Redis;
+  private _listeners: Map<string, (payload: unknown) => Promise<void>> = new Map();
+  private _messageHandler: ZodMessageHandler<TMessageCatalog>;
+
+  constructor(
+    private readonly _pattern: string,
+    private readonly _options: ZodPubSubOptions<TMessageCatalog>,
+    private readonly _logger: Logger
+  ) {
+    this._subscriber = new Redis(_options.redis);
+    this._messageHandler = new ZodMessageHandler({
+      schema: _options.schema,
+    });
+  }
+
+  async initialize() {
+    await this._subscriber.psubscribe(this._pattern);
+    this._subscriber.on("pmessage", this.#onMessage.bind(this));
+  }
+
+  public on<K extends keyof TMessageCatalog>(
+    eventName: K,
+    listener: (payload: z.infer<TMessageCatalog[K]>) => Promise<void>
+  ): void {
+    this._listeners.set(eventName as string, listener);
+  }
+
+  public async stopListening(): Promise<void> {
+    this._listeners.clear();
+    await this._subscriber.unsubscribe();
+  }
+
+  async #onMessage(pattern: string, channel: string, serializedMessage: string) {
+    if (pattern !== this._pattern) {
+      return;
+    }
+
+    const parsedMessage = safeJsonParse(serializedMessage);
+
+    if (!parsedMessage) {
+      return;
+    }
+
+    const message = this._messageHandler.parseMessage(parsedMessage);
+
+    if (typeof message.type !== "string") {
+      return;
+    }
+
+    const listener = this._listeners.get(message.type);
+
+    if (!listener) {
+      this._logger.debug(`No listener for message type: ${message.type}`, { parsedMessage });
+
+      return;
+    }
+
+    try {
+      await listener(message.payload);
+    } catch (error) {
+      this._logger.error("Error handling message", { error, message });
+    }
+  }
+}
+
+export class ZodPubSub<TMessageCatalog extends ZodMessageCatalogSchema> {
+  private _publisher: Redis;
+  private _logger = logger.child({ module: "ZodPubSub" });
+
+  constructor(private _options: ZodPubSubOptions<TMessageCatalog>) {
+    this._publisher = new Redis(_options.redis);
+  }
+
+  public async publish<K extends keyof TMessageCatalog>(
+    channel: string,
+    type: K,
+    payload: z.input<TMessageCatalog[K]>
+  ): Promise<void> {
+    try {
+      await this._publisher.publish(channel, JSON.stringify({ type, payload, version: "v1" }));
+    } catch (e) {
+      logger.error("Failed to publish message", { channel, type, payload, error: e });
+    }
+  }
+
+  public async subscribe(channel: string): Promise<ZodSubscriber<TMessageCatalog>> {
+    const subscriber = new RedisZodSubscriber(channel, this._options, this._logger);
+
+    await subscriber.initialize();
+
+    return subscriber;
+  }
+}

packages/cli-v3/src/workers/dev/backgroundWorker.ts

@@ -108,9 +108,23 @@ export class BackgroundWorkerCoordinator {
   }
 
   async handleMessage(id: string, message: BackgroundWorkerServerMessages) {
+    logger.debug(`Received message from worker ${id}`, { workerMessage: message });
+
     switch (message.type) {
       case "EXECUTE_RUNS": {
         await Promise.all(message.payloads.map((payload) => this.#executeTaskRun(id, payload)));
+        break;
+      }
+      case "CANCEL_ATTEMPT": {
+        // Need to cancel the attempt somehow here
+        const worker = this._backgroundWorkers.get(id);
+
+        if (!worker) {
+          logger.error(`Could not find worker ${id}`);
+          return;
+        }
+
+        await worker.cancelRun(message.taskRunId);
       }
     }
   }
@@ -158,7 +172,8 @@ export class BackgroundWorkerCoordinator {
 
     const resultText = !completion.ok
       ? completion.error.type === "INTERNAL_ERROR" &&
-        completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED
+        (completion.error.code === TaskRunErrorCodes.TASK_EXECUTION_ABORTED ||
+          completion.error.code === TaskRunErrorCodes.TASK_RUN_CANCELLED)
         ? chalk.yellow("cancelled")
         : chalk.red(`error${retryingText}`)
       : chalk.green("success");
@@ -214,6 +229,14 @@ class CleanupProcessError extends Error {
   }
 }
 
+class CancelledProcessError extends Error {
+  constructor() {
+    super("Cancelled");
+
+    this.name = "CancelledProcessError";
+  }
+}
+
 export type BackgroundWorkerParams = {
   env: Record<string, string>;
   dependencies?: Record<string, string>;
@@ -367,6 +390,16 @@ export class BackgroundWorker {
     return this._taskRunProcesses.get(payload.execution.run.id) as TaskRunProcess;
   }
 
+  async cancelRun(taskRunId: string) {
+    const taskRunProcess = this._taskRunProcesses.get(taskRunId);
+
+    if (!taskRunProcess) {
+      return;
+    }
+
+    await taskRunProcess.cancel();
+  }
+
   // We need to fork the process before we can execute any tasks
   async executeTaskRun(payload: TaskRunExecutionPayload): Promise<TaskRunExecutionResult> {
     try {
@@ -393,6 +426,18 @@ export class BackgroundWorker {
 
       return result;
     } catch (e) {
+      if (e instanceof CancelledProcessError) {
+        return {
+          id: payload.execution.attempt.id,
+          ok: false,
+          retry: undefined,
+          error: {
+            type: "INTERNAL_ERROR",
+            code: TaskRunErrorCodes.TASK_RUN_CANCELLED,
+          },
+        };
+      }
+
       if (e instanceof CleanupProcessError) {
         return {
           id: payload.execution.attempt.id,
@@ -464,6 +509,7 @@ class TaskRunProcess {
   private _attemptStatuses: Map<string, "PENDING" | "REJECTED" | "RESOLVED"> = new Map();
   private _currentExecution: TaskRunExecution | undefined;
   private _isBeingKilled: boolean = false;
+  private _isBeingCancelled: boolean = false;
   public onTaskHeartbeat: Evt<string> = new Evt();
   public onExit: Evt<number> = new Evt();
 
@@ -483,6 +529,12 @@ class TaskRunProcess {
     });
   }
 
+  async cancel() {
+    this._isBeingCancelled = true;
+
+    await this.cleanup(true);
+  }
+
   async initialize() {
     this._child = fork(this.path, {
       stdio: [/*stdin*/ "ignore", /*stdout*/ "pipe", /*stderr*/ "pipe", "ipc"],
@@ -617,7 +669,9 @@ class TaskRunProcess {
 
         const { rejecter } = attemptPromise;
 
-        if (this._isBeingKilled) {
+        if (this._isBeingCancelled) {
+          rejecter(new CancelledProcessError());
+        } else if (this._isBeingKilled) {
           rejecter(new CleanupProcessError());
         } else {
           rejecter(new UnexpectedExitError(code));

packages/cli-v3/src/workers/prod/backgroundWorker.ts

@@ -36,6 +36,14 @@ class CleanupProcessError extends Error {
   }
 }
 
+class CancelledProcessError extends Error {
+  constructor() {
+    super("Cancelled");
+
+    this.name = "CancelledProcessError";
+  }
+}
+
 type BackgroundWorkerParams = {
   env: Record<string, string>;
   projectConfig: Config;
@@ -58,7 +66,7 @@ export class ProdBackgroundWorker {
 
   public tasks: Array<TaskMetadataWithFilePath> = [];
 
-  _taskRunProcesses: Map<string, TaskRunProcess> = new Map();
+  _taskRunProcess: TaskRunProcess | undefined;
 
   private _closed: boolean = false;
 
@@ -77,9 +85,7 @@ export class ProdBackgroundWorker {
     this.onTaskHeartbeat.detach();
 
     // We need to close all the task run processes
-    for (const taskRunProcess of this._taskRunProcesses.values()) {
-      taskRunProcess.cleanup(true);
-    }
+    this._taskRunProcess?.cleanup(true);
 
     // Delete worker files
     this._onClose.post();
@@ -173,14 +179,11 @@ export class ProdBackgroundWorker {
     completion: TaskRunExecutionResult,
     execution: TaskRunExecution
   ) {
-    for (const taskRunProcess of this._taskRunProcesses.values()) {
-      taskRunProcess.taskRunCompletedNotification(completion, execution);
-    }
+    this._taskRunProcess?.taskRunCompletedNotification(completion, execution);
   }
+
   async waitCompletedNotification() {
-    for (const taskRunProcess of this._taskRunProcesses.values()) {
-      taskRunProcess.waitCompletedNotification();
-    }
+    this._taskRunProcess?.waitCompletedNotification();
   }
 
   async #initializeTaskRunProcess(payload: ProdTaskRunExecutionPayload): Promise<TaskRunProcess> {
@@ -189,47 +192,45 @@ export class ProdBackgroundWorker {
       payload.execution.worker.version
     );
 
-    if (!this._taskRunProcesses.has(payload.execution.run.id)) {
-      const taskRunProcess = new TaskRunProcess(
-        this.path,
-        {
-          ...this.params.env,
-          ...(payload.environment ?? {}),
-        },
-        metadata,
-        this.params
-      );
+    const taskRunProcess = new TaskRunProcess(
+      this.path,
+      {
+        ...this.params.env,
+        ...(payload.environment ?? {}),
+      },
+      metadata,
+      this.params
+    );
 
-      taskRunProcess.onExit.attach(() => {
-        this._taskRunProcesses.delete(payload.execution.run.id);
-      });
+    this._taskRunProcess = taskRunProcess;
 
-      taskRunProcess.onTaskHeartbeat.attach((id) => {
-        this.onTaskHeartbeat.post(id);
-      });
+    taskRunProcess.onExit.attach(() => {
+      this._taskRunProcess = undefined;
+    });
 
-      taskRunProcess.onWaitForBatch.attach((message) => {
-        this.onWaitForBatch.post(message);
-      });
+    taskRunProcess.onTaskHeartbeat.attach((id) => {
+      this.onTaskHeartbeat.post(id);
+    });
 
-      taskRunProcess.onWaitForDuration.attach((message) => {
-        this.onWaitForDuration.post(message);
-      });
+    taskRunProcess.onWaitForBatch.attach((message) => {
+      this.onWaitForBatch.post(message);
+    });
 
-      taskRunProcess.onWaitForTask.attach((message) => {
-        this.onWaitForTask.post(message);
-      });
+    taskRunProcess.onWaitForDuration.attach((message) => {
+      this.onWaitForDuration.post(message);
+    });
 
-      this.preCheckpointNotification.attach((message) => {
-        taskRunProcess.preCheckpointNotification.post(message);
-      });
+    taskRunProcess.onWaitForTask.attach((message) => {
+      this.onWaitForTask.post(message);
+    });
 
-      await taskRunProcess.initialize();
+    this.preCheckpointNotification.attach((message) => {
+      taskRunProcess.preCheckpointNotification.post(message);
+    });
 
-      this._taskRunProcesses.set(payload.execution.run.id, taskRunProcess);
-    }
+    await taskRunProcess.initialize();
 
-    return this._taskRunProcesses.get(payload.execution.run.id) as TaskRunProcess;
+    return taskRunProcess;
   }
 
   // We need to fork the process before we can execute any tasks
@@ -259,6 +260,18 @@ export class ProdBackgroundWorker {
 
       return result;
     } catch (e) {
+      if (e instanceof CancelledProcessError) {
+        return {
+          id: payload.execution.attempt.id,
+          ok: false,
+          retry: undefined,
+          error: {
+            type: "INTERNAL_ERROR",
+            code: TaskRunErrorCodes.TASK_RUN_CANCELLED,
+          },
+        };
+      }
+
       if (e instanceof CleanupProcessError) {
         return {
           id: payload.execution.attempt.id,
@@ -295,6 +308,10 @@ export class ProdBackgroundWorker {
     }
   }
 
+  async cancelAttempt(attemptId: string) {
+    await this._taskRunProcess?.cancel();
+  }
+
   async #correctError(
     error: TaskRunBuiltInError,
     execution: TaskRunExecution
@@ -320,6 +337,7 @@ class TaskRunProcess {
   private _attemptStatuses: Map<string, "PENDING" | "REJECTED" | "RESOLVED"> = new Map();
   private _currentExecution: TaskRunExecution | undefined;
   private _isBeingKilled: boolean = false;
+  private _isBeingCancelled: boolean = false;
 
   public onTaskHeartbeat: Evt<string> = new Evt();
   public onExit: Evt<number> = new Evt();
@@ -405,6 +423,12 @@ class TaskRunProcess {
     this._child.stderr?.on("data", this.#handleStdErr.bind(this));
   }
 
+  async cancel() {
+    this._isBeingCancelled = true;
+
+    await this.cleanup(true);
+  }
+
   async cleanup(kill: boolean = false) {
     if (kill && this._isBeingKilled) {
       return;
@@ -484,7 +508,9 @@ class TaskRunProcess {
 
         const { rejecter } = attemptPromise;
 
-        if (this._isBeingKilled) {
+        if (this._isBeingCancelled) {
+          rejecter(new CancelledProcessError());
+        } else if (this._isBeingKilled) {
           rejecter(new CleanupProcessError());
         } else {
           rejecter(new UnexpectedExitError(code));

packages/cli-v3/src/workers/prod/entry-point.ts

@@ -200,6 +200,16 @@ class ProdWorker {
 
           process.exit(0);
         },
+        REQUEST_ATTEMPT_CANCELLATION: async (message) => {
+          if (!this.executing) {
+            return;
+          }
+
+          await this.#backgroundWorker.cancelAttempt(message.attemptId);
+        },
+        REQUEST_EXIT: async () => {
+          process.exit(0);
+        },
       },
       onConnection: async (socket, handler, sender, logger) => {
         if (process.env.INDEX_TASKS === "true") {

packages/core-backend/src/logger.ts

@@ -37,6 +37,16 @@ export class Logger {
     this.#additionalFields = additionalFields ?? (() => ({}));
   }
 
+  child(fields: Record<string, unknown>) {
+    return new Logger(
+      this.#name,
+      logLevels[this.#level],
+      this.#filteredKeys,
+      this.#jsonReplacer,
+      () => ({ ...this.#additionalFields(), ...fields })
+    );
+  }
+
   // Return a new Logger instance with the same name and a new log level
   // but filter out the keys from the log messages (at any level)
   filter(...keys: string[]) {

packages/core/src/v3/schemas/common.ts

@@ -30,6 +30,7 @@ export const TaskRunErrorCodes = {
   TASK_EXECUTION_FAILED: "TASK_EXECUTION_FAILED",
   TASK_EXECUTION_ABORTED: "TASK_EXECUTION_ABORTED",
   TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE: "TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE",
+  TASK_RUN_CANCELLED: "TASK_RUN_CANCELLED",
 } as const;
 
 export const TaskRunInternalError = z.object({
@@ -41,6 +42,7 @@ export const TaskRunInternalError = z.object({
     "TASK_EXECUTION_FAILED",
     "TASK_EXECUTION_ABORTED",
     "TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE",
+    "TASK_RUN_CANCELLED",
   ]),
 });
 

packages/core/src/v3/schemas/messages.ts

@@ -32,6 +32,11 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [
     type: z.literal("EXECUTE_RUNS"),
     payloads: z.array(TaskRunExecutionPayload),
   }),
+  z.object({
+    type: z.literal("CANCEL_ATTEMPT"),
+    taskAttemptId: z.string(),
+    taskRunId: z.string(),
+  }),
   z.object({
     type: z.literal("SCHEDULE_ATTEMPT"),
     id: z.string(),

packages/core/src/v3/schemas/schemas.ts

@@ -222,6 +222,12 @@ export const PlatformToCoordinatorMessages = {
       attemptId: z.string(),
     }),
   },
+  REQUEST_ATTEMPT_CANCELLATION: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+      attemptId: z.string(),
+    }),
+  },
 };
 
 export const ClientToSharedQueueMessages = {
@@ -375,6 +381,17 @@ export const CoordinatorToProdWorkerMessages = {
       executionPayload: ProdTaskRunExecutionPayload,
     }),
   },
+  REQUEST_ATTEMPT_CANCELLATION: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+      attemptId: z.string(),
+    }),
+  },
+  REQUEST_EXIT: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+    }),
+  },
 };
 
 export const ProdWorkerSocketData = z.object({

packages/core/src/v3/zodMessageHandler.ts

@@ -17,7 +17,7 @@ export type ZodMessageHandlerOptions<TMessageCatalog extends ZodMessageCatalogSc
   messages?: ZodMessageHandlers<TMessageCatalog>;
 };
 
-type MessageFromSchema<
+export type MessageFromSchema<
   K extends keyof TMessageCatalog,
   TMessageCatalog extends ZodMessageCatalogSchema,
 > = {
@@ -25,11 +25,11 @@ type MessageFromSchema<
   payload: z.input<TMessageCatalog[K]>;
 };
 
-type MessageFromCatalog<TMessageCatalog extends ZodMessageCatalogSchema> = {
+export type MessageFromCatalog<TMessageCatalog extends ZodMessageCatalogSchema> = {
   [K in keyof TMessageCatalog]: MessageFromSchema<K, TMessageCatalog>;
 }[keyof TMessageCatalog];
 
-const messageSchema = z.object({
+export const ZodMessageSchema = z.object({
   version: z.literal("v1").default("v1"),
   type: z.string(),
   payload: z.unknown(),
@@ -68,7 +68,7 @@ export class ZodMessageHandler<TMessageCatalog extends ZodMessageCatalogSchema>
   }
 
   public parseMessage(message: unknown): MessageFromCatalog<TMessageCatalog> {
-    const parsedMessage = messageSchema.safeParse(message);
+    const parsedMessage = ZodMessageSchema.safeParse(message);
 
     if (!parsedMessage.success) {
       throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
@@ -161,6 +161,32 @@ export class ZodMessageSender<TMessageCatalog extends ZodMessageCatalogSchema> {
 
     await this.#sender({ type, payload, version: "v1" });
   }
+
+  public async forwardMessage(message: unknown) {
+    const parsedMessage = ZodMessageSchema.safeParse(message);
+
+    if (!parsedMessage.success) {
+      throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
+    }
+
+    const schema = this.#schema[parsedMessage.data.type];
+
+    if (!schema) {
+      throw new Error(`Unknown message type: ${parsedMessage.data.type}`);
+    }
+
+    const parsedPayload = schema.safeParse(parsedMessage.data.payload);
+
+    if (!parsedPayload.success) {
+      throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
+    }
+
+    await this.#sender({
+      type: parsedMessage.data.type,
+      payload: parsedPayload.data,
+      version: "v1",
+    });
+  }
 }
 
 export type MessageCatalogToSocketIoEvents<TCatalog extends ZodMessageCatalogSchema> = {

packages/database/prisma/migrations/20240312095501_add_status_to_task_runs/migration.sql

@@ -0,0 +1,5 @@
+-- CreateEnum
+CREATE TYPE "TaskRunStatus" AS ENUM ('PENDING', 'EXECUTING', 'WAITING_TO_RESUME', 'RETRYING_AFTER_FAILURE', 'PAUSED', 'CANCELED', 'COMPLETED_SUCCESSFULLY', 'COMPLETED_WITH_ERRORS');
+
+-- AlterTable
+ALTER TABLE "TaskRun" ADD COLUMN     "status" "TaskRunStatus" NOT NULL DEFAULT 'PENDING';

packages/database/prisma/migrations/20240312095826_add_interrupted_status_to_runs/migration.sql

@@ -0,0 +1,2 @@
+-- AlterEnum
+ALTER TYPE "TaskRunStatus" ADD VALUE 'INTERRUPTED';

packages/database/prisma/migrations/20240312105844_add_system_failure_status_to_task_run_status/migration.sql

@@ -0,0 +1,2 @@
+-- AlterEnum
+ALTER TYPE "TaskRunStatus" ADD VALUE 'SYSTEM_FAILURE';

packages/database/prisma/migrations/20240312125252_add_status_to_batch_run/migration.sql

@@ -0,0 +1,5 @@
+-- CreateEnum
+CREATE TYPE "BatchTaskRunStatus" AS ENUM ('PENDING', 'COMPLETED');
+
+-- AlterTable
+ALTER TABLE "BatchTaskRun" ADD COLUMN     "status" "BatchTaskRunStatus" NOT NULL DEFAULT 'PENDING';

packages/database/prisma/migrations/20240312131122_add_task_run_attempt_to_batch_run_items/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "BatchTaskRunItem" ADD COLUMN     "taskRunAttemptId" TEXT;
+
+-- AddForeignKey
+ALTER TABLE "BatchTaskRunItem" ADD CONSTRAINT "BatchTaskRunItem_taskRunAttemptId_fkey" FOREIGN KEY ("taskRunAttemptId") REFERENCES "TaskRunAttempt"("id") ON DELETE SET NULL ON UPDATE CASCADE;

packages/database/prisma/migrations/20240313110150_add_environment_to_task_run_attempts/migration.sql

@@ -0,0 +1,11 @@
+/*
+  Warnings:
+
+  - Added the required column `runtimeEnvironmentId` to the `TaskRunAttempt` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- AlterTable
+ALTER TABLE "TaskRunAttempt" ADD COLUMN     "runtimeEnvironmentId" TEXT NOT NULL;
+
+-- AddForeignKey
+ALTER TABLE "TaskRunAttempt" ADD CONSTRAINT "TaskRunAttempt_runtimeEnvironmentId_fkey" FOREIGN KEY ("runtimeEnvironmentId") REFERENCES "RuntimeEnvironment"("id") ON DELETE CASCADE ON UPDATE CASCADE;

packages/database/prisma/schema.prisma

@@ -393,6 +393,7 @@ model RuntimeEnvironment {
   checkpoints                Checkpoint[]
   workerDeployments          WorkerDeployment[]
   workerDeploymentPromotions WorkerDeploymentPromotion[]
+  taskRunAttempts            TaskRunAttempt[]
 
   @@unique([projectId, slug, orgMemberId])
   @@unique([projectId, shortcode])
@@ -1564,6 +1565,8 @@ model TaskRun {
   number     Int    @default(0)
   friendlyId String @unique
 
+  status TaskRunStatus @default(PENDING)
+
   idempotencyKey String
   taskIdentifier String
 
@@ -1606,6 +1609,38 @@ model TaskRun {
   @@unique([runtimeEnvironmentId, idempotencyKey])
 }
 
+enum TaskRunStatus {
+  /// Task is waiting to be executed by a worker
+  PENDING
+
+  /// Task is currently being executed by a worker
+  EXECUTING
+
+  /// Task has been paused by the system, and will be resumed by the system
+  WAITING_TO_RESUME
+
+  /// Task has failed and is waiting to be retried
+  RETRYING_AFTER_FAILURE
+
+  /// Task has been paused by the user, and can be resumed by the user
+  PAUSED
+
+  /// Task has been canceled by the user
+  CANCELED
+
+  /// Task was interrupted during execution, mostly this happens in development environments
+  INTERRUPTED
+
+  /// Task has been completed successfully
+  COMPLETED_SUCCESSFULLY
+
+  /// Task has been completed with errors
+  COMPLETED_WITH_ERRORS
+
+  /// Task has failed to complete, due to an error in the system
+  SYSTEM_FAILURE
+}
+
 model TaskRunDependency {
   id String @id @default(cuid())
 
@@ -1660,6 +1695,9 @@ model TaskRunAttempt {
   backgroundWorkerTask   BackgroundWorkerTask @relation(fields: [backgroundWorkerTaskId], references: [id], onDelete: Cascade, onUpdate: Cascade)
   backgroundWorkerTaskId String
 
+  runtimeEnvironment   RuntimeEnvironment @relation(fields: [runtimeEnvironmentId], references: [id], onDelete: Cascade, onUpdate: Cascade)
+  runtimeEnvironmentId String
+
   queue   TaskQueue @relation(fields: [queueId], references: [id], onDelete: Cascade, onUpdate: Cascade)
   queueId String
 
@@ -1678,7 +1716,8 @@ model TaskRunAttempt {
   taskRunDependency      TaskRunDependency? @relation("dependentAttempt")
   batchTaskRunDependency BatchTaskRun?
 
-  checkpoints Checkpoint[]
+  checkpoints       Checkpoint[]
+  batchTaskRunItems BatchTaskRunItem[]
 
   @@unique([taskRunId, number])
 }
@@ -1826,6 +1865,8 @@ model BatchTaskRun {
 
   friendlyId String @unique
 
+  status BatchTaskRunStatus @default(PENDING)
+
   idempotencyKey String
   taskIdentifier String
 
@@ -1844,6 +1885,11 @@ model BatchTaskRun {
   @@unique([runtimeEnvironmentId, idempotencyKey])
 }
 
+enum BatchTaskRunStatus {
+  PENDING
+  COMPLETED
+}
+
 model BatchTaskRunItem {
   id String @id @default(cuid())
 
@@ -1855,6 +1901,9 @@ model BatchTaskRunItem {
   taskRun   TaskRun @relation(fields: [taskRunId], references: [id], onDelete: Cascade, onUpdate: Cascade)
   taskRunId String  @unique
 
+  taskRunAttempt   TaskRunAttempt? @relation(fields: [taskRunAttemptId], references: [id], onDelete: SetNull, onUpdate: Cascade)
+  taskRunAttemptId String?
+
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
 

references/v3-catalog/src/trigger/longRunning.ts

@@ -13,3 +13,16 @@ export const longRunning = task({
     };
   },
 });
+
+export const longRunningParent = task({
+  id: "long-running-parent",
+  run: async (payload: { message: string }) => {
+    logger.info("Long running parent", { payload });
+
+    await longRunning.triggerAndWait({ payload: { message: "child" } });
+
+    return {
+      finished: new Date().toISOString(),
+    };
+  },
+});

references/v3-catalog/src/trigger/retries.ts

@@ -4,6 +4,9 @@ import { interceptor } from "./utils/interceptor";
 
 export const taskWithRetries = task({
   id: "task-with-retries",
+  retry: {
+    maxAttempts: 4,
+  },
   run: async (payload: any, { ctx }) => {
     const result = await retry.onThrow(
       async ({ attempt }) => {
@@ -43,6 +46,13 @@ export const taskWithRetries = task({
   },
 });
 
+export const taskThatErrors = task({
+  id: "task-that-errors",
+  run: async (payload: any, { ctx }) => {
+    throw new Error("failed");
+  },
+});
+
 export const taskWithFetchRetries = task({
   id: "task-with-fetch-retries",
   middleware: (payload: any, { next }) => {