feat(run-engine): report a failed restore through a fenced worker action

Add a worker action that lets a supervisor report a failed checkpoint
restore, so the run is requeued or failed promptly instead of waiting
out its stall timeouts.

Mono-RevId: 2a28ec932fcd5256c60369d14d7f793693d19142

Author: 1stvamp

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.restore-outcome.ts

@@ -0,0 +1,80 @@
+import type { TypedResponse } from "@remix-run/server-runtime";
+import { json } from "@remix-run/server-runtime";
+import type { WorkerApiRunRestoreOutcomeResponseBody } from "@trigger.dev/core/v3/workers";
+import { WorkerApiRunRestoreOutcomeRequestBody } from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
+import { logger } from "~/services/logger.server";
+import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+
+export const action = createActionWorkerApiRoute(
+  {
+    params: z.object({
+      runFriendlyId: z.string(),
+      snapshotFriendlyId: z.string(),
+    }),
+    body: z.compile(WorkerApiRunRestoreOutcomeRequestBody),
+  },
+  async ({
+    authenticatedWorker,
+    params,
+    body,
+    runnerId,
+    environmentId,
+  }): Promise<TypedResponse<WorkerApiRunRestoreOutcomeResponseBody>> => {
+    const { runFriendlyId, snapshotFriendlyId } = params;
+
+    logger.debug("Reporting restore outcome", { runFriendlyId, snapshotFriendlyId, body });
+
+    let result: Awaited<ReturnType<typeof authenticatedWorker.reportRestoreOutcome>>;
+    try {
+      result = await authenticatedWorker.reportRestoreOutcome({
+        runFriendlyId,
+        snapshotFriendlyId,
+        outcome: body.outcome,
+        reason: body.reason,
+        message: body.message,
+        snapshotRoute: body.snapshotRoute,
+        runnerId,
+        environmentId,
+      });
+    } catch (error) {
+      if (error instanceof Response) {
+        throw error;
+      }
+
+      logger.warn("Failed to report restore outcome", {
+        runFriendlyId,
+        snapshotFriendlyId,
+        environmentId,
+        outcome: body.outcome,
+        reason: body.reason,
+        error,
+      });
+
+      // Left to the generic 500 handler, which the worker's client retries.
+      throw error;
+    }
+
+    if (!result.ok) {
+      logger.info("Restore outcome not applied, the run has moved on", {
+        runFriendlyId,
+        snapshotFriendlyId,
+        environmentId,
+        outcome: body.outcome,
+        reason: body.reason,
+        latestExecutionStatus: result.latestExecutionStatus,
+      });
+
+      // The client retries a 409 by default; a stale report will never apply, so say not to.
+      throw json(
+        {
+          error: "Snapshot is no longer the run's restore snapshot",
+          latestExecutionStatus: result.latestExecutionStatus,
+        },
+        { status: 409, headers: { "x-should-retry": "false" } }
+      );
+    }
+
+    return json({ ok: true });
+  }
+);

apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts

@@ -687,6 +687,39 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
     });
   }
 
+  async reportRestoreOutcome({
+    runFriendlyId,
+    snapshotFriendlyId,
+    outcome,
+    reason,
+    message,
+    runnerId,
+    environmentId,
+    snapshotRoute,
+  }: {
+    runFriendlyId: string;
+    snapshotFriendlyId: string;
+    outcome: "requeue" | "fail";
+    reason: string;
+    message?: string;
+    runnerId?: string;
+    environmentId?: string;
+    // Carried back from the restore DequeuedMessage on the supervisor's report; honors durable residency.
+    snapshotRoute?: SnapshotRouteWire;
+  }) {
+    return await this._engine.reportRestoreOutcome({
+      runId: fromFriendlyId(runFriendlyId),
+      snapshotId: fromFriendlyId(snapshotFriendlyId),
+      outcome,
+      reason,
+      message,
+      workerId: this.workerInstanceId,
+      runnerId,
+      environmentId,
+      snapshotRoute,
+    });
+  }
+
   async getSnapshotsSince({
     runFriendlyId,
     snapshotId,

internal-packages/run-engine/src/engine/index.ts

@@ -105,6 +105,8 @@ import type {
   EngineWorker,
   HeartbeatTimeouts,
   ReportableQueue,
+  RestoreOutcome,
+  RestoreOutcomeResult,
   RunEngineOptions,
   TriggerParams,
   WorkerQueueDequeueOptions,
@@ -2362,6 +2364,46 @@ export class RunEngine {
     });
   }
 
+  /**
+   * Called by the worker that dequeued a run for a restore when the restore did not start.
+   * `requeue` gives the run back with the same checkpoint; `fail` fails the attempt as a lost
+   * checkpoint. A conflict is returned if `snapshotId` is no longer the restore's snapshot.
+   */
+  async reportRestoreOutcome({
+    runId,
+    snapshotId,
+    outcome,
+    reason,
+    message,
+    workerId,
+    runnerId,
+    environmentId,
+    snapshotRoute,
+  }: {
+    runId: string;
+    snapshotId: string;
+    outcome: RestoreOutcome;
+    reason: string;
+    message?: string;
+    workerId?: string;
+    runnerId?: string;
+    environmentId?: string;
+    // Carried from the restore DequeuedMessage; resolved durably when absent.
+    snapshotRoute?: SnapshotRouteWire;
+  }): Promise<RestoreOutcomeResult> {
+    return this.runAttemptSystem.reportRestoreOutcome({
+      runId,
+      snapshotId,
+      outcome,
+      reason,
+      message,
+      workerId,
+      runnerId,
+      environmentId,
+      snapshotRoute,
+    });
+  }
+
   /**
   Send a heartbeat to signal the the run is still executing.
   If a heartbeat isn't received, after a while the run is considered "stalled"

internal-packages/run-engine/src/engine/retrying.ts

@@ -137,33 +137,12 @@ export async function retryOutcomeFromCompletion(
       return { outcome: "fail_run", sanitizedError };
     }
 
-    const retryConfig = run.lockedRetryConfig;
-
-    if (!retryConfig) {
-      return { outcome: "fail_run", sanitizedError };
-    }
-
-    const parsedRetryConfig = NullishRetryOptions.safeParse(retryConfig);
-
-    if (!parsedRetryConfig.success) {
-      return { outcome: "fail_run", sanitizedError };
-    }
-
-    if (!parsedRetryConfig.data) {
-      return { outcome: "fail_run", sanitizedError };
-    }
-
-    const nextDelay = calculateNextRetryDelay(parsedRetryConfig.data, attemptNumber ?? 1);
+    const retrySettings = retrySettingsFromLockedConfig(run.lockedRetryConfig, attemptNumber);
 
-    if (!nextDelay) {
+    if (!retrySettings) {
       return { outcome: "fail_run", sanitizedError };
     }
 
-    const retrySettings = {
-      timestamp: Date.now() + nextDelay,
-      delay: nextDelay,
-    };
-
     return {
       outcome: "retry",
       method: "queue", // we'll always retry on the queue because usually having no settings means something bad happened
@@ -184,6 +163,30 @@ export async function retryOutcomeFromCompletion(
   };
 }
 
+/** The next retry the run's locked retry config allows after `attemptNumber`, if any. */
+export function retrySettingsFromLockedConfig(
+  lockedRetryConfig: unknown,
+  attemptNumber: number | null
+): TaskRunExecutionRetry | undefined {
+  if (!lockedRetryConfig) {
+    return;
+  }
+
+  const parsedRetryConfig = NullishRetryOptions.safeParse(lockedRetryConfig);
+
+  if (!parsedRetryConfig.success || !parsedRetryConfig.data) {
+    return;
+  }
+
+  const nextDelay = calculateNextRetryDelay(parsedRetryConfig.data, attemptNumber ?? 1);
+
+  if (!nextDelay) {
+    return;
+  }
+
+  return { timestamp: Date.now() + nextDelay, delay: nextDelay };
+}
+
 async function retryOOMOnMachine(
   prisma: PrismaClientOrTransaction,
   runStore: RunStore,