From e98d7799f0c0e8506dc191d0875cb3bfed4645d1 Mon Sep 17 00:00:00 2001
From: Eric Allam <eallam@icloud.com>
Date: Wed, 10 Sep 2025 14:43:34 +0100
Subject: [PATCH] fix(run-engine): Preserve snapshot checkpoint ID when a
 PENDING_EXECUTING snapshot stalls

---
 internal-packages/run-engine/src/engine/index.ts               | 1 +
 .../run-engine/src/engine/systems/runAttemptSystem.ts          | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts
index a551e6ec12..e81aedfe8b 100644
--- a/internal-packages/run-engine/src/engine/index.ts
+++ b/internal-packages/run-engine/src/engine/index.ts
@@ -1251,6 +1251,7 @@ export class RunEngine {
             },
             orgId: latestSnapshot.organizationId,
             projectId: latestSnapshot.projectId,
+            checkpointId: latestSnapshot.checkpointId ?? undefined,
             error: {
               type: "INTERNAL_ERROR",
               code: "TASK_RUN_DEQUEUED_MAX_RETRIES",
diff --git a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
index 942c22f6be..842eaefff1 100644
--- a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
+++ b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
@@ -1129,6 +1129,7 @@ export class RunAttemptSystem {
     error,
     workerId,
     runnerId,
+    checkpointId,
     tx,
   }: {
     run: TaskRun;
@@ -1142,6 +1143,7 @@ export class RunAttemptSystem {
     error: TaskRunInternalError;
     workerId?: string;
     runnerId?: string;
+    checkpointId?: string;
     tx?: PrismaClientOrTransaction;
   }): Promise<{ wasRequeued: boolean } & ExecutionResult> {
     const prisma = tx ?? this.$.prisma;
@@ -1189,6 +1191,7 @@ export class RunAttemptSystem {
         organizationId: orgId,
         workerId,
         runnerId,
+        checkpointId,
       });
 
       return {