Retry on ratelimit status

src/client/dlq.test.ts

@@ -36,7 +36,7 @@ describe("DLQ", () => {
 
       await sleep(10_000);
 
-      const dlqLogs = await client.dlq.listMessages();
+      const dlqLogs = await client.dlq.listMessages({ filter: { messageId: message.messageId } });
       expect(dlqLogs.messages.map((dlq) => dlq.messageId)).toContain(message.messageId);
     },
     { timeout: 20_000 }

src/client/error.ts

@@ -1,7 +1,7 @@
 import type { ChatRateLimit, RateLimit } from "./types";
 import type { FailureFunctionPayload, Step } from "./workflow/types";
 
-const RATELIMIT_STATUS = 429;
+export const RATELIMIT_STATUS = 429;
 
 /**
  * Result of 500 Internal Server Error

src/client/http.test.ts

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-magic-numbers */
-import { describe, test, expect } from "bun:test";
+import { describe, test, expect, spyOn } from "bun:test";
 import { Client } from "./client";
+import { MOCK_QSTASH_SERVER_URL, mockQStashServer } from "./workflow/test-utils";
 
 describe("http", () => {
   test("should terminate after sleeping 5 times", () => {
@@ -22,4 +23,53 @@ describe("http", () => {
     // if the Promise.race doesn't throw, that means the retries took longer than 4.5s
     expect(throws).toThrow("Was there a typo in the url or port?");
   });
+
+  test("should backoff for seconds in ratelimit", async () => {
+    const qstashToken = "my-token";
+    const retries = 3;
+
+    const spy = spyOn(console, "warn");
+
+    let callCount = 0;
+    const client = new Client({
+      baseUrl: MOCK_QSTASH_SERVER_URL,
+      token: qstashToken,
+      retry: {
+        retries,
+        ratelimitBackoff: () => {
+          callCount += 1;
+          return 250;
+        },
+      },
+    });
+
+    await mockQStashServer({
+      execute: async () => {
+        await client.publishJSON({
+          url: "https://requestcatcher.com",
+        });
+        expect(callCount).toBe(retries);
+      },
+      receivesRequest: {
+        method: "POST",
+        token: qstashToken,
+        url: "http://localhost:8080/v2/publish/https://requestcatcher.com",
+      },
+      responseFields: {
+        status: 429,
+        headers: {
+          "Burst-RateLimit-Limit": "100",
+          "Burst-RateLimit-Remaining": "0",
+          "Burst-RateLimit-Reset": "213123",
+        },
+        body: "ratelimited",
+      },
+    });
+
+    expect(callCount).toBe(retries);
+    expect(spy).toHaveBeenCalledTimes(3);
+    expect(spy).toHaveBeenLastCalledWith(
+      'QStash Ratelimit Exceeded. Retrying after 250 milliseconds. Exceeded burst rate limit. {"limit":"100","remaining":"0","reset":"213123"}'
+    );
+  });
 });

src/client/http.ts

@@ -4,6 +4,7 @@ import {
   QstashRatelimitError,
   QstashChatRatelimitError,
   QstashDailyRatelimitError,
+  RATELIMIT_STATUS,
 } from "./error";
 import type { BodyInit, HeadersInit, HTTPMethods, RequestOptions } from "./types";
 import type { ChatCompletionChunk } from "./llm/types";
@@ -75,6 +76,19 @@ export type RetryConfig =
        * ```
        */
       backoff?: (retryCount: number) => number;
+      /**
+       * A backoff function receives the current retry cound and returns a number in milliseconds to wait before retrying.
+       *
+       * Applied when the response has 429 status, indicating a ratelimit.
+       *
+       * Initial `lastBackoff` value is 0.
+       *
+       * @default
+       * ```ts
+       * ((lastBackoff) => Math.max(lastBackoff, Math.random() * 4000) + 1000),
+       * ```
+       */
+      ratelimitBackoff?: (lastBackoff: number) => number;
     };
 
 export type HttpClientConfig = {
@@ -93,6 +107,7 @@ export class HttpClient implements Requester {
   public retry: {
     attempts: number;
     backoff: (retryCount: number) => number;
+    ratelimitBackoff: (lastBackoff: number) => number;
   };
 
   public constructor(config: HttpClientConfig) {
@@ -106,10 +121,14 @@ export class HttpClient implements Requester {
         ? {
             attempts: 1,
             backoff: () => 0,
+            ratelimitBackoff: () => 0,
           }
         : {
             attempts: config.retry?.retries ?? 5,
             backoff: config.retry?.backoff ?? ((retryCount) => Math.exp(retryCount) * 50),
+            ratelimitBackoff:
+              config.retry?.ratelimitBackoff ??
+              ((lastBackoff) => Math.max(lastBackoff, Math.random() * 4000) + 1000),
           };
   }
 
@@ -172,23 +191,33 @@ export class HttpClient implements Requester {
 
     let response: Response | undefined = undefined;
     let error: Error | undefined = undefined;
+    let ratelimitBackoff = 0;
     for (let index = 0; index <= this.retry.attempts; index++) {
       try {
         response = await fetch(url.toString(), requestOptions);
+        await this.checkResponse(response);
         break;
       } catch (error_) {
         error = error_ as Error;
 
-        // Only sleep if this is not the last attempt
         if (index < this.retry.attempts) {
-          await new Promise((r) => setTimeout(r, this.retry.backoff(index)));
+          // Only sleep if this is not the last attempt
+
+          if (error instanceof QstashError && error.status === RATELIMIT_STATUS) {
+            ratelimitBackoff = this.retry.ratelimitBackoff(index);
+            console.warn(
+              `QStash Ratelimit Exceeded. Retrying after ${ratelimitBackoff} milliseconds. ${error.message}`
+            );
+            await new Promise((r) => setTimeout(r, ratelimitBackoff));
+          } else {
+            await new Promise((r) => setTimeout(r, this.retry.backoff(index)));
+          }
         }
       }
     }
     if (!response) {
       throw error ?? new Error("Exhausted all retries");
     }
-    await this.checkResponse(response);
 
     return {
       response,
@@ -221,7 +250,7 @@ export class HttpClient implements Requester {
   };
 
   private async checkResponse(response: Response) {
-    if (response.status === 429) {
+    if (response.status === RATELIMIT_STATUS) {
       if (response.headers.get("x-ratelimit-limit-requests")) {
         throw new QstashChatRatelimitError({
           "limit-requests": response.headers.get("x-ratelimit-limit-requests"),

src/client/workflow.test.ts

@@ -5,10 +5,9 @@ import { triggerFirstInvocation } from "./workflow/workflow-requests";
 import { WorkflowContext } from "./workflow/context";
 import { nanoid } from "nanoid";
 import { Client } from "./client";
-import { QstashError } from "./error";
 
 describe("workflow tests", () => {
-  const qstashClient = new Client({ token: process.env.QSTASH_TOKEN! });
+  const qstashClient = new Client({ token: process.env.QSTASH_TOKEN!, retry: false });
   test("should delete workflow succesfully", async () => {
     const workflowRunId = `wfr-${nanoid()}`;
     const result = await triggerFirstInvocation(
@@ -17,7 +16,7 @@ describe("workflow tests", () => {
         workflowRunId,
         headers: new Headers({}) as Headers,
         steps: [],
-        url: "https://some-url.com",
+        url: "https://requestcatcher.com",
         initialPayload: undefined,
       }),
       3
@@ -27,11 +26,5 @@ describe("workflow tests", () => {
     // eslint-disable-next-line @typescript-eslint/no-deprecated
     const cancelResult = await qstashClient.workflow.cancel(workflowRunId);
     expect(cancelResult).toBeTrue();
-
-    // eslint-disable-next-line @typescript-eslint/no-deprecated
-    const throws = qstashClient.workflow.cancel(workflowRunId);
-    expect(throws).rejects.toThrow(
-      new QstashError(`{"error":"workflowRun ${workflowRunId} not found"}`, 404)
-    );
   });
 });

src/client/workflow/test-utils.ts

@@ -16,6 +16,7 @@ export const WORKFLOW_ENDPOINT = "https://www.my-website.com/api";
 export type ResponseFields = {
   body: unknown;
   status: number;
+  headers?: Record<string, string>;
 };
 
 export type RequestFields = {
@@ -84,6 +85,7 @@ export const mockQStashServer = async ({
       }
       return new Response(JSON.stringify(responseFields.body), {
         status: responseFields.status,
+        headers: responseFields.headers,
       });
     },
     port: MOCK_QSTASH_SERVER_PORT,