Workflow (Legacy) with Conditional Branching (experimental)

import { Mastra } from "@mastra/core";
import { LegacyStep, LegacyWorkflow } from "@mastra/core/workflows/legacy";
import { z } from "zod";
 
// Step that provides the initial value
const startStep = new LegacyStep({
  id: "start",
  outputSchema: z.object({
    value: z.number(),
  }),
  execute: async ({ context }) => {
    // Get the value from the trigger data
    const value = context.triggerData.inputValue;
    return { value };
  },
});
 
// Step that handles high values
const highValueStep = new LegacyStep({
  id: "highValue",
  outputSchema: z.object({
    result: z.string(),
  }),
  execute: async ({ context }) => {
    const value = context.getStepResult<{ value: number }>("start")?.value;
    return { result: `High value processed: ${value}` };
  },
});
 
// Step that handles low values
const lowValueStep = new LegacyStep({
  id: "lowValue",
  outputSchema: z.object({
    result: z.string(),
  }),
  execute: async ({ context }) => {
    const value = context.getStepResult<{ value: number }>("start")?.value;
    return { result: `Low value processed: ${value}` };
  },
});
 
// Final step that summarizes the result
const finalStep = new LegacyStep({
  id: "final",
  outputSchema: z.object({
    summary: z.string(),
  }),
  execute: async ({ context }) => {
    // Get the result from whichever branch executed
    const highResult = context.getStepResult<{ result: string }>(
      "highValue",
    )?.result;
    const lowResult = context.getStepResult<{ result: string }>(
      "lowValue",
    )?.result;
 
    const result = highResult || lowResult;
    return { summary: `Processing complete: ${result}` };
  },
});
 
// Build the workflow with conditional branching
const conditionalWorkflow = new LegacyWorkflow({
  name: "conditional-workflow",
  triggerSchema: z.object({
    inputValue: z.number(),
  }),
});
 
conditionalWorkflow
  .step(startStep)
  .if(async ({ context }) => {
    const value = context.getStepResult<{ value: number }>("start")?.value ?? 0;
    return value >= 10; // Condition: value is 10 or greater
  })
  .then(highValueStep)
  .then(finalStep)
  .else()
  .then(lowValueStep)
  .then(finalStep) // Both branches converge on the final step
  .commit();
 
// Register the workflow
const mastra = new Mastra({
  legacy_workflows: { conditionalWorkflow },
});
 
// Example usage
async function runWorkflow(inputValue: number) {
  const workflow = mastra.legacy_getWorkflow("conditionalWorkflow");
  const { start } = workflow.createRun();
 
  const result = await start({
    triggerData: { inputValue },
  });
 
  console.log("Workflow result:", result.results);
  return result;
}
 
// Run with a high value (follows the "if" branch)
const result1 = await runWorkflow(15);
// Run with a low value (follows the "else" branch)
const result2 = await runWorkflow(5);
 
console.log("Result 1:", result1);
console.log("Result 2:", result2);

// Using reference-based conditions instead of functions
conditionalWorkflow
  .step(startStep)
  .if({
    ref: { step: startStep, path: "value" },
    query: { $gte: 10 }, // Condition: value is 10 or greater
  })
  .then(highValueStep)
  .then(finalStep)
  .else()
  .then(lowValueStep)
  .then(finalStep)
  .commit();