DX-1614: Add flow control parameters (#66)

* feat: add flow control settings

* fix: add flowControl to getHeaders

* fix: update docstring

* fix: bump qstash dependency to include flow control

* fix: add test for failureFunction in logs

Author: CahidArda

bun.lockb

@@ -102,7 +102,7 @@
   },
   "dependencies": {
     "@ai-sdk/openai": "^1.0.15",
-    "@upstash/qstash": "^2.7.20",
+    "@upstash/qstash": "^2.7.22",
     "ai": "^4.0.30",
     "zod": "^3.24.1"
   },

package.json

@@ -340,7 +340,7 @@ describe("workflow client", () => {
           const postCancelLogs = await liveClient.logs({
             workflowRunId
           })
-    
+
           expect(postCancelLogs.cursor).toBe("")
           expect(postCancelLogs.runs.length).toBe(1)
           expect(postCancelLogs.runs[0]).toEqual({
@@ -378,5 +378,95 @@ describe("workflow client", () => {
     }, {
       timeout: 60000
     })
+
+    test("should include failure logs in case of failure", async () => {
+      const qstashClient = new QStashClient({
+        baseUrl: process.env.QSTASH_URL,
+        token: process.env.QSTASH_TOKEN!,
+      });
+      const liveClient = new Client({
+        baseUrl: process.env.QSTASH_URL,
+        token: process.env.QSTASH_TOKEN!,
+      })
+
+      const body = "some-body"
+      const workflowRunId = "wfr_some-workflow-run-id-" + nanoid()
+
+      const result = await triggerFirstInvocation({
+        workflowContext: new WorkflowContext({
+          qstashClient,
+          headers: new Headers({}) as Headers,
+          initialPayload: body,
+          workflowRunId,
+          steps: [],
+          url: "https://httpstat.us/400",
+          failureUrl: "https://httpstat.us/200",
+          retries: 0
+        })
+      })
+      expect(result.isOk()).toBe(true)
+
+      await eventually(
+        async () => {
+          const logs = await liveClient.logs({
+            workflowRunId
+          })
+
+          expect(logs.cursor).toBe("")
+          expect(logs.runs.length).toBe(1)
+          expect(logs.runs[0]).toEqual({
+            workflowRunId,
+            workflowUrl: "https://httpstat.us/400",
+            workflowState: "RUN_FAILED",
+            workflowRunCreatedAt: expect.any(Number),
+            workflowRunCompletedAt: expect.any(Number),
+            failureFunction: {
+              messageId: expect.any(String),
+              dlqId: expect.any(String),
+              failResponse: "400 Bad Request",
+              failStatus: 400,
+              url: "https://httpstat.us/400",
+              state: "DELIVERED",
+              failHeaders: expect.any(Object),
+            },
+            steps: [
+              {
+                steps: [
+                  {
+                    callType: "step",
+                    concurrent: 1,
+                    createdAt: expect.any(Number),
+                    headers: {
+                      "Upstash-Workflow-Sdk-Version": [
+                        "1"
+                      ],
+                    },
+                    messageId: expect.any(String),
+                    out: "some-body",
+                    state: "STEP_SUCCESS",
+                    stepName: "init",
+                    stepType: "Initial",
+                  },
+                ],
+                type: "sequential",
+              },
+              {
+                steps: [
+                  {
+                    state: "STEP_FAILED",
+                    messageId: expect.any(String),
+                  },
+                ],
+                type: "next",
+              },
+            ],
+          })
+        },
+        { timeout: 30_000, interval: 100 }
+      );
+
+    }, {
+      timeout: 60000
+    })
   })
 });

src/client/index.test.ts

@@ -1,5 +1,5 @@
 import { NotifyResponse, Waiter } from "../types";
-import { Client as QStashClient } from "@upstash/qstash";
+import { FlowControl, Client as QStashClient } from "@upstash/qstash";
 import { makeGetWaitersRequest, makeNotifyRequest } from "./utils";
 import { getWorkflowRunId } from "../utils";
 import { triggerFirstInvocation } from "../workflow-requests";
@@ -198,12 +198,14 @@ export class Client {
     headers,
     workflowRunId,
     retries,
+    flowControl,
   }: {
     url: string;
     body?: unknown;
     headers?: Record<string, string>;
     workflowRunId?: string;
     retries?: number;
+    flowControl?: FlowControl
   }): Promise<{ workflowRunId: string }> {
     const finalWorkflowRunId = getWorkflowRunId(workflowRunId);
     const context = new WorkflowContext({
@@ -216,6 +218,7 @@ export class Client {
       workflowRunId: finalWorkflowRunId,
       retries,
       telemetry: undefined, // can't know workflow telemetry here
+      flowControl,
     });
     const result = await triggerFirstInvocation({
       workflowContext: context,

src/client/index.ts

@@ -58,7 +58,7 @@ export class AutoExecutor {
     if (this.executingStep) {
       throw new WorkflowError(
         "A step can not be run inside another step." +
-          ` Tried to run '${stepInfo.stepName}' inside '${this.executingStep}'`
+        ` Tried to run '${stepInfo.stepName}' inside '${this.executingStep}'`
       );
     }
 
@@ -171,7 +171,7 @@ export class AutoExecutor {
       // user has added/removed a parallel step
       throw new WorkflowError(
         `Incompatible number of parallel steps when call state was '${parallelCallState}'.` +
-          ` Expected ${parallelSteps.length}, got ${plannedParallelStepCount} from the request.`
+        ` Expected ${parallelSteps.length}, got ${plannedParallelStepCount} from the request.`
       );
     }
 
@@ -206,7 +206,7 @@ export class AutoExecutor {
         if (!planStep || planStep.targetStep === undefined) {
           throw new WorkflowError(
             `There must be a last step and it should have targetStep larger than 0.` +
-              `Received: ${JSON.stringify(planStep)}`
+            `Received: ${JSON.stringify(planStep)}`
           );
         }
         const stepIndex = planStep.targetStep - initialStepCount;
@@ -350,6 +350,7 @@ export class AutoExecutor {
         failureUrl: this.context.failureUrl,
         retries: this.context.retries,
         telemetry: this.telemetry,
+        flowControl: this.context.flowControl
       });
 
       // call wait
@@ -393,37 +394,39 @@ export class AutoExecutor {
           callRetries: lazyStep instanceof LazyCallStep ? lazyStep.retries : undefined,
           callTimeout: lazyStep instanceof LazyCallStep ? lazyStep.timeout : undefined,
           telemetry: this.telemetry,
+          flowControl: this.context.flowControl,
+          callFlowControl: lazyStep instanceof LazyCallStep ? lazyStep.flowControl : undefined,
         });
 
         // if the step is a single step execution or a plan step, we can add sleep headers
         const willWait = singleStep.concurrent === NO_CONCURRENCY || singleStep.stepId === 0;
 
         singleStep.out = JSON.stringify(singleStep.out);
 
-        return singleStep.callUrl
+        return singleStep.callUrl && lazyStep instanceof LazyCallStep
           ? // if the step is a third party call, we call the third party
-            // url (singleStep.callUrl) and pass information about the workflow
-            // in the headers (handled in getHeaders). QStash makes the request
-            // to callUrl and returns the result to Workflow endpoint.
-            // handleThirdPartyCallResult method sends the result of the third
-            // party call to QStash.
-            {
-              headers,
-              method: singleStep.callMethod,
-              body: singleStep.callBody,
-              url: singleStep.callUrl,
-            }
+          // url (singleStep.callUrl) and pass information about the workflow
+          // in the headers (handled in getHeaders). QStash makes the request
+          // to callUrl and returns the result to Workflow endpoint.
+          // handleThirdPartyCallResult method sends the result of the third
+          // party call to QStash.
+          {
+            headers,
+            method: singleStep.callMethod,
+            body: singleStep.callBody,
+            url: singleStep.callUrl,
+          }
           : // if the step is not a third party call, we use workflow
-            // endpoint (context.url) as URL when calling QStash. QStash
-            // calls us back with the updated steps list.
-            {
-              headers,
-              method: "POST",
-              body: singleStep,
-              url: this.context.url,
-              notBefore: willWait ? singleStep.sleepUntil : undefined,
-              delay: willWait ? singleStep.sleepFor : undefined,
-            };
+          // endpoint (context.url) as URL when calling QStash. QStash
+          // calls us back with the updated steps list.
+          {
+            headers,
+            method: "POST",
+            body: singleStep,
+            url: this.context.url,
+            notBefore: willWait ? singleStep.sleepUntil : undefined,
+            delay: willWait ? singleStep.sleepFor : undefined,
+          };
       })
     );
 
@@ -496,14 +499,14 @@ const validateStep = (lazyStep: BaseLazyStep, stepFromRequest: Step): void => {
   if (lazyStep.stepName !== stepFromRequest.stepName) {
     throw new WorkflowError(
       `Incompatible step name. Expected '${lazyStep.stepName}',` +
-        ` got '${stepFromRequest.stepName}' from the request`
+      ` got '${stepFromRequest.stepName}' from the request`
     );
   }
   // check type name
   if (lazyStep.stepType !== stepFromRequest.stepType) {
     throw new WorkflowError(
       `Incompatible step type. Expected '${lazyStep.stepType}',` +
-        ` got '${stepFromRequest.stepType}' from the request`
+      ` got '${stepFromRequest.stepType}' from the request`
     );
   }
 };
@@ -531,10 +534,10 @@ const validateParallelSteps = (lazySteps: BaseLazyStep[], stepsFromRequest: Step
       const requestStepTypes = stepsFromRequest.map((step) => step.stepType);
       throw new WorkflowError(
         `Incompatible steps detected in parallel execution: ${error.message}` +
-          `\n  > Step Names from the request: ${JSON.stringify(requestStepNames)}` +
-          `\n    Step Types from the request: ${JSON.stringify(requestStepTypes)}` +
-          `\n  > Step Names expected: ${JSON.stringify(lazyStepNames)}` +
-          `\n    Step Types expected: ${JSON.stringify(lazyStepTypes)}`
+        `\n  > Step Names from the request: ${JSON.stringify(requestStepNames)}` +
+        `\n    Step Types from the request: ${JSON.stringify(requestStepTypes)}` +
+        `\n  > Step Names expected: ${JSON.stringify(lazyStepNames)}` +
+        `\n    Step Types expected: ${JSON.stringify(lazyStepTypes)}`
       );
     }
     throw error;

src/context/auto-executor.ts

@@ -24,6 +24,7 @@ import { WorkflowAbort } from "../error";
 import type { Duration } from "../types";
 import { WorkflowApi } from "./api";
 import { WorkflowAgents } from "../agents";
+import { FlowControl } from "@upstash/qstash";
 
 /**
  * Upstash Workflow context
@@ -152,6 +153,11 @@ export class WorkflowContext<TInitialPayload = unknown> {
    * Number of retries
    */
   public readonly retries: number;
+  /**
+   * Settings for controlling the number of active requests
+   * and number of requests per second with the same key.
+   */
+  public readonly flowControl?: FlowControl
 
   constructor({
     qstashClient,
@@ -165,6 +171,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     env,
     retries,
     telemetry,
+    flowControl
   }: {
     qstashClient: WorkflowClient;
     workflowRunId: string;
@@ -177,6 +184,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     env?: Record<string, string | undefined>;
     retries?: number;
     telemetry?: Telemetry;
+    flowControl?: FlowControl
   }) {
     this.qstashClient = qstashClient;
     this.workflowRunId = workflowRunId;
@@ -187,6 +195,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     this.requestPayload = initialPayload;
     this.env = env ?? {};
     this.retries = retries ?? DEFAULT_RETRIES;
+    this.flowControl = flowControl
 
     this.executor = new AutoExecutor(this, this.steps, telemetry, debug);
   }
@@ -303,7 +312,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     stepName: string,
     settings: CallSettings<TBody>
   ): Promise<CallResponse<TResult>> {
-    const { url, method = "GET", body, headers = {}, retries = 0, timeout } = settings;
+    const { url, method = "GET", body, headers = {}, retries = 0, timeout, flowControl } = settings;
 
     const result = await this.addStep(
       new LazyCallStep<CallResponse<string> | string>(
@@ -313,7 +322,8 @@ export class WorkflowContext<TInitialPayload = unknown> {
         body,
         headers,
         retries,
-        timeout
+        timeout,
+        flowControl
       )
     );
 

src/context/context.ts

@@ -9,7 +9,7 @@ import {
 } from "./steps";
 import { nanoid } from "../utils";
 import type { NotifyResponse, NotifyStepResponse, Step } from "../types";
-import { Client } from "@upstash/qstash";
+import { Client, FlowControl } from "@upstash/qstash";
 import { MOCK_QSTASH_SERVER_URL, mockQStashServer } from "../test-utils";
 import { WorkflowError } from "../error";
 
@@ -146,11 +146,18 @@ describe("test steps", () => {
     const callHeaders = {
       "my-header": headerValue,
     };
-    const step = new LazyCallStep(stepName, callUrl, callMethod, callBody, callHeaders, 14, 30);
+    const flowControl: FlowControl = {
+      key: "my-key",
+      parallelism: 3
+    }
+    const step = new LazyCallStep(stepName, callUrl, callMethod, callBody, callHeaders, 14, 30, flowControl);
 
     test("should set correct fields", () => {
       expect(step.stepName).toBe(stepName);
       expect(step.stepType).toBe("Call");
+      expect(step.flowControl).toEqual(flowControl);
+      expect(step.retries).toBe(14);
+      expect(step.timeout).toBe(30);
     });
     test("should create plan step", () => {
       expect(step.getPlanStep(concurrent, targetStep)).toEqual({
@@ -299,7 +306,7 @@ describe("test steps", () => {
     });
 
     test("should throw when step name is empty string ", () => {
-      const throws = () => new LazyFunctionStep("", () => {});
+      const throws = () => new LazyFunctionStep("", () => { });
       expect(throws).toThrow(
         new WorkflowError(
           "A workflow step name cannot be undefined or an empty string. Please provide a name for your workflow step."