v3: fix dependency checkpoint race (#1171)

* consolidate task statuses and utils

* check dependency completion when creating checkpoints

* add changeset

.changeset/silly-buses-obey.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+Add callback to checkpoint created message

apps/coordinator/src/index.ts

@@ -1162,13 +1162,7 @@ class TaskCoordinator {
             return;
           }
 
-          if (!checkpoint.docker || !willSimulate) {
-            socket.emit("REQUEST_EXIT", {
-              version: "v1",
-            });
-          }
-
-          this.#platformSocket?.send("CHECKPOINT_CREATED", {
+          const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
             version: "v1",
             attemptFriendlyId: message.attemptFriendlyId,
             docker: checkpoint.docker,
@@ -1179,6 +1173,17 @@ class TaskCoordinator {
               now: message.now,
             },
           });
+
+          if (ack?.keepRunAlive) {
+            logger.log("keeping run alive after duration checkpoint", { runId: socket.data.runId });
+            return;
+          }
+
+          if (!checkpoint.docker || !willSimulate) {
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+          }
         });
 
         socket.on("WAIT_FOR_TASK", async (message, callback) => {
@@ -1205,13 +1210,7 @@ class TaskCoordinator {
             return;
           }
 
-          if (!checkpoint.docker || !willSimulate) {
-            socket.emit("REQUEST_EXIT", {
-              version: "v1",
-            });
-          }
-
-          this.#platformSocket?.send("CHECKPOINT_CREATED", {
+          const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
             version: "v1",
             attemptFriendlyId: message.attemptFriendlyId,
             docker: checkpoint.docker,
@@ -1221,6 +1220,17 @@ class TaskCoordinator {
               friendlyId: message.friendlyId,
             },
           });
+
+          if (ack?.keepRunAlive) {
+            logger.log("keeping run alive after task checkpoint", { runId: socket.data.runId });
+            return;
+          }
+
+          if (!checkpoint.docker || !willSimulate) {
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+          }
         });
 
         socket.on("WAIT_FOR_BATCH", async (message, callback) => {
@@ -1247,13 +1257,7 @@ class TaskCoordinator {
             return;
           }
 
-          if (!checkpoint.docker || !willSimulate) {
-            socket.emit("REQUEST_EXIT", {
-              version: "v1",
-            });
-          }
-
-          this.#platformSocket?.send("CHECKPOINT_CREATED", {
+          const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
             version: "v1",
             attemptFriendlyId: message.attemptFriendlyId,
             docker: checkpoint.docker,
@@ -1264,6 +1268,17 @@ class TaskCoordinator {
               runFriendlyIds: message.runFriendlyIds,
             },
           });
+
+          if (ack?.keepRunAlive) {
+            logger.log("keeping run alive after batch checkpoint", { runId: socket.data.runId });
+            return;
+          }
+
+          if (!checkpoint.docker || !willSimulate) {
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+          }
         });
 
         socket.on("INDEX_TASKS", async (message, callback) => {

apps/webapp/app/presenters/v3/RunListPresenter.server.ts

@@ -4,8 +4,8 @@ import { Direction } from "~/components/runs/RunStatuses";
 import { FINISHED_STATUSES } from "~/components/runs/v3/TaskRunStatus";
 import { sqlDatabaseSchema } from "~/db.server";
 import { displayableEnvironment } from "~/models/runtimeEnvironment.server";
-import { CANCELLABLE_STATUSES } from "~/v3/services/cancelTaskRun.server";
 import { BasePresenter } from "./basePresenter.server";
+import { isCancellableRunStatus } from "~/v3/taskStatus";
 
 export type RunListOptions = {
   userId?: string;
@@ -291,7 +291,7 @@ export class RunListPresenter extends BasePresenter {
           taskIdentifier: run.taskIdentifier,
           spanId: run.spanId,
           isReplayable: true,
-          isCancellable: CANCELLABLE_STATUSES.includes(run.status),
+          isCancellable: isCancellableRunStatus(run.status),
           environment: displayableEnvironment(environment, userId),
           idempotencyKey: run.idempotencyKey ? run.idempotencyKey : undefined,
         };

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

@@ -138,8 +138,19 @@ function createCoordinatorNamespace(io: Server) {
         await sharedQueueTasks.taskRunHeartbeat(message.runId);
       },
       CHECKPOINT_CREATED: async (message) => {
-        const createCheckpoint = new CreateCheckpointService();
-        await createCheckpoint.call(message);
+        try {
+          const createCheckpoint = new CreateCheckpointService();
+          const result = await createCheckpoint.call(message);
+
+          return { keepRunAlive: result?.keepRunAlive ?? false };
+        } catch (error) {
+          logger.error("Error while creating checkpoint", {
+            rawMessage: message,
+            error: error instanceof Error ? error.message : error,
+          });
+
+          return { keepRunAlive: false };
+        }
       },
       CREATE_WORKER: async (message) => {
         try {

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

@@ -17,7 +17,6 @@ import {
   BackgroundWorkerTask,
   RuntimeEnvironment,
   TaskRun,
-  TaskRunAttemptStatus,
   TaskRunStatus,
 } from "@trigger.dev/database";
 import { z } from "zod";
@@ -43,6 +42,7 @@ import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { EnvironmentVariable } from "../environmentVariables/repository";
 import { machinePresetFromConfig } from "../machinePresets.server";
 import { env } from "~/env.server";
+import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
 
 const WithTraceContext = z.object({
   traceparent: z.string().optional(),
@@ -962,19 +962,7 @@ class SharedQueueTasks {
     }
 
     if (setToExecuting) {
-      const FINAL_RUN_STATUSES: TaskRunStatus[] = [
-        "CANCELED",
-        "COMPLETED_SUCCESSFULLY",
-        "COMPLETED_WITH_ERRORS",
-        "INTERRUPTED",
-        "SYSTEM_FAILURE",
-      ];
-      const FINAL_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["CANCELED", "COMPLETED", "FAILED"];
-
-      if (
-        FINAL_ATTEMPT_STATUSES.includes(attempt.status) ||
-        FINAL_RUN_STATUSES.includes(attempt.taskRun.status)
-      ) {
+      if (isFinalAttemptStatus(attempt.status) || isFinalRunStatus(attempt.taskRun.status)) {
         logger.error("Status already in final state", {
           attempt: {
             id: attempt.id,

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

@@ -6,7 +6,7 @@ import { logger } from "~/services/logger.server";
 
 import { PrismaClientOrTransaction, prisma } from "~/db.server";
 import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
-import { CANCELLABLE_STATUSES } from "./cancelTaskRun.server";
+import { isCancellableRunStatus } from "../taskStatus";
 
 export class CancelAttemptService extends BaseService {
   public async call(
@@ -55,7 +55,7 @@ export class CancelAttemptService extends BaseService {
           taskRun: {
             update: {
               data: {
-                status: CANCELLABLE_STATUSES.includes(taskRunAttempt.taskRun.status)
+                status: isCancellableRunStatus(taskRunAttempt.taskRun.status)
                   ? "INTERRUPTED"
                   : undefined,
               },

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

@@ -1,4 +1,4 @@
-import { Prisma, TaskRun, TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database";
+import { Prisma, TaskRun } from "@trigger.dev/database";
 import assertNever from "assert-never";
 import { logger } from "~/services/logger.server";
 import { marqs } from "~/v3/marqs/index.server";
@@ -7,22 +7,7 @@ import { socketIo } from "../handleSocketIo.server";
 import { devPubSub } from "../marqs/devPubSub.server";
 import { BaseService } from "./baseService.server";
 import { CancelAttemptService } from "./cancelAttempt.server";
-
-export const CANCELLABLE_STATUSES: Array<TaskRunStatus> = [
-  "PENDING",
-  "WAITING_FOR_DEPLOY",
-  "EXECUTING",
-  "PAUSED",
-  "WAITING_TO_RESUME",
-  "PAUSED",
-  "RETRYING_AFTER_FAILURE",
-];
-
-const CANCELLABLE_ATTEMPT_STATUSES: Array<TaskRunAttemptStatus> = [
-  "EXECUTING",
-  "PAUSED",
-  "PENDING",
-];
+import { CANCELLABLE_ATTEMPT_STATUSES, isCancellableRunStatus } from "../taskStatus";
 
 type ExtendedTaskRun = Prisma.TaskRunGetPayload<{
   include: {
@@ -53,7 +38,11 @@ export class CancelTaskRunService extends BaseService {
     };
 
     // Make sure the task run is in a cancellable state
-    if (!CANCELLABLE_STATUSES.includes(taskRun.status)) {
+    if (!isCancellableRunStatus(taskRun.status)) {
+      logger.error("Task run is not in a cancellable state", {
+        runId: taskRun.id,
+        status: taskRun.status,
+      });
       return;
     }
 

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

@@ -1,27 +1,11 @@
-import {
-  TaskRun,
-  TaskRunAttempt,
-  TaskRunAttemptStatus,
-  TaskRunStatus,
-} from "@trigger.dev/database";
+import { TaskRun, TaskRunAttempt } from "@trigger.dev/database";
 import { eventRepository } from "../eventRepository.server";
 import { marqs } from "~/v3/marqs/index.server";
 import { BaseService } from "./baseService.server";
 import { logger } from "~/services/logger.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
-
-export const CRASHABLE_RUN_STATUSES: Array<TaskRunStatus> = [
-  "PENDING",
-  "WAITING_FOR_DEPLOY",
-  "EXECUTING",
-  "PAUSED",
-  "WAITING_TO_RESUME",
-  "PAUSED",
-  "RETRYING_AFTER_FAILURE",
-];
-
-const CRASHABLE_ATTEMPT_STATUSES: Array<TaskRunAttemptStatus> = ["EXECUTING", "PAUSED", "PENDING"];
+import { CRASHABLE_ATTEMPT_STATUSES, isCrashableRunStatus } from "../taskStatus";
 
 export type CrashTaskRunServiceOptions = {
   reason?: string;
@@ -52,7 +36,8 @@ export class CrashTaskRunService extends BaseService {
     }
 
     // Make sure the task run is in a crashable state
-    if (!CRASHABLE_RUN_STATUSES.includes(taskRun.status)) {
+    if (!isCrashableRunStatus(taskRun.status)) {
+      logger.error("Task run is not in a crashable state", { runId, status: taskRun.status });
       return;
     }
 

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

@@ -1,19 +1,13 @@
 import { CoordinatorToPlatformMessages } from "@trigger.dev/core/v3";
 import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket";
-import type {
-  CheckpointRestoreEvent,
-  TaskRunAttemptStatus,
-  TaskRunStatus,
-} from "@trigger.dev/database";
+import type { Checkpoint, CheckpointRestoreEvent } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
 import { generateFriendlyId } from "../friendlyIdentifiers";
 import { marqs } from "~/v3/marqs/index.server";
 import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
 import { BaseService } from "./baseService.server";
 import { CrashTaskRunService } from "./crashTaskRun.server";
-
-const FREEZABLE_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "RETRYING_AFTER_FAILURE"];
-const FREEZABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["EXECUTING", "FAILED"];
+import { isFinalRunStatus, isFreezableAttemptStatus, isFreezableRunStatus } from "../taskStatus";
 
 export class CreateCheckpointService extends BaseService {
   public async call(
@@ -21,7 +15,14 @@ export class CreateCheckpointService extends BaseService {
       InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "CHECKPOINT_CREATED">,
       "version"
     >
-  ) {
+  ): Promise<
+    | {
+        checkpoint: Checkpoint;
+        event: CheckpointRestoreEvent;
+        keepRunAlive: boolean;
+      }
+    | undefined
+  > {
     logger.debug(`Creating checkpoint`, params);
 
     const attempt = await this._prisma.taskRunAttempt.findUnique({
@@ -49,8 +50,8 @@ export class CreateCheckpointService extends BaseService {
     }
 
     if (
-      !FREEZABLE_ATTEMPT_STATUSES.includes(attempt.status) ||
-      !FREEZABLE_RUN_STATUSES.includes(attempt.taskRun.status)
+      !isFreezableAttemptStatus(attempt.status) ||
+      !isFreezableRunStatus(attempt.taskRun.status)
     ) {
       logger.error("Unfreezable state", {
         attempt: {
@@ -115,7 +116,9 @@ export class CreateCheckpointService extends BaseService {
     });
 
     const { reason } = params;
+
     let checkpointEvent: CheckpointRestoreEvent | undefined;
+    let keepRunAlive = false;
 
     switch (reason.type) {
       case "WAIT_FOR_DURATION": {
@@ -131,7 +134,12 @@ export class CreateCheckpointService extends BaseService {
           dependencyFriendlyRunId: reason.friendlyId,
         });
 
-        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        keepRunAlive = await this.#isRunCompleted(reason.friendlyId);
+
+        if (!keepRunAlive) {
+          await marqs?.acknowledgeMessage(attempt.taskRunId);
+        }
+
         break;
       }
       case "WAIT_FOR_BATCH": {
@@ -140,7 +148,12 @@ export class CreateCheckpointService extends BaseService {
           batchDependencyFriendlyId: reason.batchFriendlyId,
         });
 
-        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        keepRunAlive = await this.#isBatchCompleted(reason.batchFriendlyId);
+
+        if (!keepRunAlive) {
+          await marqs?.acknowledgeMessage(attempt.taskRunId);
+        }
+
         break;
       }
       case "RETRYING_AFTER_FAILURE": {
@@ -180,6 +193,37 @@ export class CreateCheckpointService extends BaseService {
     return {
       checkpoint,
       event: checkpointEvent,
+      keepRunAlive,
     };
   }
+
+  async #isBatchCompleted(friendlyId: string): Promise<boolean> {
+    const batch = await this._prisma.batchTaskRun.findUnique({
+      where: {
+        friendlyId,
+      },
+    });
+
+    if (!batch) {
+      logger.error("Batch not found", { friendlyId });
+      return false;
+    }
+
+    return batch.status === "COMPLETED";
+  }
+
+  async #isRunCompleted(friendlyId: string): Promise<boolean> {
+    const run = await this._prisma.taskRun.findUnique({
+      where: {
+        friendlyId,
+      },
+    });
+
+    if (!run) {
+      logger.error("Run not found", { friendlyId });
+      return false;
+    }
+
+    return isFinalRunStatus(run.status);
+  }
 }

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

@@ -1,12 +1,10 @@
-import { TaskRunAttemptStatus, TaskRunStatus, type Checkpoint } from "@trigger.dev/database";
+import { type Checkpoint } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
 import { socketIo } from "../handleSocketIo.server";
 import { machinePresetFromConfig } from "../machinePresets.server";
 import { BaseService } from "./baseService.server";
 import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
-
-const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"];
-const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"];
+import { isRestorableAttemptStatus, isRestorableRunStatus } from "../taskStatus";
 
 export class RestoreCheckpointService extends BaseService {
   public async call(params: {
@@ -51,10 +49,7 @@ export class RestoreCheckpointService extends BaseService {
 
     const checkpoint = checkpointEvent.checkpoint;
 
-    const runIsRestorable = RESTORABLE_RUN_STATUSES.includes(checkpoint.run.status);
-    const attemptIsRestorable = RESTORABLE_ATTEMPT_STATUSES.includes(checkpoint.attempt.status);
-
-    if (!runIsRestorable) {
+    if (!isRestorableRunStatus(checkpoint.run.status)) {
       logger.error("Run is unrestorable", {
         eventId: params.eventId,
         runId: checkpoint.runId,
@@ -64,7 +59,7 @@ export class RestoreCheckpointService extends BaseService {
       return;
     }
 
-    if (!attemptIsRestorable && !params.isRetry) {
+    if (!isRestorableAttemptStatus(checkpoint.attempt.status) && !params.isRetry) {
       logger.error("Attempt is unrestorable", {
         eventId: params.eventId,
         runId: checkpoint.runId,

apps/webapp/app/v3/taskStatus.ts

@@ -0,0 +1,69 @@
+import { TaskRunAttemptStatus, TaskRunStatus } from "@trigger.dev/database";
+
+export const CANCELLABLE_RUN_STATUSES: TaskRunStatus[] = [
+  "PENDING",
+  "WAITING_FOR_DEPLOY",
+  "EXECUTING",
+  "PAUSED",
+  "WAITING_TO_RESUME",
+  "PAUSED",
+  "RETRYING_AFTER_FAILURE",
+];
+export const CANCELLABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = [
+  "EXECUTING",
+  "PAUSED",
+  "PENDING",
+];
+
+export function isCancellableRunStatus(status: TaskRunStatus): boolean {
+  return CANCELLABLE_RUN_STATUSES.includes(status);
+}
+export function isCancellableAttemptStatus(status: TaskRunAttemptStatus): boolean {
+  return CANCELLABLE_ATTEMPT_STATUSES.includes(status);
+}
+
+export const CRASHABLE_RUN_STATUSES: TaskRunStatus[] = CANCELLABLE_RUN_STATUSES;
+export const CRASHABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = CANCELLABLE_ATTEMPT_STATUSES;
+
+export function isCrashableRunStatus(status: TaskRunStatus): boolean {
+  return CRASHABLE_RUN_STATUSES.includes(status);
+}
+export function isCrashableAttemptStatus(status: TaskRunAttemptStatus): boolean {
+  return CRASHABLE_ATTEMPT_STATUSES.includes(status);
+}
+
+export const FINAL_RUN_STATUSES: TaskRunStatus[] = [
+  "CANCELED",
+  "COMPLETED_SUCCESSFULLY",
+  "COMPLETED_WITH_ERRORS",
+  "INTERRUPTED",
+  "SYSTEM_FAILURE",
+];
+export const FINAL_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["CANCELED", "COMPLETED", "FAILED"];
+
+export const FREEZABLE_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "RETRYING_AFTER_FAILURE"];
+export const FREEZABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["EXECUTING", "FAILED"];
+
+export function isFreezableRunStatus(status: TaskRunStatus): boolean {
+  return FREEZABLE_RUN_STATUSES.includes(status);
+}
+export function isFreezableAttemptStatus(status: TaskRunAttemptStatus): boolean {
+  return FREEZABLE_ATTEMPT_STATUSES.includes(status);
+}
+
+export function isFinalRunStatus(status: TaskRunStatus): boolean {
+  return FINAL_RUN_STATUSES.includes(status);
+}
+export function isFinalAttemptStatus(status: TaskRunAttemptStatus): boolean {
+  return FINAL_ATTEMPT_STATUSES.includes(status);
+}
+
+export const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"];
+export const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"];
+
+export function isRestorableRunStatus(status: TaskRunStatus): boolean {
+  return RESTORABLE_RUN_STATUSES.includes(status);
+}
+export function isRestorableAttemptStatus(status: TaskRunAttemptStatus): boolean {
+  return RESTORABLE_ATTEMPT_STATUSES.includes(status);
+}

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

@@ -554,6 +554,10 @@ export const CoordinatorToPlatformMessages = {
         }),
       ]),
     }),
+    callback: z.object({
+      version: z.literal("v1").default("v1"),
+      keepRunAlive: z.boolean(),
+    }),
   },
   INDEXING_FAILED: {
     message: z.object({