failed run service now respects retries

Author: nicktrn

apps/webapp/app/v3/failedTaskRun.server.ts

@@ -1,22 +1,47 @@
-import { sanitizeError, TaskRunFailedExecutionResult } from "@trigger.dev/core/v3";
+import {
+  calculateNextRetryDelay,
+  RetryOptions,
+  sanitizeError,
+  TaskRunExecution,
+  TaskRunExecutionRetry,
+  TaskRunFailedExecutionResult,
+} from "@trigger.dev/core/v3";
 import { logger } from "~/services/logger.server";
 import { createExceptionPropertiesFromError, eventRepository } from "./eventRepository.server";
 import { BaseService } from "./services/baseService.server";
 import { FinalizeTaskRunService } from "./services/finalizeTaskRun.server";
-import { FAILABLE_RUN_STATUSES } from "./taskStatus";
+import { isFailableRunStatus, isFinalAttemptStatus } from "./taskStatus";
+import { Prisma } from "@trigger.dev/database";
+import { CompleteAttemptService } from "./services/completeAttempt.server";
+import { CreateTaskRunAttemptService } from "./services/createTaskRunAttempt.server";
+import { sharedQueueTasks } from "./marqs/sharedQueueConsumer.server";
+
+const includeAttempts = {
+  attempts: {
+    orderBy: {
+      createdAt: "desc",
+    },
+    take: 1,
+  },
+  lockedBy: true,
+} satisfies Prisma.TaskRunInclude;
+
+type TaskRunWithAttempts = Prisma.TaskRunGetPayload<{
+  include: typeof includeAttempts;
+}>;
 
-/**
- * 
- */
 export class FailedTaskRunService extends BaseService {
   public async call(anyRunId: string, completion: TaskRunFailedExecutionResult) {
+    logger.debug("[FailedTaskRunService] Handling failed task run", { anyRunId, completion });
+
     const isFriendlyId = anyRunId.startsWith("run_");
 
     const taskRun = await this._prisma.taskRun.findUnique({
       where: {
         friendlyId: isFriendlyId ? anyRunId : undefined,
         id: !isFriendlyId ? anyRunId : undefined,
       },
+      include: includeAttempts,
     });
 
     if (!taskRun) {
@@ -28,7 +53,7 @@ export class FailedTaskRunService extends BaseService {
       return;
     }
 
-    if (!FAILABLE_RUN_STATUSES.includes(taskRun.status)) {
+    if (!isFailableRunStatus(taskRun.status)) {
       logger.error("[FailedTaskRunService] Task run is not in a failable state", {
         taskRun,
         completion,
@@ -37,7 +62,28 @@ export class FailedTaskRunService extends BaseService {
       return;
     }
 
-    // No more retries, we need to fail the task run
+    const retriableExecution = await this.#getRetriableAttemptExecution(taskRun, completion);
+
+    if (retriableExecution) {
+      logger.debug("[FailedTaskRunService] Completing attempt", { taskRun, completion });
+
+      const executionRetry =
+        completion.retry ?? (await this.#getExecutionRetry(taskRun, retriableExecution));
+
+      const completeAttempt = new CompleteAttemptService(this._prisma);
+      await completeAttempt.call({
+        completion: {
+          ...completion,
+          retry: executionRetry,
+        },
+        execution: retriableExecution,
+        isSystemFailure: true,
+      });
+
+      return;
+    }
+
+    // No retriable execution, so we need to fail the task run
     logger.debug("[FailedTaskRunService] Failing task run", { taskRun, completion });
 
     const finalizeService = new FinalizeTaskRunService();
@@ -66,4 +112,87 @@ export class FailedTaskRunService extends BaseService {
       ],
     });
   }
+
+  async #getRetriableAttemptExecution(
+    run: TaskRunWithAttempts,
+    completion: TaskRunFailedExecutionResult
+  ): Promise<TaskRunExecution | undefined> {
+    let attempt = run.attempts[0];
+
+    // We need to create an attempt if:
+    // - None exists yet
+    // - The last attempt has a final status, e.g. we failed between attempts
+    if (!attempt || isFinalAttemptStatus(attempt.status)) {
+      logger.error("[FailedTaskRunService] No attempts found", {
+        run,
+        completion,
+      });
+
+      const createAttempt = new CreateTaskRunAttemptService(this._prisma);
+
+      try {
+        const { execution } = await createAttempt.call(run.id);
+        return execution;
+      } catch (error) {
+        logger.error("[FailedTaskRunService] Failed to create attempt", {
+          run,
+          completion,
+          error,
+        });
+
+        return;
+      }
+    }
+
+    // We already have an attempt with non-final status, let's use it
+    try {
+      const executionPayload = await sharedQueueTasks.getExecutionPayloadFromAttempt(
+        attempt.id,
+        undefined,
+        undefined,
+        true
+      );
+      return executionPayload?.execution;
+    } catch (error) {
+      logger.error("[FailedTaskRunService] Failed to get execution payload", {
+        run,
+        completion,
+        error,
+      });
+
+      return;
+    }
+  }
+
+  async #getExecutionRetry(
+    run: TaskRunWithAttempts,
+    execution: TaskRunExecution
+  ): Promise<TaskRunExecutionRetry | undefined> {
+    const parsedRetryConfig = RetryOptions.safeParse(run.lockedBy?.retryConfig);
+
+    if (!parsedRetryConfig.success) {
+      logger.error("[FailedTaskRunService] Invalid retry config", {
+        run,
+        execution,
+      });
+
+      return;
+    }
+
+    const delay = calculateNextRetryDelay(parsedRetryConfig.data, execution.attempt.number);
+
+    if (!delay) {
+      logger.debug("[FailedTaskRunService] No more retries", {
+        run,
+        execution,
+      });
+
+      return;
+    }
+
+    return {
+      timestamp: Date.now() + delay,
+      delay,
+    };
+  }
 }

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -945,7 +945,8 @@ class SharedQueueTasks {
   async getExecutionPayloadFromAttempt(
     id: string,
     setToExecuting?: boolean,
-    isRetrying?: boolean
+    isRetrying?: boolean,
+    skipStatusChecks?: boolean
   ): Promise<ProdTaskRunExecutionPayload | undefined> {
     const attempt = await prisma.taskRunAttempt.findUnique({
       where: {
@@ -979,27 +980,29 @@ class SharedQueueTasks {
       return;
     }
 
-    switch (attempt.status) {
-      case "CANCELED":
-      case "EXECUTING": {
-        logger.error("Invalid attempt status for execution payload retrieval", {
-          attemptId: id,
-          status: attempt.status,
-        });
-        return;
+    if (!skipStatusChecks) {
+      switch (attempt.status) {
+        case "CANCELED":
+        case "EXECUTING": {
+          logger.error("Invalid attempt status for execution payload retrieval", {
+            attemptId: id,
+            status: attempt.status,
+          });
+          return;
+        }
       }
-    }
 
-    switch (attempt.taskRun.status) {
-      case "CANCELED":
-      case "EXECUTING":
-      case "INTERRUPTED": {
-        logger.error("Invalid run status for execution payload retrieval", {
-          attemptId: id,
-          runId: attempt.taskRunId,
-          status: attempt.taskRun.status,
-        });
-        return;
+      switch (attempt.taskRun.status) {
+        case "CANCELED":
+        case "EXECUTING":
+        case "INTERRUPTED": {
+          logger.error("Invalid run status for execution payload retrieval", {
+            attemptId: id,
+            runId: attempt.taskRunId,
+            status: attempt.taskRun.status,
+          });
+          return;
+        }
       }
     }
 

apps/webapp/app/v3/services/completeAttempt.server.ts

@@ -39,12 +39,14 @@ export class CompleteAttemptService extends BaseService {
     env,
     checkpoint,
     supportsRetryCheckpoints,
+    isSystemFailure,
   }: {
     completion: TaskRunExecutionResult;
     execution: TaskRunExecution;
     env?: AuthenticatedEnvironment;
     checkpoint?: CheckpointData;
     supportsRetryCheckpoints?: boolean;
+    isSystemFailure?: boolean;
   }): Promise<"COMPLETED" | "RETRIED"> {
     const taskRunAttempt = await findAttempt(this._prisma, execution.attempt.id);
 
@@ -110,6 +112,7 @@ export class CompleteAttemptService extends BaseService {
         env,
         checkpoint,
         supportsRetryCheckpoints,
+        isSystemFailure,
       });
     }
   }
@@ -170,13 +173,15 @@ export class CompleteAttemptService extends BaseService {
     env,
     checkpoint,
     supportsRetryCheckpoints,
+    isSystemFailure,
   }: {
     completion: TaskRunFailedExecutionResult;
     execution: TaskRunExecution;
     taskRunAttempt: NonNullable<FoundAttempt>;
     env?: AuthenticatedEnvironment;
     checkpoint?: CheckpointData;
     supportsRetryCheckpoints?: boolean;
+    isSystemFailure?: boolean;
   }): Promise<"COMPLETED" | "RETRIED"> {
     if (
       completion.error.type === "INTERNAL_ERROR" &&
@@ -199,14 +204,6 @@ export class CompleteAttemptService extends BaseService {
 
     const sanitizedError = sanitizeError(completion.error);
 
-    // TODO: make this handle the case where the current attempt is unknown, with only a run id
-
-    // 1. Get the task run
-
-    // 2. Get the most recent attempt
-
-    // 3. Get the retry config
-
     await this._prisma.taskRunAttempt.update({
       where: { id: taskRunAttempt.id },
       data: {
@@ -258,10 +255,11 @@ export class CompleteAttemptService extends BaseService {
       },
     });
 
-    const status =
-      sanitizedError.type === "INTERNAL_ERROR" && sanitizedError.code === "MAX_DURATION_EXCEEDED"
-        ? "TIMED_OUT"
-        : "COMPLETED_WITH_ERRORS";
+    const status = isSystemFailure
+      ? "SYSTEM_FAILURE"
+      : sanitizedError.type === "INTERNAL_ERROR" && sanitizedError.code === "MAX_DURATION_EXCEEDED"
+      ? "TIMED_OUT"
+      : "COMPLETED_WITH_ERRORS";
 
     const finalizeService = new FinalizeTaskRunService();
     await finalizeService.call({

apps/webapp/app/v3/taskStatus.ts

@@ -88,6 +88,10 @@ export function isCrashableAttemptStatus(status: TaskRunAttemptStatus): boolean
   return CRASHABLE_ATTEMPT_STATUSES.includes(status);
 }
 
+export function isFailableRunStatus(status: TaskRunStatus): boolean {
+  return FAILABLE_RUN_STATUSES.includes(status);
+}
+
 export function isFreezableRunStatus(status: TaskRunStatus): boolean {
   return FREEZABLE_RUN_STATUSES.includes(status);
 }