manual checkpoints

nicktrn · nicktrn · commit 113413bf99ad · 2025-02-13T16:05:09.000Z
diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts
@@ -1492,6 +1492,95 @@ class TaskCoordinator {
             });
           }
         }
+        case "/checkpoint/manual": {
+          try {
+            const body = await getTextBody(req);
+            const json = safeJsonParse(body);
+
+            if (typeof json !== "object" || !json) {
+              return reply.text("Invalid body", 400);
+            }
+
+            if (!("runId" in json) || typeof json.runId !== "string") {
+              return reply.text("Missing or invalid: runId", 400);
+            }
+
+            let restoreAtUnixTimeMs: number | undefined;
+            if ("restoreAtUnixTimeMs" in json && typeof json.restoreAtUnixTimeMs === "number") {
+              restoreAtUnixTimeMs = json.restoreAtUnixTimeMs;
+            }
+
+            let keepRunAlive = false;
+            if ("keepRunAlive" in json && typeof json.keepRunAlive === "boolean") {
+              keepRunAlive = json.keepRunAlive;
+            }
+
+            const { runId } = json;
+
+            if (!runId) {
+              return reply.text("Missing runId", 400);
+            }
+
+            const runSocket = await this.#getRunSocket(runId);
+            if (!runSocket) {
+              return reply.text("Run socket not found", 404);
+            }
+
+            const { data } = runSocket;
+
+            console.log("Manual checkpoint", data);
+
+            const checkpoint = await this.#checkpointer.checkpointAndPush({
+              runId: data.runId,
+              projectRef: data.projectRef,
+              deploymentVersion: data.deploymentVersion,
+              attemptNumber: data.attemptNumber ? parseInt(data.attemptNumber) : undefined,
+            });
+
+            if (!checkpoint) {
+              return reply.text("Failed to checkpoint", 500);
+            }
+
+            if (!data.attemptFriendlyId) {
+              return reply.text("Socket data missing attemptFriendlyId", 500);
+            }
+
+            const ack = await this.#platformSocket?.sendWithAck("CHECKPOINT_CREATED", {
+              version: "v1",
+              runId,
+              attemptFriendlyId: data.attemptFriendlyId,
+              docker: checkpoint.docker,
+              location: checkpoint.location,
+              reason: {
+                type: "MANUAL",
+                restoreAtUnixTimeMs,
+              },
+            });
+
+            if (ack?.keepRunAlive || keepRunAlive) {
+              return reply.json({
+                message: `keeping run ${runId} alive after checkpoint`,
+                checkpoint,
+                requestJson: json,
+              });
+            }
+
+            runSocket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+
+            return reply.json({
+              message: `checkpoint created for run ${runId}`,
+              checkpoint,
+              requestJson: json,
+            });
+          } catch (error) {
+            return reply.json({
+              message: `error`,
+              error,
+            });
+          }
+        }
         default: {
           return reply.empty(404);
         }
diff --git a/apps/webapp/app/database-types.ts b/apps/webapp/app/database-types.ts
@@ -69,3 +69,11 @@ export const RuntimeEnvironmentType = {
   DEVELOPMENT: "DEVELOPMENT",
   PREVIEW: "PREVIEW",
 } as const satisfies Record<RuntimeEnvironmentTypeType, RuntimeEnvironmentTypeType>;
+
+export function isTaskRunAttemptStatus(value: string): value is keyof typeof TaskRunAttemptStatus {
+  return Object.values(TaskRunAttemptStatus).includes(value as keyof typeof TaskRunAttemptStatus);
+}
+
+export function isTaskRunStatus(value: string): value is keyof typeof TaskRunStatus {
+  return Object.values(TaskRunStatus).includes(value as keyof typeof TaskRunStatus);
+}
diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts
@@ -1,4 +1,4 @@
-import { CoordinatorToPlatformMessages } from "@trigger.dev/core/v3";
+import { CoordinatorToPlatformMessages, ManualCheckpointMetadata } from "@trigger.dev/core/v3";
 import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket";
 import type { Checkpoint, CheckpointRestoreEvent } from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
@@ -101,6 +101,19 @@ export class CreateCheckpointService extends BaseService {
     //   setTimeout(resolve, waitSeconds * 1000);
     // });
 
+    let metadata: string;
+
+    if (params.reason.type === "MANUAL") {
+      metadata = JSON.stringify({
+        ...params.reason,
+        attemptId: attempt.id,
+        previousAttemptStatus: attempt.status,
+        previousRunStatus: attempt.taskRun.status,
+      } satisfies ManualCheckpointMetadata);
+    } else {
+      metadata = JSON.stringify(params.reason);
+    }
+
     const checkpoint = await this._prisma.checkpoint.create({
       data: {
         friendlyId: generateFriendlyId("checkpoint"),
@@ -112,7 +125,7 @@ export class CreateCheckpointService extends BaseService {
         location: params.location,
         type: params.docker ? "DOCKER" : "KUBERNETES",
         reason: params.reason.type,
-        metadata: JSON.stringify(params.reason),
+        metadata,
         imageRef,
       },
     });
@@ -138,7 +151,17 @@ export class CreateCheckpointService extends BaseService {
     let checkpointEvent: CheckpointRestoreEvent | undefined;
 
     switch (reason.type) {
+      case "MANUAL":
       case "WAIT_FOR_DURATION": {
+        let restoreAtUnixTimeMs: number;
+
+        if (reason.type === "MANUAL") {
+          // Restore immediately if not specified, useful for live migration
+          restoreAtUnixTimeMs = reason.restoreAtUnixTimeMs ?? Date.now();
+        } else {
+          restoreAtUnixTimeMs = reason.now + reason.ms;
+        }
+
         checkpointEvent = await eventService.checkpoint({
           checkpointId: checkpoint.id,
         });
@@ -151,7 +174,7 @@ export class CreateCheckpointService extends BaseService {
               resumableAttemptId: attempt.id,
               checkpointEventId: checkpointEvent.id,
             },
-            reason.now + reason.ms
+            restoreAtUnixTimeMs
           );
 
           return {
diff --git a/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts b/apps/webapp/app/v3/services/createCheckpointRestoreEvent.server.ts
@@ -1,6 +1,13 @@
-import type { CheckpointRestoreEvent, CheckpointRestoreEventType } from "@trigger.dev/database";
+import type {
+  Checkpoint,
+  CheckpointRestoreEvent,
+  CheckpointRestoreEventType,
+} from "@trigger.dev/database";
 import { logger } from "~/services/logger.server";
 import { BaseService } from "./baseService.server";
+import { ManualCheckpointMetadata } from "@trigger.dev/core/v3";
+import { isTaskRunAttemptStatus, isTaskRunStatus, TaskRunAttemptStatus } from "~/database-types";
+import { safeJsonParse } from "~/utils/json";
 
 interface CheckpointRestoreEventCallParams {
   checkpointId: string;
@@ -39,6 +46,13 @@ export class CreateCheckpointRestoreEventService extends BaseService {
       return;
     }
 
+    if (params.type === "RESTORE" && checkpoint.reason === "MANUAL") {
+      const manualRestoreSuccess = await this.#handleManualCheckpointRestore(checkpoint);
+      if (!manualRestoreSuccess) {
+        return;
+      }
+    }
+
     logger.debug(`Creating checkpoint/restore event`, { params });
 
     let taskRunDependencyId: string | undefined;
@@ -99,4 +113,81 @@ export class CreateCheckpointRestoreEventService extends BaseService {
 
     return checkpointEvent;
   }
+
+  async #handleManualCheckpointRestore(checkpoint: Checkpoint): Promise<boolean> {
+    const json = checkpoint.metadata ? safeJsonParse(checkpoint.metadata) : undefined;
+
+    // We need to restore the previous run and attempt status as saved in the metadata
+    const metadata = ManualCheckpointMetadata.safeParse(json);
+
+    if (!metadata.success) {
+      logger.error("Invalid metadata", { metadata });
+      return false;
+    }
+
+    const { attemptId, previousAttemptStatus, previousRunStatus } = metadata.data;
+
+    if (!isTaskRunAttemptStatus(previousAttemptStatus)) {
+      logger.error("Invalid previous attempt status", { previousAttemptStatus });
+      return false;
+    }
+
+    if (!isTaskRunStatus(previousRunStatus)) {
+      logger.error("Invalid previous run status", { previousRunStatus });
+      return false;
+    }
+
+    try {
+      const updatedAttempt = await this._prisma.taskRunAttempt.update({
+        where: {
+          id: attemptId,
+        },
+        data: {
+          status: previousAttemptStatus,
+          taskRun: {
+            update: {
+              data: {
+                status: previousRunStatus,
+              },
+            },
+          },
+        },
+        select: {
+          id: true,
+          status: true,
+          taskRun: {
+            select: {
+              id: true,
+              status: true,
+            },
+          },
+        },
+      });
+
+      logger.debug("Set post resume statuses after manual checkpoint", {
+        run: {
+          id: updatedAttempt.taskRun.id,
+          status: updatedAttempt.taskRun.status,
+        },
+        attempt: {
+          id: updatedAttempt.id,
+          status: updatedAttempt.status,
+        },
+      });
+
+      return true;
+    } catch (error) {
+      logger.error("Failed to set post resume statuses", {
+        error:
+          error instanceof Error
+            ? {
+                name: error.name,
+                message: error.message,
+                stack: error.stack,
+              }
+            : error,
+      });
+      return false;
+    }
+  }
 }
diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts
@@ -479,6 +479,11 @@ export const CoordinatorToPlatformMessages = {
           type: z.literal("RETRYING_AFTER_FAILURE"),
           attemptNumber: z.number(),
         }),
+        z.object({
+          type: z.literal("MANUAL"),
+          /** If unspecified it will be restored immediately, e.g. for live migration */
+          restoreAtUnixTimeMs: z.number().optional(),
+        }),
       ]),
     }),
     callback: z.object({
diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts
@@ -250,3 +250,12 @@ export const TaskRunExecutionLazyAttemptPayload = z.object({
 });
 
 export type TaskRunExecutionLazyAttemptPayload = z.infer<typeof TaskRunExecutionLazyAttemptPayload>;
+
+export const ManualCheckpointMetadata = z.object({
+  /** NOT a friendly ID */
+  attemptId: z.string(),
+  previousRunStatus: z.string(),
+  previousAttemptStatus: z.string(),
+});
+
+export type ManualCheckpointMetadata = z.infer<typeof ManualCheckpointMetadata>;
