v2: When a run hits the rate limit reschedule the re-execution (#1125)

* Fix: API rate limit error has the correct seconds until reset

* When a v2 run hits the rate limit, reschedule using the reset timestamp

* Still throw AutoYieldRateLimitErrors

* Reschedule runs from the rate limit

* The stress test timeout should be inside the task

* If the rate limit error is thrown, don’t retry the API request

Author: matt-aitken

.changeset/strong-owls-know.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core": patch
+---
+
+When a v2 run hits the rate limit, reschedule with the reset date

apps/webapp/app/services/apiRateLimit.server.ts

@@ -157,16 +157,18 @@ export function authorizationRateLimitMiddleware({
     }
 
     res.setHeader("Content-Type", "application/problem+json");
+    const secondsUntilReset = Math.max(0, (reset - new Date().getTime()) / 1000);
     return res.status(429).send(
       JSON.stringify(
         {
           title: "Rate Limit Exceeded",
           status: 429,
           type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429",
-          detail: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry after ${reset} seconds.`,
-          reset: reset,
-          limit: limit,
-          error: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry after ${reset} seconds.`,
+          detail: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
+          reset,
+          limit,
+          secondsUntilReset,
+          error: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
         },
         null,
         2

apps/webapp/app/services/runs/performRunExecutionV3.server.ts

@@ -441,6 +441,10 @@ export class PerformRunExecutionV3Service {
           await this.#resumeAutoYieldedRunWithCompletedTask(run, safeBody.data, durationInMs);
           break;
         }
+        case "AUTO_YIELD_RATE_LIMIT": {
+          await this.#rescheduleRun(run, safeBody.data.reset, durationInMs);
+          break;
+        }
         case "RESUME_WITH_PARALLEL_TASK": {
           await this.#resumeParallelRunWithTask(run, safeBody.data, durationInMs);
 
@@ -667,6 +671,10 @@ export class PerformRunExecutionV3Service {
 
           break;
         }
+        case "AUTO_YIELD_RATE_LIMIT": {
+          await this.#rescheduleRun(run, childError.reset, durationInMs);
+          break;
+        }
         case "CANCELED": {
           break;
         }
@@ -801,9 +809,9 @@ export class PerformRunExecutionV3Service {
     });
   }
 
-  async #resumeAutoYieldedRun(
+  async #rescheduleRun(
     run: FoundRun,
-    data: AutoYieldMetadata,
+    reset: number,
     durationInMs: number,
     executionCount: number = 1
   ) {
@@ -820,24 +828,14 @@ export class PerformRunExecutionV3Service {
           executionCount: {
             increment: executionCount,
           },
-          autoYieldExecution: {
-            create: [
-              {
-                location: data.location,
-                timeRemaining: data.timeRemaining,
-                timeElapsed: data.timeElapsed,
-                limit: data.limit ?? 0,
-              },
-            ],
-          },
           forceYieldImmediately: false,
         },
         select: {
           executionCount: true,
         },
       });
 
-      await ResumeRunService.enqueue(run, tx);
+      await ResumeRunService.enqueue(run, tx, new Date(reset));
     });
   }
 
@@ -888,6 +886,46 @@ export class PerformRunExecutionV3Service {
     });
   }
 
+  async #resumeAutoYieldedRun(
+    run: FoundRun,
+    data: AutoYieldMetadata,
+    durationInMs: number,
+    executionCount: number = 1
+  ) {
+    await $transaction(this.#prismaClient, async (tx) => {
+      await tx.jobRun.update({
+        where: {
+          id: run.id,
+        },
+        data: {
+          status: "WAITING_TO_EXECUTE",
+          executionDuration: {
+            increment: durationInMs,
+          },
+          executionCount: {
+            increment: executionCount,
+          },
+          autoYieldExecution: {
+            create: [
+              {
+                location: data.location,
+                timeRemaining: data.timeRemaining,
+                timeElapsed: data.timeElapsed,
+                limit: data.limit ?? 0,
+              },
+            ],
+          },
+          forceYieldImmediately: false,
+        },
+        select: {
+          executionCount: true,
+        },
+      });
+
+      await ResumeRunService.enqueue(run, tx);
+    });
+  }
+
   async #retryRunWithTask(
     run: FoundRun,
     data: RunJobRetryWithTask,

packages/core/src/schemas/api.ts

@@ -671,6 +671,13 @@ export type RunJobAutoYieldWithCompletedTaskExecutionError = z.infer<
   typeof RunJobAutoYieldWithCompletedTaskExecutionErrorSchema
 >;
 
+export const RunJobAutoYieldRateLimitErrorSchema = z.object({
+  status: z.literal("AUTO_YIELD_RATE_LIMIT"),
+  reset: z.coerce.number(),
+});
+
+export type RunJobAutoYieldRateLimitError = z.infer<typeof RunJobAutoYieldRateLimitErrorSchema>;
+
 export const RunJobInvalidPayloadErrorSchema = z.object({
   status: z.literal("INVALID_PAYLOAD"),
   errors: z.array(SchemaErrorSchema),
@@ -719,6 +726,7 @@ export const RunJobErrorResponseSchema = z.union([
   RunJobAutoYieldExecutionErrorSchema,
   RunJobAutoYieldWithCompletedTaskExecutionErrorSchema,
   RunJobYieldExecutionErrorSchema,
+  RunJobAutoYieldRateLimitErrorSchema,
   RunJobErrorSchema,
   RunJobUnresolvedAuthErrorSchema,
   RunJobInvalidPayloadErrorSchema,
@@ -741,6 +749,7 @@ export const RunJobResponseSchema = z.discriminatedUnion("status", [
   RunJobAutoYieldExecutionErrorSchema,
   RunJobAutoYieldWithCompletedTaskExecutionErrorSchema,
   RunJobYieldExecutionErrorSchema,
+  RunJobAutoYieldRateLimitErrorSchema,
   RunJobErrorSchema,
   RunJobUnresolvedAuthErrorSchema,
   RunJobInvalidPayloadErrorSchema,

packages/trigger-sdk/src/apiClient.ts

@@ -44,6 +44,7 @@ import { env } from "node:process";
 
 import { z } from "zod";
 import { KeyValueStoreClient } from "./store/keyValueStoreClient";
+import { AutoYieldRateLimitError } from "./errors";
 
 export type ApiClientOptions = {
   apiKey?: string;
@@ -818,6 +819,15 @@ async function zodfetchWithVersions<
       return;
     }
 
+    //rate limit, so we want to reschedule
+    if (response.status === 429) {
+      //unix timestamp in milliseconds
+      const retryAfter = response.headers.get("x-ratelimit-reset");
+      if (retryAfter) {
+        throw new AutoYieldRateLimitError(parseInt(retryAfter));
+      }
+    }
+
     if (response.status >= 400 && response.status < 500) {
       const rawBody = await safeResponseText(response);
       const body = safeJsonParse(rawBody);
@@ -894,7 +904,7 @@ async function zodfetchWithVersions<
       body: versionedSchema.parse(jsonBody),
     };
   } catch (error) {
-    if (error instanceof UnknownVersionError) {
+    if (error instanceof UnknownVersionError || error instanceof AutoYieldRateLimitError) {
       throw error;
     }
 
@@ -987,6 +997,15 @@ async function zodfetch<TResponseSchema extends z.ZodTypeAny, TOptional extends
       return;
     }
 
+    //rate limit, so we want to reschedule
+    if (response.status === 429) {
+      //unix timestamp in milliseconds
+      const retryAfter = response.headers.get("x-ratelimit-reset");
+      if (retryAfter) {
+        throw new AutoYieldRateLimitError(parseInt(retryAfter));
+      }
+    }
+
     if (response.status >= 400 && response.status < 500) {
       const body = await response.json();
 
@@ -1012,6 +1031,10 @@ async function zodfetch<TResponseSchema extends z.ZodTypeAny, TOptional extends
 
     return schema.parse(jsonBody);
   } catch (error) {
+    if (error instanceof AutoYieldRateLimitError) {
+      throw error;
+    }
+
     if (retryCount < MAX_RETRIES) {
       // retry with exponential backoff and jitter
       const delay = exponentialBackoff(retryCount + 1);

packages/trigger-sdk/src/errors.ts

@@ -45,6 +45,10 @@ export class AutoYieldWithCompletedTaskExecutionError {
   ) {}
 }
 
+export class AutoYieldRateLimitError {
+  constructor(public resetAtTimestamp: number) {}
+}
+
 export class ParsedPayloadSchemaError {
   constructor(public schemaErrors: SchemaError[]) {}
 }
@@ -56,6 +60,7 @@ export type TriggerInternalError =
   | YieldExecutionError
   | AutoYieldExecutionError
   | AutoYieldWithCompletedTaskExecutionError
+  | AutoYieldRateLimitError
   | ResumeWithParallelTaskError;
 
 /** Use this function if you're using a `try/catch` block to catch errors.
@@ -72,6 +77,7 @@ export function isTriggerError(err: unknown): err is TriggerInternalError {
     err instanceof YieldExecutionError ||
     err instanceof AutoYieldExecutionError ||
     err instanceof AutoYieldWithCompletedTaskExecutionError ||
+    err instanceof AutoYieldRateLimitError ||
     err instanceof ResumeWithParallelTaskError
   );
 }

packages/trigger-sdk/src/io.ts

@@ -28,6 +28,7 @@ import { webcrypto } from "node:crypto";
 import { ApiClient } from "./apiClient";
 import {
   AutoYieldExecutionError,
+  AutoYieldRateLimitError,
   AutoYieldWithCompletedTaskExecutionError,
   CanceledWithTaskError,
   ErrorWithTask,
@@ -1460,6 +1461,14 @@ export class IO {
         cachedTasksCursor: this._cachedTasksCursor,
       });
     } catch (error) {
+      if (error instanceof AutoYieldRateLimitError) {
+        this._logger.debug("AutoYieldRateLimitError", {
+          error,
+        });
+
+        throw error;
+      }
+
       return;
     }
   }

packages/trigger-sdk/src/triggerClient.ts

@@ -52,6 +52,7 @@ import { ApiClient } from "./apiClient";
 import { ConcurrencyLimit, ConcurrencyLimitOptions } from "./concurrencyLimit";
 import {
   AutoYieldExecutionError,
+  AutoYieldRateLimitError,
   AutoYieldWithCompletedTaskExecutionError,
   CanceledWithTaskError,
   ErrorWithTask,
@@ -1235,6 +1236,13 @@ export class TriggerClient {
       };
     }
 
+    if (error instanceof AutoYieldRateLimitError) {
+      return {
+        status: "AUTO_YIELD_RATE_LIMIT",
+        reset: error.resetAtTimestamp,
+      };
+    }
+
     if (error instanceof YieldExecutionError) {
       return { status: "YIELD_EXECUTION", key: error.key };
     }

references/job-catalog/src/stressTest.ts

@@ -22,14 +22,14 @@ client.defineJob({
       await io.runTask(
         `task-${i}`,
         async (task) => {
+          await new Promise((resolve) => setTimeout(resolve, 2000));
+
           return {
             output: "a".repeat(30),
           };
         },
         { name: `Task ${i}` }
       );
-
-      await new Promise((resolve) => setTimeout(resolve, 2000));
     }
 
     // Now do a wait for 5 seconds