feat(core): support multi-step doWhile/doUntil loops (#1064)

Author: omeraplak

.changeset/giant-coins-twirl.md

@@ -0,0 +1,10 @@
+---
+"@voltagent/core": patch
+---
+
+Add multi-step loop bodies for `andDoWhile` and `andDoUntil`.
+
+- Loop steps now accept either a single `step` or a sequential `steps` array.
+- When `steps` is provided, each iteration runs the steps in order and feeds each output into the next step.
+- Workflow step serialization now includes loop `subSteps` when a loop has multiple steps.
+- Added runtime and type tests for chained loop steps.

packages/core/src/workflow/core.spec-d.ts

@@ -1,5 +1,5 @@
 import { describe, expectTypeOf, it } from "vitest";
-import { andForEach, andThen } from "./";
+import { andDoWhile, andForEach, andThen } from "./";
 import type { WorkflowExecuteContext } from "./internal/types";
 import type { WorkflowStateStore, WorkflowStateUpdater } from "./types";
 
@@ -55,6 +55,44 @@ describe("non-chaining API type inference", () => {
     });
   });
 
+  describe("andDoWhile", () => {
+    it("should infer condition data from the last chained loop step", () => {
+      const step = andDoWhile({
+        id: "loop",
+        steps: [
+          andThen({
+            id: "step-1",
+            execute: async (
+              context: WorkflowExecuteContext<
+                { input: string },
+                { value: number },
+                unknown,
+                unknown
+              >,
+            ) => ({ value: context.data.value + 1 }),
+          }),
+          andThen({
+            id: "step-2",
+            execute: async (
+              context: WorkflowExecuteContext<
+                { input: string },
+                { value: number },
+                unknown,
+                unknown
+              >,
+            ) => context.data.value,
+          }),
+        ],
+        condition: async (context) => {
+          expectTypeOf(context.data).toEqualTypeOf<number>();
+          return context.data < 3;
+        },
+      });
+
+      expectTypeOf(step).not.toBeNever();
+    });
+  });
+
   describe("type parity with chaining API", () => {
     it("WorkflowExecuteContext should match chaining API context shape", () => {
       // Verify WorkflowExecuteContext has all expected properties

packages/core/src/workflow/core.ts

@@ -2642,9 +2642,17 @@ export function serializeWorkflowStep(step: BaseStep, index: number): Serialized
     case "loop": {
       const loopStep = step as WorkflowStep<unknown, unknown, unknown, unknown> & {
         step?: BaseStep;
+        steps?: BaseStep[];
         condition?: (...args: any[]) => unknown;
         loopType?: "dowhile" | "dountil";
       };
+      const serializedSteps =
+        loopStep.steps && Array.isArray(loopStep.steps)
+          ? loopStep.steps.map((subStep, subIndex) => serializeWorkflowStep(subStep, subIndex))
+          : loopStep.step
+            ? [serializeWorkflowStep(loopStep.step, 0)]
+            : [];
+
       return {
         ...baseStep,
         ...(loopStep.condition && {
@@ -2653,8 +2661,12 @@ export function serializeWorkflowStep(step: BaseStep, index: number): Serialized
         ...(loopStep.loopType && {
           loopType: loopStep.loopType,
         }),
-        ...(loopStep.step && {
-          nestedStep: serializeWorkflowStep(loopStep.step, 0),
+        ...(serializedSteps.length === 1 && {
+          nestedStep: serializedSteps[0],
+        }),
+        ...(serializedSteps.length > 1 && {
+          subSteps: serializedSteps,
+          subStepsCount: serializedSteps.length,
         }),
       };
     }

packages/core/src/workflow/steps/and-loop.spec.ts

@@ -41,4 +41,39 @@ describe("andLoop", () => {
 
     expect(result).toBe(2);
   });
+
+  it("supports multiple chained steps in each loop iteration", async () => {
+    const step = andDoWhile({
+      id: "loop",
+      steps: [
+        andThen({
+          id: "increment",
+          execute: async ({ data }) => data + 1,
+        }),
+        andThen({
+          id: "double",
+          execute: async ({ data }) => data * 2,
+        }),
+      ],
+      condition: async ({ data }) => data < 20,
+    });
+
+    const result = await step.execute(
+      createMockWorkflowExecuteContext({
+        data: 1,
+      }),
+    );
+
+    expect(result).toBe(22);
+  });
+
+  it("throws when loop steps array is empty", () => {
+    expect(() =>
+      andDoUntil({
+        id: "loop",
+        steps: [] as never,
+        condition: async () => true,
+      }),
+    ).toThrow("andDoWhile/andDoUntil requires at least one step");
+  });
 });

packages/core/src/workflow/steps/and-loop.ts

@@ -2,60 +2,100 @@ import type { Span } from "@opentelemetry/api";
 import { defaultStepConfig } from "../internal/utils";
 import { matchStep } from "./helpers";
 import { throwIfAborted } from "./signal";
-import type { WorkflowStepLoop, WorkflowStepLoopConfig } from "./types";
+import type { WorkflowStepLoop, WorkflowStepLoopConfig, WorkflowStepLoopSteps } from "./types";
 
 type LoopType = "dowhile" | "dountil";
+type LoopStepMetadata = { id: string; name?: string; purpose?: string; retries?: number };
+
+function splitLoopConfig<INPUT, DATA, RESULT>(config: WorkflowStepLoopConfig<INPUT, DATA, RESULT>) {
+  if ("step" in config) {
+    const { step: _step, condition, ...stepConfig } = config;
+    return { condition, stepConfig: stepConfig as LoopStepMetadata };
+  }
+
+  const { steps: _steps, condition, ...stepConfig } = config;
+  return { condition, stepConfig: stepConfig as LoopStepMetadata };
+}
+
+function getLoopSteps<INPUT, DATA, RESULT>(
+  config: WorkflowStepLoopConfig<INPUT, DATA, RESULT>,
+): WorkflowStepLoopSteps<INPUT, DATA, RESULT> {
+  if ("steps" in config && config.steps) {
+    if (config.steps.length === 0) {
+      throw new Error("andDoWhile/andDoUntil requires at least one step");
+    }
+    return config.steps;
+  }
+
+  return [config.step];
+}
 
 const createLoopStep = <INPUT, DATA, RESULT>(
   loopType: LoopType,
-  { step, condition, ...config }: WorkflowStepLoopConfig<INPUT, DATA, RESULT>,
+  config: WorkflowStepLoopConfig<INPUT, DATA, RESULT>,
 ) => {
-  const finalStep = matchStep(step);
+  const { condition, stepConfig } = splitLoopConfig(config);
+  const loopSteps = getLoopSteps(config);
+  const resolvedSteps = loopSteps.map((loopStep) => matchStep(loopStep));
+  const firstStep = loopSteps[0];
+
+  if (!firstStep) {
+    throw new Error("andDoWhile/andDoUntil requires at least one step");
+  }
 
   return {
-    ...defaultStepConfig(config),
+    ...defaultStepConfig(stepConfig),
     type: "loop",
     loopType,
-    step,
+    step: firstStep,
+    steps: loopSteps,
     condition,
     execute: async (context) => {
       const { state } = context;
       const traceContext = state.workflowContext?.traceContext;
       let currentData = context.data as DATA | RESULT;
       let iteration = 0;
 
-      while (true) {
-        throwIfAborted(state.signal);
+      const runResolvedStep = async (stepIndex: number) => {
+        const resolvedStep = resolvedSteps[stepIndex];
+        if (!resolvedStep) {
+          return;
+        }
 
         let childSpan: Span | undefined;
         if (traceContext) {
+          const isSingleLoopStep = resolvedSteps.length === 1;
           childSpan = traceContext.createStepSpan(
-            iteration,
-            finalStep.type,
-            finalStep.name || finalStep.id || `Loop ${iteration + 1}`,
+            iteration * resolvedSteps.length + stepIndex,
+            resolvedStep.type,
+            resolvedStep.name ||
+              resolvedStep.id ||
+              (isSingleLoopStep
+                ? `Loop ${iteration + 1}`
+                : `Loop ${iteration + 1}.${stepIndex + 1}`),
             {
-              parentStepId: config.id,
-              parallelIndex: iteration,
+              parentStepId: stepConfig.id,
+              parallelIndex: isSingleLoopStep ? iteration : stepIndex,
               input: currentData,
               attributes: {
                 "workflow.step.loop": true,
                 "workflow.step.parent_type": "loop",
                 "workflow.step.loop_type": loopType,
+                "workflow.step.loop_iteration": iteration,
+                "workflow.step.loop_step_index": stepIndex,
               },
             },
           );
         }
 
-        const subState = {
-          ...state,
-          workflowContext: undefined,
-        };
-
         const executeStep = () =>
-          finalStep.execute({
+          resolvedStep.execute({
             ...context,
-            data: currentData as DATA,
-            state: subState,
+            data: currentData as never,
+            state: {
+              ...state,
+              workflowContext: undefined,
+            },
           });
 
         try {
@@ -73,6 +113,15 @@ const createLoopStep = <INPUT, DATA, RESULT>(
           }
           throw error;
         }
+      };
+
+      while (true) {
+        throwIfAborted(state.signal);
+
+        for (let stepIndex = 0; stepIndex < resolvedSteps.length; stepIndex += 1) {
+          throwIfAborted(state.signal);
+          await runResolvedStep(stepIndex);
+        }
 
         iteration += 1;
         throwIfAborted(state.signal);

packages/core/src/workflow/steps/index.ts

@@ -32,6 +32,7 @@ export type {
   WorkflowStepForEachItemsFunc,
   WorkflowStepForEachMapFunc,
   WorkflowStepLoopConfig,
+  WorkflowStepLoopSteps,
   WorkflowStepBranchConfig,
   WorkflowStepMapConfig,
   WorkflowStepMapEntry,

packages/core/src/workflow/steps/types.ts

@@ -195,16 +195,44 @@ export interface WorkflowStepForEach<INPUT, DATA, ITEM, RESULT, MAP_DATA = ITEM>
   map?: WorkflowStepForEachMapFunc<INPUT, DATA, ITEM, MAP_DATA>;
 }
 
-export type WorkflowStepLoopConfig<INPUT, DATA, RESULT> = InternalWorkflowStepConfig<{
-  step: InternalAnyWorkflowStep<INPUT, DATA, RESULT>;
+export type WorkflowStepLoopSteps<INPUT, DATA, RESULT> =
+  | readonly [InternalAnyWorkflowStep<INPUT, DATA, RESULT>]
+  | readonly [
+      InternalAnyWorkflowStep<INPUT, DATA, DangerouslyAllowAny>,
+      ...InternalAnyWorkflowStep<INPUT, DangerouslyAllowAny, DangerouslyAllowAny>[],
+      InternalAnyWorkflowStep<INPUT, DangerouslyAllowAny, RESULT>,
+    ];
+
+type WorkflowStepLoopBaseConfig<INPUT, RESULT> = InternalWorkflowStepConfig<{
   condition: InternalWorkflowFunc<INPUT, RESULT, boolean, any, any>;
 }>;
 
+type WorkflowStepLoopSingleStepConfig<INPUT, DATA, RESULT> = WorkflowStepLoopBaseConfig<
+  INPUT,
+  RESULT
+> & {
+  step: InternalAnyWorkflowStep<INPUT, DATA, RESULT>;
+  steps?: never;
+};
+
+type WorkflowStepLoopMultiStepConfig<INPUT, DATA, RESULT> = WorkflowStepLoopBaseConfig<
+  INPUT,
+  RESULT
+> & {
+  steps: WorkflowStepLoopSteps<INPUT, DATA, RESULT>;
+  step?: never;
+};
+
+export type WorkflowStepLoopConfig<INPUT, DATA, RESULT> =
+  | WorkflowStepLoopSingleStepConfig<INPUT, DATA, RESULT>
+  | WorkflowStepLoopMultiStepConfig<INPUT, DATA, RESULT>;
+
 export interface WorkflowStepLoop<INPUT, DATA, RESULT>
   extends InternalBaseWorkflowStep<INPUT, DATA, RESULT, any, any> {
   type: "loop";
   loopType: "dowhile" | "dountil";
-  step: InternalAnyWorkflowStep<INPUT, DATA, RESULT>;
+  step: InternalAnyWorkflowStep<INPUT, DATA, DangerouslyAllowAny>;
+  steps: WorkflowStepLoopSteps<INPUT, DATA, RESULT>;
   condition: InternalWorkflowFunc<INPUT, RESULT, boolean, any, any>;
 }
 

website/docs/workflows/steps/and-loop.md

@@ -1,6 +1,6 @@
 # andLoop
 
-> Repeat a step with do-while or do-until semantics.
+> Repeat one or more steps with do-while or do-until semantics.
 
 ## Quick Start
 
@@ -15,10 +15,16 @@ const workflow = createWorkflowChain({
   input: z.number(),
 }).andDoWhile({
   id: "increment-until-3",
-  step: andThen({
-    id: "increment",
-    execute: async ({ data }) => data + 1,
-  }),
+  steps: [
+    andThen({
+      id: "increment",
+      execute: async ({ data }) => data + 1,
+    }),
+    andThen({
+      id: "double",
+      execute: async ({ data }) => data * 2,
+    }),
+  ],
   condition: ({ data }) => data < 3,
 });
 ```
@@ -47,7 +53,8 @@ const workflow = createWorkflowChain({
 ```typescript
 .andDoWhile({
   id: string,
-  step: Step,
+  step?: Step,
+  steps?: Step[],
   condition: (ctx) => boolean | Promise<boolean>,
   retries?: number,
   name?: string,
@@ -56,7 +63,8 @@ const workflow = createWorkflowChain({
 
 .andDoUntil({
   id: string,
-  step: Step,
+  step?: Step,
+  steps?: Step[],
   condition: (ctx) => boolean | Promise<boolean>,
   retries?: number,
   name?: string,
@@ -66,5 +74,7 @@ const workflow = createWorkflowChain({
 
 ## Notes
 
-- The step runs at least once.
+- Provide either `step` (single step) or `steps` (multiple sequential steps).
+- The configured step(s) run at least once.
+- When `steps` is provided, steps run in order on each iteration.
 - The loop continues until the condition fails (do-while) or succeeds (do-until).