feat(orchestra): add custom commands with typed arguments

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add support for declaring custom commands in YAML flows. A flow file can declare itself as a reusable command via `command:` + `arguments:` in its config block; callers invoke it as if it were a built-in command (e.g. `- takeScreenshotWithContext: { label: "login" }`) instead of `runFlow:` with `env:`. Files placed under a sibling `subflows/` folder are auto-discovered in single-file runs; in workspace runs, any file with a `command:` header is registered regardless of folder.
+
 ## 2.6.0
 
 - Add Maestro Viewer - a desktop visualizer for live hierarchy and commands

maestro-cli/src/main/java/maestro/cli/command/TestCommand.kt

@@ -61,6 +61,7 @@ import maestro.cli.auth.Auth
 import maestro.cli.model.FlowStatus
 import maestro.cli.view.cyan
 import maestro.cli.promotion.PromotionStateManager
+import maestro.orchestra.CustomCommandDef
 import maestro.orchestra.error.ValidationError
 import maestro.orchestra.workspace.WorkspaceExecutionPlanner
 import maestro.orchestra.workspace.WorkspaceExecutionPlanner.ExecutionPlan
@@ -527,9 +528,10 @@ class TestCommand : Callable<Int> {
                         authToken,
                         testOutputDir,
                         deviceId,
+                        executionPlan.customCommands,
                     )
                 } else {
-                    runSingleFlow(maestro, device, flowFile, debugOutputPath, testOutputDir, deviceId)
+                    runSingleFlow(maestro, device, flowFile, debugOutputPath, testOutputDir, deviceId, executionPlan.customCommands)
                 }
             }
         }
@@ -555,6 +557,7 @@ class TestCommand : Callable<Int> {
         debugOutputPath: Path,
         testOutputDir: Path?,
         deviceId: String?,
+        customCommands: Map<String, CustomCommandDef> = emptyMap(),
     ): Triple<Int, Int, Nothing?> {
         val resultView =
             if (DisableAnsiMixin.ansiEnabled) {
@@ -579,6 +582,7 @@ class TestCommand : Callable<Int> {
             apiKey = authToken,
             testOutputDir = testOutputDir,
             deviceId = deviceId,
+            customCommands = customCommands,
         )
         val duration = System.currentTimeMillis() - startTime
 
@@ -665,7 +669,7 @@ class TestCommand : Callable<Int> {
             .groupBy { it.index % effectiveShards }
             .map { (_, files) ->
                 val flowsToRun = files.map { it.value }
-                ExecutionPlan(flowsToRun, plan.sequence, plan.workspaceConfig)
+                ExecutionPlan(flowsToRun, plan.sequence, plan.workspaceConfig, plan.customCommands)
             }
     }
 

maestro-cli/src/main/java/maestro/cli/runner/TestRunner.kt

@@ -17,6 +17,7 @@ import maestro.cli.runner.resultview.ResultView
 import maestro.cli.runner.resultview.UiState
 import maestro.cli.util.PrintUtils
 import maestro.cli.view.ErrorViewUtils
+import maestro.orchestra.CustomCommandDef
 import maestro.orchestra.MaestroCommand
 import maestro.orchestra.debug.FlowDebugOutput
 import maestro.orchestra.util.Env.withEnv
@@ -51,6 +52,7 @@ object TestRunner {
         apiKey: String? = null,
         testOutputDir: Path?,
         deviceId: String?,
+        customCommands: Map<String, CustomCommandDef> = emptyMap(),
     ): Int {
         val debugOutput = FlowDebugOutput()
         var aiOutput = FlowAIOutput(
@@ -63,7 +65,7 @@ object TestRunner {
             .withDefaultEnvVars(flowFile, deviceId)
 
         val result = runCatching(resultView, maestro) {
-            val commands = YamlCommandReader.readCommands(flowFile.toPath())
+            val commands = YamlCommandReader.readCommands(flowFile.toPath(), customCommands)
                 .withEnv(updatedEnv)
 
             val flowName = YamlCommandReader.getConfig(commands)?.name ?: flowFile.nameWithoutExtension
@@ -127,6 +129,7 @@ object TestRunner {
         apiKey: String? = null,
         testOutputDir: Path?,
         deviceId: String?,
+        customCommands: Map<String, CustomCommandDef> = emptyMap(),
     ): Nothing {
         val resultView = AnsiResultView("> Press [ENTER] to restart the Flow\n\n")
 
@@ -147,7 +150,7 @@ object TestRunner {
                     .withDefaultEnvVars(flowFile, deviceId)
 
                 val commands = YamlCommandReader
-                    .readCommands(flowFile.toPath())
+                    .readCommands(flowFile.toPath(), customCommands)
                     .withEnv(updatedEnv)
 
                 val flowName = YamlCommandReader.getConfig(commands)?.name

maestro-cli/src/main/java/maestro/cli/runner/TestSuiteInteractor.kt

@@ -16,6 +16,7 @@ import maestro.cli.util.TimeUtils
 import maestro.cli.view.ErrorViewUtils
 import maestro.cli.view.TestSuiteStatusView
 import maestro.cli.view.TestSuiteStatusView.TestSuiteViewModel
+import maestro.orchestra.CustomCommandDef
 import maestro.orchestra.Orchestra
 import maestro.orchestra.debug.CommandDebugMetadata
 import maestro.orchestra.debug.CommandStatus
@@ -77,7 +78,7 @@ class TestSuiteInteractor(
             val updatedEnv = env
                 .withInjectedShellEnvVars()
                 .withDefaultEnvVars(flowFile, deviceId, shardIndex)
-            val (result, aiOutput) = runFlow(flowFile, updatedEnv, maestro, debugOutputPath, testOutputDir)
+            val (result, aiOutput) = runFlow(flowFile, updatedEnv, maestro, debugOutputPath, testOutputDir, executionPlan.customCommands)
             flowResults.add(result)
             aiOutputs.add(aiOutput)
 
@@ -97,7 +98,7 @@ class TestSuiteInteractor(
             val updatedEnv = env
                 .withInjectedShellEnvVars()
                 .withDefaultEnvVars(flowFile, deviceId, shardIndex)
-            val (result, aiOutput) = runFlow(flowFile, updatedEnv, maestro, debugOutputPath, testOutputDir)
+            val (result, aiOutput) = runFlow(flowFile, updatedEnv, maestro, debugOutputPath, testOutputDir, executionPlan.customCommands)
             aiOutputs.add(aiOutput)
 
             if (result.status == FlowStatus.ERROR) {
@@ -158,7 +159,8 @@ class TestSuiteInteractor(
         env: Map<String, String>,
         maestro: Maestro,
         debugOutputPath: Path,
-        testOutputDir: Path? = null
+        testOutputDir: Path? = null,
+        customCommands: Map<String, CustomCommandDef> = emptyMap(),
     ): Pair<TestExecutionSummary.FlowResult, FlowAIOutput> {
         // TODO(bartekpacia): merge TestExecutionSummary with AI suggestions
         //  (i.e. consider them also part of the test output)
@@ -173,7 +175,7 @@ class TestSuiteInteractor(
             flowFile = flowFile,
         )
         val commands = YamlCommandReader
-            .readCommands(flowFile.toPath())
+            .readCommands(flowFile.toPath(), customCommands)
             .withEnv(env)
 
         val maestroConfig = YamlCommandReader.getConfig(commands)

maestro-orchestra-models/src/main/java/maestro/orchestra/Commands.kt

@@ -812,6 +812,8 @@ data class RunFlowCommand(
     val config: MaestroConfig?,
     override val label: String? = null,
     override val optional: Boolean = false,
+    /** Non-null only when this RunFlowCommand was produced by expanding a custom-command call. */
+    val customCommandName: String? = null,
 ) : CompositeCommand {
 
     override fun subCommands(): List<MaestroCommand> {

maestro-orchestra-models/src/main/java/maestro/orchestra/CustomCommandDef.kt

@@ -0,0 +1,19 @@
+package maestro.orchestra
+
+import com.fasterxml.jackson.annotation.JsonIgnore
+import java.nio.file.Path
+
+enum class ArgumentType { STRING, NUMBER, BOOLEAN }
+
+data class CustomCommandArgument(
+    val name: String,
+    val type: ArgumentType = ArgumentType.STRING,
+    val required: Boolean = false,
+    val default: String? = null,
+)
+
+data class CustomCommandDef(
+    val name: String,
+    @field:JsonIgnore val sourceFile: Path,
+    val arguments: List<CustomCommandArgument>,
+)

maestro-orchestra-models/src/main/java/maestro/orchestra/MaestroConfig.kt

@@ -13,6 +13,8 @@ data class MaestroConfig(
     val onFlowStart: MaestroOnFlowStart? = null,
     val onFlowComplete: MaestroOnFlowComplete? = null,
     val properties: Map<String, String> = emptyMap(),
+    val command: String? = null,
+    val arguments: List<CustomCommandArgument>? = null,
 ) {
 
     fun evaluateScripts(jsEngine: JsEngine): MaestroConfig {

maestro-orchestra/src/main/java/maestro/orchestra/workspace/WorkspaceExecutionPlanner.kt

@@ -1,9 +1,12 @@
 package maestro.orchestra.workspace
 
+import maestro.orchestra.CustomCommandDef
 import maestro.orchestra.MaestroCommand
 import maestro.orchestra.WorkspaceConfig
+import maestro.orchestra.error.SyntaxError
 import maestro.orchestra.error.ValidationError
 import maestro.orchestra.workspace.ExecutionOrderPlanner.getFlowsToRunInSequence
+import maestro.orchestra.yaml.MaestroFlowParser
 import maestro.orchestra.yaml.YamlCommandReader
 import org.slf4j.LoggerFactory
 import java.nio.file.Files
@@ -29,7 +32,9 @@ object WorkspaceExecutionPlanner {
         }
 
         if (input.isRegularFile) {
-            validateFlowFile(input.first())
+            val flow = input.first()
+            val singleFileCustomCommands = discoverCustomCommandsFromSubflowsDir(flow)
+            validateFlowFile(flow, singleFileCustomCommands)
             val workspaceConfig = if (config != null) {
                 YamlCommandReader.readWorkspaceConfig(config.absolute())
             } else {
@@ -38,7 +43,8 @@ object WorkspaceExecutionPlanner {
             return ExecutionPlan(
                 flowsToRun = input.toList(),
                 sequence = FlowSequence(emptyList()),
-                workspaceConfig = workspaceConfig
+                workspaceConfig = workspaceConfig,
+                customCommands = singleFileCustomCommands,
             )
         }
 
@@ -62,6 +68,11 @@ object WorkspaceExecutionPlanner {
             """.trimIndent())
         }
 
+        val customCommands = discoverCustomCommands(flowFiles + flowFilesInDirs)
+
+        // Command-definition files (any file declaring `command:`) are not runnable flows.
+        val customCommandFiles = customCommands.values.map { it.sourceFile }.toSet()
+
         // Filter flows based on flows config
 
         val workspaceConfig =
@@ -74,9 +85,9 @@ object WorkspaceExecutionPlanner {
             directories.map { it.fileSystem.getPathMatcher(escapeSlashesForWindows("glob:${it.pathString}${it.fileSystem.separator}$glob")) }
         }
 
-        val unsortedFlowFiles = flowFiles + flowFilesInDirs.filter { path ->
+        val unsortedFlowFiles = (flowFiles + flowFilesInDirs.filter { path ->
             matchers.any { matcher -> matcher.matches(path) }
-        }.toList()
+        }).filter { it.absolute() !in customCommandFiles }
 
         if (unsortedFlowFiles.isEmpty()) {
             if ("*" == globs.singleOrNull()) {
@@ -99,7 +110,7 @@ object WorkspaceExecutionPlanner {
         // Filter flows based on tags
 
         val configPerFlowFile = unsortedFlowFiles.associateWith {
-            val commands = validateFlowFile(it)
+            val commands = validateFlowFile(it, customCommands)
             YamlCommandReader.getConfig(commands)
         }
 
@@ -140,7 +151,7 @@ object WorkspaceExecutionPlanner {
         // validation of media files for add media command
         allFlows.forEach {
             val commands = YamlCommandReader
-                .readCommands(it)
+                .readCommands(it, customCommands)
                 .mapNotNull { maestroCommand -> maestroCommand.addMediaCommand }
             val mediaPaths = commands.flatMap { addMediaCommand -> addMediaCommand.mediaPaths }
             YamlCommandsPathValidator.validatePathsExistInWorkspace(input, it, mediaPaths)
@@ -153,15 +164,96 @@ object WorkspaceExecutionPlanner {
                 workspaceConfig.executionOrder?.continueOnFailure
             ),
             workspaceConfig = workspaceConfig,
+            customCommands = customCommands,
         )
 
         logger.info("Created execution plan: $executionPlan")
 
         return executionPlan
     }
 
-    private fun validateFlowFile(topLevelFlowPath: Path): List<MaestroCommand> {
-        return YamlCommandReader.readCommands(topLevelFlowPath)
+    private fun validateFlowFile(
+        topLevelFlowPath: Path,
+        customCommands: Map<String, CustomCommandDef> = emptyMap(),
+    ): List<MaestroCommand> {
+        return YamlCommandReader.readCommands(topLevelFlowPath, customCommands)
+    }
+
+    internal fun discoverCustomCommandsFromSubflowsDir(flowFile: Path): Map<String, CustomCommandDef> {
+        val parent = flowFile.absolute().parent ?: return emptyMap()
+        val cmdDir = parent.resolve("subflows")
+        if (!cmdDir.isDirectory()) return emptyMap()
+        val files = Files.walk(cmdDir).use { stream ->
+            stream.filter { it.isRegularFile() && isYamlFile(it) }.toList()
+        }
+        return discoverCustomCommands(files)
+    }
+
+    private fun isYamlFile(path: Path): Boolean {
+        val name = path.fileName.toString().lowercase()
+        return name.endsWith(".yaml") || name.endsWith(".yml")
+    }
+
+    internal fun discoverCustomCommands(files: Iterable<Path>): Map<String, CustomCommandDef> {
+        val registry = mutableMapOf<String, CustomCommandDef>()
+        val seenFiles = mutableSetOf<Path>()
+        for (file in files) {
+            val canonical = file.absolute().normalize()
+            if (!seenFiles.add(canonical)) continue
+            val config = try {
+                YamlCommandReader.readConfig(file)
+            } catch (e: Exception) {
+                // Skip files that fail to parse during the pre-pass. If the file is a
+                // runnable flow, the real parse error will surface during validation;
+                // if it was meant as a command-definition file, callers see "Invalid
+                // Command" — log a hint so the cause is recoverable from DEBUG output.
+                logger.debug("Skipping {} during custom-command discovery: {}", file, e.message)
+                continue
+            }
+            val name = config.command ?: continue
+            val arguments = config.toCustomCommandArguments().orEmpty()
+
+            if (MaestroFlowParser.isBuiltInCommand(name)) {
+                throw SyntaxError(
+                    "Custom command name '$name' (declared in ${file.absolutePathString()}) " +
+                        "collides with a built-in Maestro command."
+                )
+            }
+            val existing = registry[name]
+            if (existing != null) {
+                throw SyntaxError(
+                    "Duplicate custom command name '$name': declared in both " +
+                        "${existing.sourceFile.absolutePathString()} and ${file.absolutePathString()}."
+                )
+            }
+            registry[name] = CustomCommandDef(
+                name = name,
+                sourceFile = file.absolute(),
+                arguments = arguments,
+            )
+        }
+
+        // Reject nesting: a custom-command body cannot invoke another custom command.
+        for (def in registry.values) {
+            val body = try {
+                YamlCommandReader.readCommands(def.sourceFile, registry)
+            } catch (e: Exception) {
+                logger.debug("Skipping nesting check for {}: {}", def.sourceFile, e.message)
+                continue
+            }
+            val nested = body.firstNotNullOfOrNull { mc ->
+                mc.runFlowCommand?.customCommandName
+                    ?.takeIf { it != def.name }
+                    ?.let(registry::get)
+            }
+            if (nested != null) {
+                throw SyntaxError(
+                    "Custom command '${nested.name}' invoked inside custom command '${def.name}' " +
+                        "— nesting is not supported."
+                )
+            }
+        }
+        return registry
     }
 
     private fun findConfigFile(input: Path): Path? {
@@ -192,5 +284,6 @@ object WorkspaceExecutionPlanner {
         val flowsToRun: List<Path>,
         val sequence: FlowSequence,
         val workspaceConfig: WorkspaceConfig = WorkspaceConfig(),
+        val customCommands: Map<String, CustomCommandDef> = emptyMap(),
     )
 }