feat: add WorkflowRetryAfterError error (#144)

* feat: add WorkflowRetryAfterError error

* fix: update zod import path to the latest version

* feat: export WorkflowTool from agents/adapters

Author: CahidArda

src/error.ts

@@ -1,5 +1,5 @@
 import { QstashError } from "@upstash/qstash";
-import type { FailureFunctionPayload, Step } from "./types";
+import type { Duration, FailureFunctionPayload, Step } from "./types";
 
 /**
  * Error raised during Workflow execution
@@ -56,6 +56,20 @@ export class WorkflowNonRetryableError extends WorkflowAbort {
   }
 }
 
+export class WorkflowRetryAfterError extends WorkflowAbort {
+  public retryAfter: number | Duration;
+  /**
+   * @param retryAfter time in seconds after which the workflow should be retried
+   * @param message error message to be displayed
+   */
+  constructor(message: string, retryAfter: number | Duration) {
+    super("retry", undefined, false);
+    this.name = "WorkflowRetryAfterError";
+    this.retryAfter = retryAfter;
+    if (message) this.message = message;
+  }
+}
+
 /**
  * Formats an unknown error to match the FailureFunctionPayload format
  *

src/index.ts

@@ -4,5 +4,10 @@ export * from "./types";
 export * from "./client/types";
 export * from "./logger";
 export * from "./client";
-export { WorkflowError, WorkflowAbort, WorkflowNonRetryableError } from "./error";
+export {
+  WorkflowError,
+  WorkflowAbort,
+  WorkflowNonRetryableError,
+  WorkflowRetryAfterError,
+} from "./error";
 export { WorkflowTool } from "./agents/adapters";

src/serve/authorization.ts

@@ -1,6 +1,11 @@
 import type { Err, Ok } from "neverthrow";
 import { err, ok } from "neverthrow";
-import { WorkflowAbort, WorkflowNonRetryableError } from "../error";
+import {
+  isInstanceOf,
+  WorkflowAbort,
+  WorkflowNonRetryableError,
+  WorkflowRetryAfterError,
+} from "../error";
 import { RouteFunction } from "../types";
 import { WorkflowContext } from "../context";
 import { BaseLazyStep } from "../context/steps";
@@ -94,8 +99,9 @@ export class DisabledWorkflowContext<
       await routeFunction(disabledContext);
     } catch (error) {
       if (
-        (error instanceof WorkflowAbort && error.stepName === this.disabledMessage) ||
-        error instanceof WorkflowNonRetryableError
+        (isInstanceOf(error, WorkflowAbort) && error.stepName === this.disabledMessage) ||
+        isInstanceOf(error, WorkflowNonRetryableError) ||
+        isInstanceOf(error, WorkflowRetryAfterError)
       ) {
         return ok("step-found");
       }

src/serve/index.ts

@@ -7,7 +7,12 @@ import {
   WORKFLOW_PROTOCOL_VERSION_HEADER,
 } from "../constants";
 import { WorkflowContext } from "../context";
-import { formatWorkflowError, isInstanceOf, WorkflowNonRetryableError } from "../error";
+import {
+  formatWorkflowError,
+  isInstanceOf,
+  WorkflowNonRetryableError,
+  WorkflowRetryAfterError,
+} from "../error";
 import { WorkflowLogger } from "../logger";
 import {
   ExclusiveValidationOptions,
@@ -241,6 +246,13 @@ export const serveBase = <
         });
       }
 
+      if (result.isOk() && isInstanceOf(result.value, WorkflowRetryAfterError)) {
+        return onStepFinish(workflowRunId, result.value, {
+          condition: "retry-after-error",
+          result: result.value,
+        });
+      }
+
       if (result.isErr()) {
         // error while running the workflow or when cleaning up
         await debug?.log("ERROR", "ERROR", { error: result.error.message });

src/serve/options.ts

@@ -73,6 +73,14 @@ export const processOptions = <TResponse extends Response = Response, TInitialPa
           },
           status: 489,
         }) as TResponse;
+      } else if (detailedFinishCondition?.condition === "retry-after-error") {
+        return new Response(JSON.stringify(formatWorkflowError(detailedFinishCondition.result)), {
+          headers: {
+            "Retry-After": detailedFinishCondition.result.retryAfter.toString(),
+            [WORKFLOW_PROTOCOL_VERSION_HEADER]: WORKFLOW_PROTOCOL_VERSION,
+          },
+          status: 429,
+        }) as TResponse;
       } else if (detailedFinishCondition?.condition === "failure-callback") {
         return new Response(
           JSON.stringify({ result: detailedFinishCondition.result ?? undefined }),

src/serve/serve.test.ts

@@ -27,7 +27,7 @@ import {
 import { AUTH_FAIL_MESSAGE, processOptions } from "./options";
 import { WorkflowLogger } from "../logger";
 import { z } from "zod";
-import { WorkflowNonRetryableError } from "../error";
+import { WorkflowNonRetryableError, WorkflowRetryAfterError } from "../error";
 
 const someWork = (input: string) => {
   return `processed '${input}'`;
@@ -913,6 +913,41 @@ describe("serve", () => {
     expect(called).toBeTrue();
   });
 
+  test("should call qstash to retry workflow on WorkflowRetryAfterError", async () => {
+    const request = getRequest(WORKFLOW_ENDPOINT, "wfr-foo-3", "my-payload", []);
+    let called = false;
+    const { handler: endpoint } = serve(
+      async (context) => {
+        called = true;
+        throw new WorkflowRetryAfterError("This is a retry-after error", 30);
+      },
+      {
+        qstashClient,
+        receiver: undefined,
+        verbose: true,
+      }
+    );
+
+    await mockQStashServer({
+      execute: async () => {
+        const response = await endpoint(request);
+
+        expect(response.status).toBe(429);
+        expect(response.headers.get("Retry-After")).toBe("30");
+
+        const body = await response.json();
+        expect(body).toEqual({
+          error: "WorkflowRetryAfterError",
+          message: "This is a retry-after error",
+          stack: expect.any(String),
+        });
+      },
+      responseFields: { body: undefined, status: 429 },
+      receivesRequest: false,
+    });
+    expect(called).toBeTrue();
+  });
+
   test("should send waitForEvent", async () => {
     const request = getRequest(WORKFLOW_ENDPOINT, "wfr-bar", "my-payload", []);
     const { handler: endpoint } = serve(

src/types.ts

@@ -4,7 +4,7 @@ import type { HTTPMethods } from "@upstash/qstash";
 import type { WorkflowContext } from "./context";
 import type { WorkflowLogger } from "./logger";
 import { z } from "zod";
-import { WorkflowNonRetryableError } from "./error";
+import { WorkflowNonRetryableError, WorkflowRetryAfterError } from "./error";
 
 /**
  * Interface for Client with required methods
@@ -145,6 +145,10 @@ export type DetailedFinishCondition =
       condition: "non-retryable-error";
       result: WorkflowNonRetryableError;
     }
+  | {
+      condition: "retry-after-error";
+      result: WorkflowRetryAfterError;
+    }
   | {
       condition: "failure-callback";
       result: string | void;

src/workflow-requests.test.ts

@@ -9,7 +9,7 @@ import {
   triggerRouteFunction,
   triggerWorkflowDelete,
 } from "./workflow-requests";
-import { WorkflowAbort, WorkflowNonRetryableError } from "./error";
+import { WorkflowAbort, WorkflowNonRetryableError, WorkflowRetryAfterError } from "./error";
 import { WorkflowContext } from "./context";
 import { Client } from "@upstash/qstash";
 import { Client as WorkflowClient } from "./client";
@@ -263,6 +263,25 @@ describe("Workflow Requests", () => {
     expect(result.value).toBeInstanceOf(WorkflowNonRetryableError);
   });
 
+  test("should retry workflow and return ok if WorkflowRetryAfterError is thrown", async () => {
+    const result = await triggerRouteFunction({
+      onStep: async () => {
+        throw new WorkflowRetryAfterError("This is a retry-after error", 5);
+      },
+      onCleanup: async () => {
+        throw new Error("shouldn't call");
+      },
+      onCancel: async () => {
+        throw new Error("shouldn't call");
+      },
+    });
+    expect(result.isOk()).toBeTrue();
+    // @ts-expect-error value will be set since result isOk
+    expect(result.value).toBeInstanceOf(WorkflowRetryAfterError);
+    // @ts-expect-error value will be set since result isOk
+    expect(result.value.retryAfter).toBe(5);
+  });
+
   test("should call onCancel if context.cancel is called inside context.run", async () => {
     const workflowRunId = nanoid();
     const token = "myToken";

src/workflow-requests.ts

@@ -1,6 +1,12 @@
 import type { Err, Ok } from "neverthrow";
 import { err, ok } from "neverthrow";
-import { isInstanceOf, WorkflowAbort, WorkflowError, WorkflowNonRetryableError } from "./error";
+import {
+  isInstanceOf,
+  WorkflowAbort,
+  WorkflowError,
+  WorkflowNonRetryableError,
+  WorkflowRetryAfterError,
+} from "./error";
 import type { WorkflowContext } from "./context";
 import {
   TELEMETRY_HEADER_FRAMEWORK,
@@ -164,7 +170,11 @@ export const triggerRouteFunction = async ({
   debug?: WorkflowLogger;
 }): Promise<
   | Ok<
-      "workflow-finished" | "step-finished" | "workflow-was-finished" | WorkflowNonRetryableError,
+      | "workflow-finished"
+      | "step-finished"
+      | "workflow-was-finished"
+      | WorkflowNonRetryableError
+      | WorkflowRetryAfterError,
       never
     >
   | Err<never, Error>
@@ -185,7 +195,10 @@ export const triggerRouteFunction = async ({
         errorMessage: error.message,
       });
       return ok("workflow-was-finished");
-    } else if (isInstanceOf(error_, WorkflowNonRetryableError)) {
+    } else if (
+      isInstanceOf(error_, WorkflowNonRetryableError) ||
+      isInstanceOf(error_, WorkflowRetryAfterError)
+    ) {
       return ok(error_);
     } else if (!isInstanceOf(error_, WorkflowAbort)) {
       return err(error_);