fix(run-engine,webapp): harden the batch seal-timeout reaper

matt-aitken · matt-aitken · commit 79bd92e17eec · 2026-06-22T18:05:35.000+01:00
Don't abort a batch whose items were all processed (processingCompletedAt set)
before the seal landed; those runs resume the parent on their own.

Keep expireBatch idempotent: an already-aborted batch still resolves its parent
waitpoint, so a retry after a mid-run crash can't leave the parent blocked.

Treat ABORTED as terminal in tryCompleteBatch so a straggler run finalizing
can't flip the batch back to COMPLETED.

Abort the batch if scheduling the reaper fails, so a blocked parent is never
stranded with nothing to free it.
diff --git a/apps/webapp/app/runEngine/services/createBatch.server.ts b/apps/webapp/app/runEngine/services/createBatch.server.ts
@@ -150,13 +150,20 @@ export class CreateBatchService extends WithRunEngine {
 
           await this._engine.initializeBatch(initOptions);
 
-          // Guard the 2-phase gap: if Phase 2 never seals this batch, the reaper
-          // aborts it after the timeout and resumes any blocked parent with an
-          // error instead of leaving it suspended forever.
-          await this._engine.scheduleExpireBatch({
-            batchId: batch.id,
-            availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
-          });
+          // Guard the gap between creating the batch and sealing it: if the item
+          // stream never seals this batch, the reaper aborts it after the timeout
+          // and resumes any blocked parent with an error instead of leaving it
+          // suspended forever. If scheduling the reaper itself fails, abort the
+          // batch now so a blocked parent can't be stranded with nothing to free it.
+          try {
+            await this._engine.scheduleExpireBatch({
+              batchId: batch.id,
+              availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
+            });
+          } catch (scheduleError) {
+            await this._engine.expireBatch({ batchId: batch.id });
+            throw scheduleError;
+          }
 
           logger.info("Batch created", {
             batchId: friendlyId,
diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.ts
@@ -62,7 +62,7 @@ export class BatchSystem {
       span.setAttribute("batchId", batchId);
 
       const batch = await this.$.prisma.batchTaskRun.findFirst({
-        select: { status: true, sealed: true },
+        select: { status: true, sealed: true, processingCompletedAt: true },
         where: { id: batchId },
       });
 
@@ -71,40 +71,58 @@ export class BatchSystem {
         return;
       }
 
-      // The stream sealed the batch, or it already progressed — nothing to fail.
-      if (batch.sealed || batch.status !== "PENDING") {
-        this.$.logger.debug("expireBatch: batch already sealed or no longer PENDING", {
+      // The stream sealed the batch, so the normal completion path owns it.
+      if (batch.sealed) {
+        this.$.logger.debug("expireBatch: batch already sealed", { batchId });
+        return;
+      }
+
+      // Terminal states other than ABORTED are done. ABORTED falls through on
+      // purpose: a previous attempt may have aborted the batch but crashed before
+      // resolving the waitpoint, and completeWaitpoint is idempotent, so retrying
+      // can't leave the parent blocked.
+      if (batch.status !== "PENDING" && batch.status !== "ABORTED") {
+        this.$.logger.debug("expireBatch: batch in terminal non-aborted state", {
           batchId,
           status: batch.status,
-          sealed: batch.sealed,
         });
         return;
       }
 
-      // Conditional update guards against racing a late seal — whichever loses no-ops.
-      const aborted = await this.$.prisma.batchTaskRun.updateMany({
-        where: { id: batchId, sealed: false, status: "PENDING" },
-        data: {
-          status: "ABORTED",
-          completedAt: new Date(),
-          processingCompletedAt: new Date(),
-        },
-      });
-
-      if (aborted.count === 0) {
-        this.$.logger.debug("expireBatch: lost race to seal, no-op", { batchId });
+      // The BatchQueue already processed every item (the completion callback set
+      // processingCompletedAt) even though the seal never landed — the runs exist
+      // and will resume the parent on their own. Aborting would fail a healthy batch.
+      if (batch.status === "PENDING" && batch.processingCompletedAt !== null) {
+        this.$.logger.debug("expireBatch: items already processed, not aborting", { batchId });
         return;
       }
 
-      // Only batchTriggerAndWait blocks a parent, so only it has a waitpoint to resolve.
+      if (batch.status === "PENDING") {
+        // Conditional update guards against racing a late seal or completion —
+        // whichever loses no-ops.
+        const aborted = await this.$.prisma.batchTaskRun.updateMany({
+          where: { id: batchId, sealed: false, status: "PENDING", processingCompletedAt: null },
+          data: {
+            status: "ABORTED",
+            completedAt: new Date(),
+          },
+        });
+
+        if (aborted.count === 0) {
+          this.$.logger.debug("expireBatch: lost race to seal/complete, no-op", { batchId });
+          return;
+        }
+      }
+
+      // Only batchTriggerAndWait blocks a parent, so only it has a waitpoint to
+      // resolve. The status filter keeps this idempotent if a previous attempt
+      // already resolved it.
       const waitpoint = await this.$.prisma.waitpoint.findFirst({
-        where: { completedByBatchId: batchId },
+        where: { completedByBatchId: batchId, status: "PENDING" },
       });
 
       if (!waitpoint) {
-        this.$.logger.debug("expireBatch: no waitpoint to resolve (fire-and-forget batch)", {
-          batchId,
-        });
+        this.$.logger.debug("expireBatch: no pending waitpoint to resolve", { batchId });
         return;
       }
 
@@ -151,8 +169,11 @@ export class BatchSystem {
         return;
       }
 
-      if (batch.status === "COMPLETED") {
-        this.$.logger.debug("#tryCompleteBatch: Batch already completed", { batchId });
+      if (batch.status === "COMPLETED" || batch.status === "ABORTED") {
+        this.$.logger.debug("#tryCompleteBatch: Batch already in a terminal status", {
+          batchId,
+          status: batch.status,
+        });
         return;
       }
 
diff --git a/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts b/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts
@@ -13,9 +13,53 @@ import type {
   BatchItem,
   InitializeBatchOptions,
 } from "../../batch-queue/types.js";
+import { type PrismaClient } from "@trigger.dev/database";
+import { type RedisOptions } from "ioredis";
 
 vi.setConfig({ testTimeout: 60_000 });
 
+function createBatchTestEngine(prisma: PrismaClient, redisOptions: RedisOptions) {
+  return new RunEngine({
+    prisma,
+    worker: {
+      redis: redisOptions,
+      workers: 1,
+      tasksPerWorker: 10,
+      pollIntervalMs: 20,
+    },
+    queue: {
+      redis: redisOptions,
+      masterQueueConsumersDisabled: true,
+      processWorkerQueueDebounceMs: 50,
+    },
+    runLock: {
+      redis: redisOptions,
+    },
+    machines: {
+      defaultMachine: "small-1x",
+      machines: {
+        "small-1x": {
+          name: "small-1x" as const,
+          cpu: 0.5,
+          memory: 0.5,
+          centsPerMs: 0.0001,
+        },
+      },
+      baseCostInCents: 0.0001,
+    },
+    batchQueue: {
+      redis: redisOptions,
+      consumerCount: 2,
+      consumerIntervalMs: 50,
+      drr: {
+        quantum: 10,
+        maxDeficit: 100,
+      },
+    },
+    tracer: trace.getTracer("test", "0.0.0"),
+  });
+}
+
 describe("RunEngine 2-Phase Batch API", () => {
   containerTest(
     "2-phase batch: initialize batch, stream items one by one, items get processed",
@@ -823,6 +867,178 @@ describe("RunEngine 2-Phase Batch API", () => {
     }
   );
 
+  containerTest(
+    "2-phase batch: expireBatch does not abort a batch whose items were all processed",
+    async ({ prisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+      const engine = createBatchTestEngine(prisma, redisOptions);
+      engine.setBatchProcessItemCallback(async () => ({ success: true, runId: "x" }));
+      engine.setBatchCompletionCallback(async () => {});
+
+      try {
+        // The BatchQueue streamed and processed every item (the completion callback
+        // set processingCompletedAt) but the seal never landed. The runs exist and
+        // will resume the parent on their own, so the reaper must NOT abort this.
+        const { id: batchId, friendlyId } = BatchId.generate();
+        await prisma.batchTaskRun.create({
+          data: {
+            id: batchId,
+            friendlyId,
+            runtimeEnvironmentId: authenticatedEnvironment.id,
+            status: "PENDING",
+            runCount: 2,
+            expectedCount: 2,
+            sealed: false,
+            batchVersion: "runengine:v2",
+            processingCompletedAt: new Date(),
+          },
+        });
+
+        await engine.expireBatch({ batchId });
+
+        const batch = await prisma.batchTaskRun.findUnique({ where: { id: batchId } });
+        expect(batch?.status).toBe("PENDING");
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
+  containerTest(
+    "2-phase batch: expireBatch resumes the parent even if a prior attempt already aborted the batch",
+    async ({ prisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+      const engine = createBatchTestEngine(prisma, redisOptions);
+      engine.setBatchProcessItemCallback(async () => ({ success: true, runId: "x" }));
+      engine.setBatchCompletionCallback(async () => {});
+
+      try {
+        const parentTask = "parent-task";
+        await setupBackgroundWorker(engine, authenticatedEnvironment, [parentTask]);
+
+        const { id: batchId, friendlyId } = BatchId.generate();
+        await prisma.batchTaskRun.create({
+          data: {
+            id: batchId,
+            friendlyId,
+            runtimeEnvironmentId: authenticatedEnvironment.id,
+            status: "PENDING",
+            runCount: 2,
+            expectedCount: 2,
+            sealed: false,
+            batchVersion: "runengine:v2",
+          },
+        });
+
+        const parentRun = await engine.trigger(
+          {
+            friendlyId: generateFriendlyId("run"),
+            environment: authenticatedEnvironment,
+            taskIdentifier: parentTask,
+            payload: "{}",
+            payloadType: "application/json",
+            context: {},
+            traceContext: {},
+            traceId: "t_parent",
+            spanId: "s_parent",
+            workerQueue: "main",
+            queue: `task/${parentTask}`,
+            isTest: false,
+            tags: [],
+          },
+          prisma
+        );
+
+        await setTimeout(500);
+        const dequeued = await engine.dequeueFromWorkerQueue({
+          consumerId: "test_12345",
+          workerQueue: "main",
+        });
+        expect(dequeued.length).toBe(1);
+
+        const initial = await engine.getRunExecutionData({ runId: parentRun.id });
+        assertNonNullable(initial);
+        await engine.startRunAttempt({ runId: parentRun.id, snapshotId: initial.snapshot.id });
+
+        await engine.blockRunWithCreatedBatch({
+          runId: parentRun.id,
+          batchId,
+          environmentId: authenticatedEnvironment.id,
+          projectId: authenticatedEnvironment.projectId,
+          organizationId: authenticatedEnvironment.organizationId,
+        });
+
+        // Simulate a prior expireBatch attempt that aborted the batch but crashed
+        // before resolving the parent's waitpoint.
+        await prisma.batchTaskRun.update({
+          where: { id: batchId },
+          data: { status: "ABORTED", completedAt: new Date() },
+        });
+
+        // The retry must still resolve the parent's waitpoint so it can't stay stuck.
+        await engine.expireBatch({ batchId });
+
+        const waitpoint = await prisma.waitpoint.findFirst({
+          where: { completedByBatchId: batchId },
+        });
+        assertNonNullable(waitpoint);
+        expect(waitpoint.status).toBe("COMPLETED");
+        expect(waitpoint.outputIsError).toBe(true);
+
+        await vi.waitFor(
+          async () => {
+            const wps = await prisma.taskRunWaitpoint.findMany({
+              where: { taskRunId: parentRun.id },
+            });
+            expect(wps.length).toBe(0);
+          },
+          { timeout: 15000 }
+        );
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
+  containerTest(
+    "2-phase batch: tryCompleteBatch leaves an ABORTED batch terminal",
+    async ({ prisma, redisOptions }) => {
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+      const engine = createBatchTestEngine(prisma, redisOptions);
+      engine.setBatchProcessItemCallback(async () => ({ success: true, runId: "x" }));
+      engine.setBatchCompletionCallback(async () => {});
+
+      try {
+        // An aborted batch is terminal: a later tryCompleteBatch (e.g. from a
+        // straggler child run finalizing) must not flip it back to COMPLETED.
+        const { id: batchId, friendlyId } = BatchId.generate();
+        await prisma.batchTaskRun.create({
+          data: {
+            id: batchId,
+            friendlyId,
+            runtimeEnvironmentId: authenticatedEnvironment.id,
+            status: "ABORTED",
+            runCount: 0,
+            expectedCount: 0,
+            sealed: false,
+            batchVersion: "runengine:v2",
+            completedAt: new Date(),
+          },
+        });
+
+        await engine.tryCompleteBatch({ batchId });
+
+        // Allow the debounced completion job to run (or correctly no-op).
+        await setTimeout(1500);
+
+        const batch = await prisma.batchTaskRun.findUnique({ where: { id: batchId } });
+        expect(batch?.status).toBe("ABORTED");
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
   containerTest(
     "2-phase batch: error if batch not initialized",
     async ({ prisma, redisOptions }) => {
