fix(webapp): tighten Phase 2 idempotency check; cover PARTIAL_FAILED race

matt-aitken · claude · matt-aitken · commit e8caa1dfd5fd · 2026-05-27T14:50:00.000-07:00
Addresses two PR review findings:

CodeRabbit: sealed=true + ABORTED would silently succeed under the
previous `if (batch.sealed || ...)` check. V2's batch completion
callback can set status=ABORTED (failedRunCount &gt; 0 &amp;&amp;
successfulRunCount === 0) on a batch that streamBatchItems already
sealed — leaving sealed=true alongside a terminally-failed batch.
A Phase 2 retry of such a batch must surface the error, not mask it.

Devin: PARTIAL_FAILED (failedRunCount &gt; 0 with at least one success)
is a real V2 completion-callback status, but neither the pre-loop
check nor the post-loop race handlers (lines 226 and 306) accepted
it as success. A retry whose original stream succeeded would either
422 at the pre-loop or hit "unexpected state" at the post-loop seal-
failed branch.

Changes:
- Pre-loop: replace the broad `sealed || PROCESSING || COMPLETED`
  check with an `isIdempotentRetrySuccess` boolean that admits
  PROCESSING, COMPLETED, PARTIAL_FAILED, or (sealed &amp;&amp; PENDING) —
  ABORTED falls through to the throw.
- Post-loop count-mismatch (line 226 region): add PARTIAL_FAILED to
  the success short-circuit alongside sealed and COMPLETED.
- Post-loop seal-failed (line 306 region): add PARTIAL_FAILED to the
  success short-circuit alongside (sealed &amp;&amp; PROCESSING) and
  COMPLETED.

Tests (TDD red-then-green):
- New: pre-loop sealed=true + ABORTED → throws (CodeRabbit's case).
- New: pre-loop PARTIAL_FAILED → returns sealed:true.
- New: post-loop seal-failed race with PARTIAL_FAILED → returns
  sealed:true (uses the same racingPrisma pattern as the existing
  COMPLETED race test).
All 34 tests in streamBatchItems.test.ts pass.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/apps/webapp/app/runEngine/services/streamBatchItems.server.ts b/apps/webapp/app/runEngine/services/streamBatchItems.server.ts
@@ -100,12 +100,26 @@ export class StreamBatchItemsService extends WithRunEngine {
           throw new ServiceValidationError(`Batch ${batchFriendlyId} not found`);
         }
 
-        // Phase 2 retry idempotency (TRI-9944): if the batch is already sealed
-        // or has moved past PENDING into PROCESSING/COMPLETED, this is a retry
-        // of a request whose response was lost — the original successful request
-        // already enqueued every item and sealed the batch. Returning sealed:true
-        // makes the SDK stop retrying instead of throwing a customer-visible 422.
-        if (batch.sealed || batch.status === "PROCESSING" || batch.status === "COMPLETED") {
+        // Phase 2 retry idempotency (TRI-9944): a successful original request
+        // sealed the batch (sealed=true, status=PROCESSING) and the V2 batch
+        // completion callback can then independently update status to:
+        //   - PENDING (all runs created — sealed stays true)
+        //   - PARTIAL_FAILED (some run creations failed — sealed stays true/false)
+        //   - COMPLETED (set by tryCompleteBatch after every run reaches a final
+        //     state — sealed is NOT set by this path)
+        // For all of these the Phase 2 stream did its job, so a retry should
+        // return sealed:true and the SDK stops retrying.
+        //
+        // ABORTED is explicitly excluded — it means every run-creation attempt
+        // failed and the batch is terminally broken; surface that as an error
+        // rather than masking it as success.
+        const isIdempotentRetrySuccess =
+          batch.status === "PROCESSING" ||
+          batch.status === "COMPLETED" ||
+          batch.status === "PARTIAL_FAILED" ||
+          (batch.sealed && batch.status === "PENDING");
+
+        if (isIdempotentRetrySuccess) {
           logger.info("Batch already sealed/completed - treating Phase 2 retry as success", {
             batchId: batchFriendlyId,
             batchSealed: batch.sealed,
@@ -239,7 +253,11 @@ export class StreamBatchItemsService extends WithRunEngine {
             select: { sealed: true, status: true },
           });
 
-          if (currentBatch?.sealed || currentBatch?.status === "COMPLETED") {
+          if (
+            currentBatch?.sealed ||
+            currentBatch?.status === "COMPLETED" ||
+            currentBatch?.status === "PARTIAL_FAILED"
+          ) {
             logger.info("Batch already sealed before count check (fast completion)", {
               batchId: batchFriendlyId,
               itemsAccepted,
@@ -321,7 +339,8 @@ export class StreamBatchItemsService extends WithRunEngine {
 
           if (
             (currentBatch?.sealed && currentBatch.status === "PROCESSING") ||
-            currentBatch?.status === "COMPLETED"
+            currentBatch?.status === "COMPLETED" ||
+            currentBatch?.status === "PARTIAL_FAILED"
           ) {
             logger.info("Batch already sealed/completed by concurrent path", {
               batchId: batchFriendlyId,
diff --git a/apps/webapp/test/engine/streamBatchItems.test.ts b/apps/webapp/test/engine/streamBatchItems.test.ts
@@ -375,6 +375,257 @@ describe("StreamBatchItemsService", () => {
     }
   );
 
+  containerTest(
+    "should throw when batch is sealed but ABORTED (callback aborted post-seal must surface as error)",
+    async ({ prisma, redisOptions }) => {
+      const engine = new RunEngine({
+        prisma,
+        worker: {
+          redis: redisOptions,
+          workers: 1,
+          tasksPerWorker: 10,
+          pollIntervalMs: 100,
+          disabled: true,
+        },
+        queue: {
+          redis: redisOptions,
+        },
+        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.0005,
+        },
+        batchQueue: {
+          redis: redisOptions,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      // V2 batch completion callback sets status=ABORTED (failedRunCount > 0 &&
+      // successfulRunCount === 0) without touching sealed=true that the seal
+      // step previously set. The Phase 2 retry must NOT mask this terminal
+      // failure as success — every run failed.
+      const batch = await createBatch(prisma, authenticatedEnvironment.id, {
+        runCount: 2,
+        status: "ABORTED",
+        sealed: true,
+      });
+
+      const service = new StreamBatchItemsService({
+        prisma,
+        engine,
+      });
+
+      await expect(
+        service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), {
+          maxItemBytes: 1024 * 1024,
+        })
+      ).rejects.toThrow(ServiceValidationError);
+
+      await engine.quit();
+    }
+  );
+
+  containerTest(
+    "should return sealed=true when batch is PARTIAL_FAILED (Phase 2 retry idempotency)",
+    async ({ prisma, redisOptions }) => {
+      const engine = new RunEngine({
+        prisma,
+        worker: {
+          redis: redisOptions,
+          workers: 1,
+          tasksPerWorker: 10,
+          pollIntervalMs: 100,
+          disabled: true,
+        },
+        queue: {
+          redis: redisOptions,
+        },
+        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.0005,
+        },
+        batchQueue: {
+          redis: redisOptions,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      // V2 completion callback sets PARTIAL_FAILED when some run-creation
+      // attempts failed but at least one succeeded. The Phase 2 stream itself
+      // did its job (items were enqueued and processed), so a retry should
+      // see this as terminal success — the per-item failures are visible on
+      // the individual run records.
+      const batch = await createBatch(prisma, authenticatedEnvironment.id, {
+        runCount: 2,
+        status: "PARTIAL_FAILED",
+        sealed: false,
+      });
+
+      const service = new StreamBatchItemsService({
+        prisma,
+        engine,
+      });
+
+      const result = await service.call(
+        authenticatedEnvironment,
+        batch.friendlyId,
+        itemsToAsyncIterable([]),
+        {
+          maxItemBytes: 1024 * 1024,
+        }
+      );
+
+      expect(result.sealed).toBe(true);
+      expect(result.id).toBe(batch.friendlyId);
+      expect(result.itemsAccepted).toBe(0);
+      expect(result.itemsDeduplicated).toBe(0);
+
+      await engine.quit();
+    }
+  );
+
+  containerTest(
+    "should return sealed=true when batch is PARTIAL_FAILED by callback before seal attempt",
+    async ({ prisma, redisOptions }) => {
+      const engine = new RunEngine({
+        prisma,
+        worker: {
+          redis: redisOptions,
+          workers: 1,
+          tasksPerWorker: 10,
+          pollIntervalMs: 100,
+          disabled: true,
+        },
+        queue: {
+          redis: redisOptions,
+        },
+        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.0005,
+        },
+        batchQueue: {
+          redis: redisOptions,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      const batch = await createBatch(prisma, authenticatedEnvironment.id, {
+        runCount: 2,
+        status: "PENDING",
+        sealed: false,
+      });
+
+      await engine.initializeBatch({
+        batchId: batch.id,
+        friendlyId: batch.friendlyId,
+        environmentId: authenticatedEnvironment.id,
+        environmentType: authenticatedEnvironment.type,
+        organizationId: authenticatedEnvironment.organizationId,
+        projectId: authenticatedEnvironment.projectId,
+        runCount: 2,
+        processingConcurrency: 10,
+      });
+
+      await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, {
+        task: "test-task",
+        payload: JSON.stringify({ data: "item1" }),
+        payloadType: "application/json",
+      });
+      await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, {
+        task: "test-task",
+        payload: JSON.stringify({ data: "item2" }),
+        payloadType: "application/json",
+      });
+
+      // Simulate the race where V2's batchCompletionCallback runs between
+      // getEnqueuedCount and the seal updateMany — some runs failed to create
+      // but at least one succeeded, so callback sets status=PARTIAL_FAILED
+      // without setting sealed=true.
+      const racingPrisma = {
+        ...prisma,
+        batchTaskRun: {
+          ...prisma.batchTaskRun,
+          findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun),
+          updateMany: async () => {
+            await prisma.batchTaskRun.update({
+              where: { id: batch.id },
+              data: {
+                status: "PARTIAL_FAILED",
+              },
+            });
+            return { count: 0 };
+          },
+          findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun),
+        },
+      } as unknown as PrismaClient;
+
+      const service = new StreamBatchItemsService({
+        prisma: racingPrisma,
+        engine,
+      });
+
+      const result = await service.call(
+        authenticatedEnvironment,
+        batch.friendlyId,
+        itemsToAsyncIterable([]),
+        {
+          maxItemBytes: 1024 * 1024,
+        }
+      );
+
+      expect(result.sealed).toBe(true);
+      expect(result.id).toBe(batch.friendlyId);
+
+      const updatedBatch = await prisma.batchTaskRun.findUnique({
+        where: { id: batch.id },
+      });
+
+      expect(updatedBatch?.status).toBe("PARTIAL_FAILED");
+      expect(updatedBatch?.sealed).toBe(false);
+
+      await engine.quit();
+    }
+  );
+
   containerTest(
     "should return sealed=true when concurrent request already sealed the batch during seal attempt",
     async ({ prisma, redisOptions }) => {
