check dependency completion when creating checkpoints

Author: nicktrn

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/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/services/createCheckpoint.server.ts

@@ -1,6 +1,6 @@
 import { CoordinatorToPlatformMessages } from "@trigger.dev/core/v3";
 import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket";
-import type { CheckpointRestoreEvent } 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";
@@ -15,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({
@@ -109,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": {
@@ -125,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": {
@@ -134,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": {
@@ -174,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);
+  }
 }

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({