feat(supervisor): make checkpoint request timeouts configurable

Author: nicktrn

apps/supervisor/src/env.ts

@@ -115,6 +115,8 @@ export const Env = z
     TRIGGER_WARM_START_URL: z.string().optional(),
     TRIGGER_WARM_START_DISPATCH_URL: z.string().optional(),
     TRIGGER_CHECKPOINT_URL: z.string().optional(),
+    TRIGGER_CHECKPOINT_TIMEOUT_MS: z.coerce.number().int().positive().default(5_000),
+    TRIGGER_CHECKPOINT_RESTORE_TIMEOUT_MS: z.coerce.number().int().positive().default(30_000),
     TRIGGER_METADATA_URL: z.string().optional(),
 
     // Warm-start delivery verification: after a warm-start hit, probe the

apps/supervisor/src/index.ts

@@ -382,6 +382,8 @@ class ManagedSupervisor {
         apiUrl: new URL(env.TRIGGER_CHECKPOINT_URL),
         workerClient: this.workerSession.httpClient,
         orchestrator: this.isKubernetes ? "KUBERNETES" : "DOCKER",
+        timeoutMs: env.TRIGGER_CHECKPOINT_TIMEOUT_MS,
+        restoreTimeoutMs: env.TRIGGER_CHECKPOINT_RESTORE_TIMEOUT_MS,
       });
     }
 

apps/supervisor/src/workloadServer/index.ts

@@ -290,12 +290,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
     checkpointDeleteRequests.inc({ result: "sent" });
   }
 
-  /**
-   * cancelCheckpoints tells the checkpoint service a resumed run's in-flight checkpoint is moot.
-   * Must be called after the reply is sent: it never delays the runner. Without it the checkpoint
-   * is abandoned by a periodic snapshot poll instead, up to that interval later.
-   */
-  private async cancelCheckpoints(runFriendlyId: string): Promise<void> {
+  private async cancelCheckpointsAfterReply(runFriendlyId: string): Promise<void> {
     if (!this.checkpointClient) {
       checkpointCancelRequests.inc({ result: "no_client" });
       return;
@@ -686,7 +681,7 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
 
                 reply.json(continuationResult.data as WorkloadContinueRunExecutionResponseBody);
 
-                await this.cancelCheckpoints(params.runFriendlyId);
+                await this.cancelCheckpointsAfterReply(params.runFriendlyId);
               }
             ),
         }

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

@@ -11,15 +11,22 @@ export type CheckpointClientOptions = {
   apiUrl: URL;
   workerClient: SupervisorHttpClient;
   orchestrator: CheckpointType;
+  timeoutMs?: number;
+  restoreTimeoutMs?: number;
 };
 
-const CANCEL_TIMEOUT_MS = 5_000;
+const DEFAULT_TIMEOUT_MS = 5_000;
+const DEFAULT_RESTORE_TIMEOUT_MS = 30_000;
 
 export class CheckpointClient {
   private readonly logger = new SimpleStructuredLogger("checkpoint-client");
 
   constructor(private readonly opts: CheckpointClientOptions) {}
 
+  private timeout(ms = this.opts.timeoutMs ?? DEFAULT_TIMEOUT_MS): AbortSignal {
+    return AbortSignal.timeout(ms);
+  }
+
   async suspendRun({
     runFriendlyId,
     snapshotFriendlyId,
@@ -43,6 +50,7 @@ export class CheckpointClient {
           type: this.opts.orchestrator,
           ...body,
         } satisfies CheckpointServiceSuspendRequestBodyInput),
+        signal: this.timeout(),
       }
     );
 
@@ -107,6 +115,7 @@ export class CheckpointClient {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(body),
+        signal: this.timeout(this.opts.restoreTimeoutMs ?? DEFAULT_RESTORE_TIMEOUT_MS),
       }
     );
 
@@ -146,6 +155,7 @@ export class CheckpointClient {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(body),
+        signal: this.timeout(),
       }
     );
 
@@ -160,19 +170,14 @@ export class CheckpointClient {
     return true;
   }
 
-  /**
-   * cancelCheckpoints returns "unsupported" when the route is absent, which is expected while a
-   * newer caller runs against an older checkpoint service. The route answers 202 even when the run
-   * has nothing in flight, so a 404 only ever means the route itself is missing.
-   */
   async cancelCheckpoints({
     runFriendlyId,
   }: {
     runFriendlyId: string;
   }): Promise<"ok" | "unsupported" | "failed"> {
     const res = await fetch(
       new URL(`/api/v1/runs/${runFriendlyId}/checkpoints/cancel`, this.opts.apiUrl),
-      { method: "POST", signal: AbortSignal.timeout(CANCEL_TIMEOUT_MS) }
+      { method: "POST", signal: this.timeout() }
     );
 
     if (res.status === 404) {