fix(run-engine): requeue a failed restore at the run's queue position (#5104)

internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts

@@ -1563,6 +1563,9 @@ export class RunAttemptSystem {
         orgId,
         messageId: run.id,
         retryAt: timestamp,
+        // A resume is enqueued at the run's original position, so a resume that fails before the
+        // run continues keeps it there instead of moving behind runs triggered after it.
+        keepQueuePosition: checkpointId !== undefined,
         resetAttemptCount: resetQueueAttempts,
         workerQueue: moveToWorkerQueue ?? workerQueue,
       });

internal-packages/run-engine/src/engine/tests/restoreOutcome.test.ts

@@ -204,6 +204,73 @@ describe("RunEngine reportRestoreOutcome", () => {
     }
   );
 
+  containerTest(
+    "requeue keeps the run ahead of runs triggered after it once, then backs off",
+    async ({ prisma, redisOptions }) => {
+      // Catches: a resume whose restore fails waiting behind every run triggered after it, and a
+      // restore that keeps failing being redelivered with no backoff.
+      const engine = createEngine(prisma, redisOptions);
+
+      try {
+        const { authenticatedEnvironment, run, restore } = await dequeueForRestore(engine, prisma);
+
+        // The restore holds the queue's only slot, so a run triggered now waits in the queue.
+        await engine.runQueue.updateQueueConcurrencyLimits(
+          authenticatedEnvironment,
+          "task/test-task",
+          1
+        );
+        const newerRun = await engine.trigger(
+          {
+            number: 2,
+            friendlyId: "run_5678",
+            environment: authenticatedEnvironment,
+            taskIdentifier: "test-task",
+            payload: "{}",
+            payloadType: "application/json",
+            context: {},
+            traceContext: {},
+            traceId: "t5678",
+            spanId: "s5678",
+            workerQueue: "main",
+            queue: "task/test-task",
+            isTest: false,
+            tags: [],
+          },
+          prisma
+        );
+
+        const result = await engine.reportRestoreOutcome({
+          runId: run.id,
+          snapshotId: restore.snapshot.id,
+          outcome: "requeue",
+          reason: "OOMKilled",
+        });
+        expect(result).toEqual({ ok: true, outcome: "requeue" });
+
+        const redelivered = await dequeue(engine, authenticatedEnvironment.id);
+        expect(redelivered.length).toBe(1);
+        assertNonNullable(redelivered[0]);
+        expect(redelivered[0].run.id).toBe(run.id);
+
+        const secondResult = await engine.reportRestoreOutcome({
+          runId: run.id,
+          snapshotId: redelivered[0].snapshot.id,
+          outcome: "requeue",
+          reason: "OOMKilled",
+        });
+        expect(secondResult).toEqual({ ok: true, outcome: "requeue" });
+
+        const afterSecondRequeue = await dequeue(engine, authenticatedEnvironment.id);
+        expect(afterSecondRequeue.length).toBe(1);
+        assertNonNullable(afterSecondRequeue[0]);
+        expect(afterSecondRequeue[0].run.id).toBe(newerRun.id);
+      } finally {
+        await engine.quit();
+      }
+    }
+  );
+
   containerTest(
     "fail fails the attempt as a crash and retries from scratch after the task's retry delay",
     async ({ prisma, redisOptions }) => {

internal-packages/run-engine/src/run-queue/index.ts

@@ -2646,6 +2646,7 @@ export class RunQueue {
     resetAttemptCount = false,
     skipDequeueProcessing = false,
     workerQueue,
+    keepQueuePosition = false,
   }: {
     orgId: string;
     messageId: string;
@@ -2660,6 +2661,11 @@ export class RunQueue {
     skipDequeueProcessing?: boolean;
     /** Fresh attempts override the destination; omit for checkpoint redelivery. */
     workerQueue?: string;
+    /**
+     * Without `retryAt`, redeliver the first reattempt at the message's original score instead of
+     * after a backoff.
+     */
+    keepQueuePosition?: boolean;
   }) {
     return this.#trace(
       "nackMessage",
@@ -2722,7 +2728,11 @@ export class RunQueue {
           });
         }
 
-        await this.#callNackMessage({ message, retryAt });
+        await this.#callNackMessage({
+          message,
+          retryAt:
+            retryAt ?? (keepQueuePosition && message.attempt <= 1 ? message.timestamp : undefined),
+        });
 
         this.#emitQueueMetric(message.queue, { op: "nack", q: message.queue });