prevent checkpoint creation for resumed task waits

Author: nicktrn

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

@@ -97,36 +97,77 @@ export class CreateCheckpointService extends BaseService {
 
     // Check if we should accept this checkpoint
     switch (reason.type) {
-      case "MANUAL":
+      case "MANUAL": {
         // Always accept manual checkpoints
         break;
-      case "WAIT_FOR_DURATION":
+      }
+      case "WAIT_FOR_DURATION": {
         // Always accept duration checkpoints
         break;
+      }
       case "WAIT_FOR_TASK": {
-        // TODO
+        const childRun = await this._prisma.taskRun.findFirst({
+          where: {
+            friendlyId: reason.friendlyId,
+          },
+          select: {
+            dependency: {
+              select: {
+                resumedAt: true,
+              },
+            },
+          },
+        });
+
+        if (!childRun) {
+          logger.error("CreateCheckpointService: Pre-check - WAIT_FOR_TASK child run not found", {
+            friendlyId: reason.friendlyId,
+            params,
+          });
+
+          return {
+            success: false,
+            keepRunAlive: false,
+          };
+        }
+
+        if (childRun.dependency?.resumedAt) {
+          logger.error("CreateCheckpointService: Child run already resumed", {
+            childRun,
+            params,
+          });
+
+          return {
+            success: false,
+            keepRunAlive: true,
+          };
+        }
+
         break;
       }
       case "WAIT_FOR_BATCH": {
         const batchRun = await this._prisma.batchTaskRun.findFirst({
           where: {
             friendlyId: reason.batchFriendlyId,
           },
+          select: {
+            resumedAt: true,
+          },
         });
 
         if (!batchRun) {
-          logger.error("CreateCheckpointService: Batch not found", {
+          logger.error("CreateCheckpointService: Pre-check - Batch not found", {
             batchFriendlyId: reason.batchFriendlyId,
             params,
           });
 
           return {
             success: false,
-            keepRunAlive: true,
+            keepRunAlive: false,
           };
         }
 
-        if (batchRun.batchVersion === "v3" && batchRun.resumedAt) {
+        if (batchRun.resumedAt) {
           logger.error("CreateCheckpointService: Batch already resumed", {
             batchRun,
             params,
@@ -140,6 +181,9 @@ export class CreateCheckpointService extends BaseService {
 
         break;
       }
+      default: {
+        break;
+      }
     }
 
     //sleep to test slow checkpoints

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

@@ -3,6 +3,7 @@ import { workerQueue } from "~/services/worker.server";
 import { marqs } from "~/v3/marqs/index.server";
 import { BaseService } from "./baseService.server";
 import { logger } from "~/services/logger.server";
+import { TaskRunDependency } from "@trigger.dev/database";
 
 export class ResumeTaskDependencyService extends BaseService {
   public async call(dependencyId: string, sourceTaskAttemptId: string) {
@@ -49,6 +50,21 @@ export class ResumeTaskDependencyService extends BaseService {
           runId: dependentRun.id,
         }
       );
+
+      const wasUpdated = await this.#setDependencyToResumedOnce(dependency);
+
+      if (!wasUpdated) {
+        logger.debug("Task dependency resume: Attempt with checkpoint was already resumed", {
+          attemptId: dependency.id,
+          dependentAttempt: dependency.dependentAttempt,
+          checkpointEventId: dependency.checkpointEventId,
+          hasCheckpointEvent: !!dependency.checkpointEventId,
+          runId: dependentRun.id,
+        });
+        return;
+      }
+
+
       await marqs?.enqueueMessage(
         dependency.taskRun.runtimeEnvironment,
         dependentRun.queue,
@@ -85,6 +101,19 @@ export class ResumeTaskDependencyService extends BaseService {
         return;
       }
 
+      const wasUpdated = await this.#setDependencyToResumedOnce(dependency);
+
+      if (!wasUpdated) {
+        logger.debug("Task dependency resume: Attempt without checkpoint was already resumed", {
+          attemptId: dependency.id,
+          dependentAttempt: dependency.dependentAttempt,
+          checkpointEventId: dependency.checkpointEventId,
+          hasCheckpointEvent: !!dependency.checkpointEventId,
+          runId: dependentRun.id,
+        });
+        return;
+      }
+
       await marqs?.replaceMessage(
         dependentRun.id,
         {
@@ -102,6 +131,26 @@ export class ResumeTaskDependencyService extends BaseService {
     }
   }
 
+  async #setDependencyToResumedOnce(dependency: TaskRunDependency) {
+    const result = await this._prisma.taskRunDependency.updateMany({
+      where: {
+        id: dependency.id,
+        resumedAt: null,
+      },
+      data: {
+        resumedAt: new Date(),
+      },
+    });
+
+    // Check if any records were updated
+    if (result.count > 0) {
+      // The status was changed, so we return true
+      return true;
+    } else {
+      return false;
+    }
+  }
+
   static async enqueue(
     dependencyId: string,
     sourceTaskAttemptId: string,

internal-packages/database/prisma/migrations/20250218124428_add_resumed_at_to_task_run_dependency/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "TaskRunDependency" ADD COLUMN     "resumedAt" TIMESTAMP(3);

internal-packages/database/prisma/schema.prisma

@@ -1915,6 +1915,7 @@ model TaskRunDependency {
 
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
+  resumedAt DateTime?
 
   @@index([dependentAttemptId])
   @@index([dependentBatchRunId])