Workflow with Conditional Branching (experimental)

import { Mastra } from '@mastra/core';
import { Step, Workflow } from '@mastra/core/workflows';
import { z } from 'zod';
 
 
// Step that provides the initial value
const startStep = new Step({
  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 Step({
  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 Step({
  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 Step({
  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 Workflow({
  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({
  workflows: { conditionalWorkflow },
});
 
// Example usage
async function runWorkflow(inputValue: number) {
  const workflow = mastra.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();