DX-1797: Add WorkflowMiddleware

src/context/auto-executor.ts

@@ -5,6 +5,7 @@ import { type BaseLazyStep } from "./steps";
 import type { WorkflowLogger } from "../logger";
 import { QstashError } from "@upstash/qstash";
 import { submitParallelSteps, submitSingleStep } from "../qstash/submit-steps";
+import { WorkflowMiddleware } from "../middleware";
 
 export class AutoExecutor {
   private context: WorkflowContext;
@@ -15,8 +16,9 @@ export class AutoExecutor {
   private readonly nonPlanStepCount: number;
   private readonly steps: Step[];
   private indexInCurrentList = 0;
-  private invokeCount: number;
-  private telemetry?: Telemetry;
+  private readonly invokeCount: number;
+  private readonly telemetry?: Telemetry;
+  private readonly middlewares?: WorkflowMiddleware[];
 
   public stepCount = 0;
   public planStepCount = 0;
@@ -28,13 +30,15 @@ export class AutoExecutor {
     steps: Step[],
     telemetry?: Telemetry,
     invokeCount?: number,
-    debug?: WorkflowLogger
+    debug?: WorkflowLogger,
+    middlewares?: WorkflowMiddleware[]
   ) {
     this.context = context;
     this.steps = steps;
     this.telemetry = telemetry;
     this.invokeCount = invokeCount ?? 0;
     this.debug = debug;
+    this.middlewares = middlewares;
 
     this.nonPlanStepCount = this.steps.filter((step) => !step.targetStep).length;
   }
@@ -133,7 +137,9 @@ export class AutoExecutor {
         step,
         stepCount: this.stepCount,
       });
-      return lazyStep.parseOut(step.out);
+      const parsedOut = lazyStep.parseOut(step.out);
+
+      return parsedOut;
     }
 
     const resultStep = await submitSingleStep({
@@ -144,6 +150,7 @@ export class AutoExecutor {
       concurrency: 1,
       telemetry: this.telemetry,
       debug: this.debug,
+      middlewares: this.middlewares,
     });
     throw new WorkflowAbort(lazyStep.stepName, resultStep);
   }
@@ -232,6 +239,7 @@ export class AutoExecutor {
             concurrency: parallelSteps.length,
             telemetry: this.telemetry,
             debug: this.debug,
+            middlewares: this.middlewares,
           });
           throw new WorkflowAbort(parallelStep.stepName, resultStep);
         } catch (error) {

src/context/context.ts

@@ -28,6 +28,7 @@ import { WorkflowApi } from "./api";
 import { WorkflowAgents } from "../agents";
 import { FlowControl } from "@upstash/qstash";
 import { getNewUrlFromWorkflowId } from "../serve/serve-many";
+import { WorkflowMiddleware } from "../middleware";
 
 /**
  * Upstash Workflow context
@@ -176,6 +177,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     telemetry,
     invokeCount,
     flowControl,
+    middlewares,
   }: {
     qstashClient: WorkflowClient;
     workflowRunId: string;
@@ -190,6 +192,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     telemetry?: Telemetry;
     invokeCount?: number;
     flowControl?: FlowControl;
+    middlewares?: WorkflowMiddleware[];
   }) {
     this.qstashClient = qstashClient;
     this.workflowRunId = workflowRunId;
@@ -202,7 +205,7 @@ export class WorkflowContext<TInitialPayload = unknown> {
     this.retries = retries ?? DEFAULT_RETRIES;
     this.flowControl = flowControl;
 
-    this.executor = new AutoExecutor(this, this.steps, telemetry, invokeCount, debug);
+    this.executor = new AutoExecutor(this, this.steps, telemetry, invokeCount, debug, middlewares);
   }
 
   /**

src/index.ts

@@ -6,3 +6,4 @@ export * from "./logger";
 export * from "./client";
 export { WorkflowError, WorkflowAbort } from "./error";
 export { WorkflowTool } from "./agents/adapters";
+export { WorkflowMiddleware, loggingMiddleware } from "./middleware";

src/middleware/index.ts

@@ -0,0 +1,2 @@
+export { WorkflowMiddleware } from "./middleware";
+export { loggingMiddleware } from "./logging";

src/middleware/logging.ts

@@ -0,0 +1,23 @@
+import { WorkflowMiddleware } from "./middleware";
+
+export const loggingMiddleware = new WorkflowMiddleware({
+  name: "logging",
+  init: () => {
+    console.log("Logging middleware initialized");
+
+    return {
+      afterExecution(params) {
+        console.log("Step executed:", params);
+      },
+      beforeExecution(params) {
+        console.log("Step execution started:", params);
+      },
+      runStarted(params) {
+        console.log("Workflow run started:", params);
+      },
+      runCompleted(params) {
+        console.log("Workflow run completed:", params);
+      },
+    };
+  },
+});