Support for the same run multiple times in the same batch

Author: matt-aitken

internal-packages/database/prisma/migrations/20250114153223_task_run_waitpoint_unique_constraint_added_batch_index/migration.sql

@@ -0,0 +1,14 @@
+/*
+  Warnings:
+
+  - A unique constraint covering the columns `[taskRunId,waitpointId,batchIndex]` on the table `TaskRunWaitpoint` will be added. If there are existing duplicate values, this will fail.
+
+*/
+-- DropIndex
+DROP INDEX "TaskRunWaitpoint_taskRunId_waitpointId_key";
+
+-- CreateIndex (multiple can have null batchIndex, so we need the other one below)
+CREATE UNIQUE INDEX "TaskRunWaitpoint_taskRunId_waitpointId_batchIndex_key" ON "TaskRunWaitpoint" ("taskRunId", "waitpointId", "batchIndex");
+
+-- CreateIndex (where batchIndex is null)
+CREATE UNIQUE INDEX "TaskRunWaitpoint_taskRunId_waitpointId_batchIndex_null_key" ON "TaskRunWaitpoint"("taskRunId", "waitpointId") WHERE "batchIndex" IS NULL;

internal-packages/database/prisma/schema.prisma

@@ -2120,7 +2120,9 @@ model TaskRunWaitpoint {
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
 
-  @@unique([taskRunId, waitpointId])
+  /// There are two constraints, the one below and also one that Prisma doesn't support
+  /// The second one implemented in SQL only prevents a TaskRun + Waitpoint with a null batchIndex
+  @@unique([taskRunId, waitpointId, batchIndex])
   @@index([taskRunId])
   @@index([waitpointId])
 }

internal-packages/run-engine/src/engine/executionSnapshots.ts

@@ -35,40 +35,55 @@ export async function getLatestExecutionSnapshot(
     ...snapshot,
     friendlyId: SnapshotId.toFriendlyId(snapshot.id),
     runFriendlyId: RunId.toFriendlyId(snapshot.runId),
-    completedWaitpoints: snapshot.completedWaitpoints.map((w) => {
-      const index = snapshot.completedWaitpointOrder.findIndex((s) => s === w.id);
+    completedWaitpoints: snapshot.completedWaitpoints.flatMap((w) => {
+      //get all indexes of the waitpoint in the completedWaitpointOrder
+      //we do this because the same run can be in a batch multiple times (i.e. same idempotencyKey)
+      let indexes: (number | undefined)[] = [];
+      for (let i = 0; i < snapshot.completedWaitpointOrder.length; i++) {
+        if (snapshot.completedWaitpointOrder[i] === w.id) {
+          indexes.push(i);
+        }
+      }
+
+      if (indexes.length === 0) {
+        indexes.push(undefined);
+      }
 
-      return {
-        id: w.id,
-        index: index === -1 ? undefined : index,
-        friendlyId: w.friendlyId,
-        type: w.type,
-        completedAt: w.completedAt ?? new Date(),
-        idempotencyKey:
-          w.userProvidedIdempotencyKey && !w.inactiveIdempotencyKey ? w.idempotencyKey : undefined,
-        completedByTaskRun: w.completedByTaskRunId
-          ? {
-              id: w.completedByTaskRunId,
-              friendlyId: RunId.toFriendlyId(w.completedByTaskRunId),
-              batch: snapshot.batchId
-                ? {
-                    id: snapshot.batchId,
-                    friendlyId: BatchId.toFriendlyId(snapshot.batchId),
-                  }
-                : undefined,
-            }
-          : undefined,
-        completedAfter: w.completedAfter ?? undefined,
-        completedByBatch: w.completedByBatchId
-          ? {
-              id: w.completedByBatchId,
-              friendlyId: BatchId.toFriendlyId(w.completedByBatchId),
-            }
-          : undefined,
-        output: w.output ?? undefined,
-        outputType: w.outputType,
-        outputIsError: w.outputIsError,
-      } satisfies CompletedWaitpoint;
+      return indexes.map((index) => {
+        return {
+          id: w.id,
+          index: index === -1 ? undefined : index,
+          friendlyId: w.friendlyId,
+          type: w.type,
+          completedAt: w.completedAt ?? new Date(),
+          idempotencyKey:
+            w.userProvidedIdempotencyKey && !w.inactiveIdempotencyKey
+              ? w.idempotencyKey
+              : undefined,
+          completedByTaskRun: w.completedByTaskRunId
+            ? {
+                id: w.completedByTaskRunId,
+                friendlyId: RunId.toFriendlyId(w.completedByTaskRunId),
+                batch: snapshot.batchId
+                  ? {
+                      id: snapshot.batchId,
+                      friendlyId: BatchId.toFriendlyId(snapshot.batchId),
+                    }
+                  : undefined,
+              }
+            : undefined,
+          completedAfter: w.completedAfter ?? undefined,
+          completedByBatch: w.completedByBatchId
+            ? {
+                id: w.completedByBatchId,
+                friendlyId: BatchId.toFriendlyId(w.completedByBatchId),
+              }
+            : undefined,
+          output: w.output ?? undefined,
+          outputType: w.outputType,
+          outputIsError: w.outputIsError,
+        } satisfies CompletedWaitpoint;
+      });
     }),
   };
 }