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

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

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,6 +22,8 @@ client.defineJob({
       await io.runTask(
         `task-${i}`,
         async (task) => {
+          await new Promise((resolve) => setTimeout(resolve, 2000));
+
           return {
             output: "a".repeat(300 * 1024),
           };