diff --git a/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts b/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts
index 71f54b988e..c6586643ba 100644
--- a/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts
+++ b/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts
@@ -194,6 +194,11 @@ export class CheckpointSystem {
             metadata: snapshot.metadata,
           },
           previousSnapshotId: snapshot.id,
+          batchId: snapshot.batchId ?? undefined,
+          completedWaitpoints: snapshot.completedWaitpoints.map((waitpoint) => ({
+            id: waitpoint.id,
+            index: waitpoint.index,
+          })),
           environmentId: snapshot.environmentId,
           environmentType: snapshot.environmentType,
           projectId: snapshot.projectId,