rework suspend completion

Author: nicktrn

apps/supervisor/package.json

@@ -6,7 +6,7 @@
   "type": "module",
   "scripts": {
     "build": "tsc",
-    "dev": "tsx --experimental-sqlite --require dotenv/config --watch src/index.ts",
+    "dev": "tsx --experimental-sqlite --require dotenv/config --watch src/index.ts || (echo '!! Remember to run: nvm use'; exit 1)",
     "start": "node --experimental-sqlite dist/index.js",
     "typecheck": "tsc --noEmit"
   },

apps/supervisor/src/workloadServer/index.ts

@@ -249,8 +249,6 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
               console.error("Failed to suspend run", { params });
               return;
             }
-
-            console.log("Suspended run", { params });
           },
         }
       )

packages/core/src/v3/runEngineWorker/supervisor/http.ts

@@ -15,6 +15,8 @@ import {
   WorkerApiRunHeartbeatResponseBody,
   WorkerApiRunLatestSnapshotResponseBody,
   WorkerApiDebugLogBody,
+  WorkerApiSuspendRunRequestBody,
+  WorkerApiSuspendRunResponseBody,
 } from "./schemas.js";
 import { SupervisorClientCommonOptions } from "./types.js";
 import { getDefaultWorkerHeaders } from "./util.js";
@@ -220,14 +222,30 @@ export class SupervisorHttpClient {
     );
   }
 
-  getSuspendCompletionUrl(runId: string, snapshotId: string, runnerId?: string) {
-    return {
-      url: `${this.apiUrl}/engine/v1/worker-actions/runs/${runId}/snapshots/${snapshotId}/suspend`,
-      headers: {
-        ...this.defaultHeaders,
-        ...this.runnerIdHeader(runnerId),
-      },
-    };
+  async submitSuspendCompletion({
+    runId,
+    snapshotId,
+    runnerId,
+    body,
+  }: {
+    runId: string;
+    snapshotId: string;
+    runnerId?: string;
+    body: WorkerApiSuspendRunRequestBody;
+  }) {
+    return wrapZodFetch(
+      WorkerApiSuspendRunResponseBody,
+      `${this.apiUrl}/engine/v1/worker-actions/runs/${runId}/snapshots/${snapshotId}/suspend`,
+      {
+        method: "POST",
+        headers: {
+          ...this.defaultHeaders,
+          ...this.runnerIdHeader(runnerId),
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(body),
+      }
+    );
   }
 
   private runnerIdHeader(runnerId?: string): Record<string, string> {

packages/core/src/v3/runEngineWorker/supervisor/schemas.ts

@@ -142,3 +142,11 @@ export const WorkerApiDebugLogBody = z.object({
   properties: Attributes.optional(),
 });
 export type WorkerApiDebugLogBody = z.infer<typeof WorkerApiDebugLogBody>;
+
+export const WorkerApiSuspendCompletionResponseBody = z.object({
+  success: z.boolean(),
+  error: z.string().optional(),
+});
+export type WorkerApiSuspendCompletionResponseBody = z.infer<
+  typeof WorkerApiSuspendCompletionResponseBody
+>;

packages/core/src/v3/serverOnly/checkpointClient.ts

@@ -21,10 +21,6 @@ export class CheckpointClient {
     return new URL("/api/v1/restore", this.apiUrl);
   }
 
-  private get suspendUrl() {
-    return new URL("/api/v1/suspend", this.apiUrl);
-  }
-
   constructor(opts: CheckpointClientOptions) {
     this.apiUrl = opts.apiUrl;
     this.workerClient = opts.workerClient;
@@ -41,26 +37,19 @@ export class CheckpointClient {
     containerId: string;
     runnerId: string;
   }): Promise<boolean> {
-    const completionUrl = this.workerClient.getSuspendCompletionUrl(
-      runFriendlyId,
-      snapshotFriendlyId,
-      runnerId
-    );
-
-    const res = await fetch(this.suspendUrl, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify({
-        type: "DOCKER",
-        containerId,
-        callbacks: {
-          completion: completionUrl.url,
-          headers: completionUrl.headers,
+    const res = await fetch(
+      new URL(`/api/v1/runs/${runFriendlyId}/snapshots/${snapshotFriendlyId}/suspend`, this.apiUrl),
+      {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
         },
-      } satisfies CheckpointServiceSuspendRequestBodyInput),
-    });
+        body: JSON.stringify({
+          type: "DOCKER",
+          containerId,
+        } satisfies CheckpointServiceSuspendRequestBodyInput),
+      }
+    );
 
     if (!res.ok) {
       this.logger.error("[CheckpointClient] Suspend request failed", {