feat: run sub recipe multiple times in parallel (Experimental feature) (#3274)

Co-authored-by: Wendy Tang <wendytang@squareup.com>

crates/goose/src/agents/recipe_tools/mod.rs

@@ -1 +1,2 @@
+pub mod param_utils;
 pub mod sub_recipe_tools;

crates/goose/src/agents/recipe_tools/param_utils/mod.rs

@@ -0,0 +1,38 @@
+use anyhow::Result;
+use serde_json::Value;
+use std::collections::HashMap;
+
+use crate::recipe::SubRecipe;
+
+pub fn prepare_command_params(
+    sub_recipe: &SubRecipe,
+    params_from_tool_call: Vec<Value>,
+) -> Result<Vec<HashMap<String, String>>> {
+    let base_params = sub_recipe.values.clone().unwrap_or_default();
+
+    if params_from_tool_call.is_empty() {
+        return Ok(vec![base_params]);
+    }
+
+    let result = params_from_tool_call
+        .into_iter()
+        .map(|tool_param| {
+            let mut param_map = base_params.clone();
+            if let Some(param_obj) = tool_param.as_object() {
+                for (key, value) in param_obj {
+                    let value_str = value
+                        .as_str()
+                        .map(String::from)
+                        .unwrap_or_else(|| value.to_string());
+                    param_map.entry(key.clone()).or_insert(value_str);
+                }
+            }
+            param_map
+        })
+        .collect();
+
+    Ok(result)
+}
+
+#[cfg(test)]
+mod tests;

crates/goose/src/agents/recipe_tools/param_utils/tests.rs

@@ -0,0 +1,140 @@
+use std::collections::HashMap;
+
+use crate::recipe::SubRecipe;
+use serde_json::json;
+
+use crate::agents::recipe_tools::param_utils::prepare_command_params;
+
+fn setup_default_sub_recipe() -> SubRecipe {
+    let sub_recipe = SubRecipe {
+        name: "test_sub_recipe".to_string(),
+        path: "test_sub_recipe.yaml".to_string(),
+        values: Some(HashMap::from([("key1".to_string(), "value1".to_string())])),
+        sequential_when_repeated: true,
+    };
+    sub_recipe
+}
+
+mod prepare_command_params_tests {
+    use super::*;
+
+    #[test]
+    fn test_return_command_param() {
+        let parameter_array = vec![json!(HashMap::from([(
+            "key2".to_string(),
+            "value2".to_string()
+        )]))];
+        let mut sub_recipe = setup_default_sub_recipe();
+        sub_recipe.values = Some(HashMap::from([("key1".to_string(), "value1".to_string())]));
+
+        let result = prepare_command_params(&sub_recipe, parameter_array).unwrap();
+        assert_eq!(
+            vec![HashMap::from([
+                ("key1".to_string(), "value1".to_string()),
+                ("key2".to_string(), "value2".to_string())
+            ]),],
+            result
+        );
+    }
+
+    #[test]
+    fn test_return_command_param_when_value_override_passed_param_value() {
+        let parameter_array = vec![json!(HashMap::from([(
+            "key2".to_string(),
+            "different_value".to_string()
+        )]))];
+        let mut sub_recipe = setup_default_sub_recipe();
+        sub_recipe.values = Some(HashMap::from([
+            ("key1".to_string(), "value1".to_string()),
+            ("key2".to_string(), "value2".to_string()),
+        ]));
+
+        let result = prepare_command_params(&sub_recipe, parameter_array).unwrap();
+        assert_eq!(
+            vec![HashMap::from([
+                ("key1".to_string(), "value1".to_string()),
+                ("key2".to_string(), "value2".to_string())
+            ]),],
+            result
+        );
+    }
+
+    #[test]
+    fn test_return_empty_command_param() {
+        let parameter_array = vec![];
+        let mut sub_recipe = setup_default_sub_recipe();
+        sub_recipe.values = None;
+
+        let result = prepare_command_params(&sub_recipe, parameter_array).unwrap();
+        assert_eq!(result, vec![HashMap::new()]);
+    }
+
+    mod multiple_tool_parameters {
+        use super::*;
+
+        #[test]
+        fn test_return_command_param_when_all_values_from_tool_call_parameters() {
+            let parameter_array = vec![
+                json!(HashMap::from([
+                    ("key1".to_string(), "key1_value1".to_string()),
+                    ("key2".to_string(), "key2_value1".to_string())
+                ])),
+                json!(HashMap::from([
+                    ("key1".to_string(), "key1_value2".to_string()),
+                    ("key2".to_string(), "key2_value2".to_string())
+                ])),
+            ];
+            let mut sub_recipe = setup_default_sub_recipe();
+            sub_recipe.values = None;
+
+            let result = prepare_command_params(&sub_recipe, parameter_array).unwrap();
+            assert_eq!(
+                vec![
+                    HashMap::from([
+                        ("key1".to_string(), "key1_value1".to_string()),
+                        ("key2".to_string(), "key2_value1".to_string()),
+                    ]),
+                    HashMap::from([
+                        ("key1".to_string(), "key1_value2".to_string()),
+                        ("key2".to_string(), "key2_value2".to_string()),
+                    ]),
+                ],
+                result
+            );
+        }
+
+        #[test]
+        fn test_merge_base_values_with_tool_parameters() {
+            let parameter_array = vec![
+                json!(HashMap::from([(
+                    "key2".to_string(),
+                    "override_value1".to_string()
+                )])),
+                json!(HashMap::from([(
+                    "key2".to_string(),
+                    "override_value2".to_string()
+                )])),
+            ];
+            let mut sub_recipe = setup_default_sub_recipe();
+            sub_recipe.values = Some(HashMap::from([
+                ("key1".to_string(), "base_value".to_string()),
+                ("key2".to_string(), "original_value".to_string()),
+            ]));
+
+            let result = prepare_command_params(&sub_recipe, parameter_array).unwrap();
+            assert_eq!(
+                vec![
+                    HashMap::from([
+                        ("key1".to_string(), "base_value".to_string()),
+                        ("key2".to_string(), "original_value".to_string()),
+                    ]),
+                    HashMap::from([
+                        ("key1".to_string(), "base_value".to_string()),
+                        ("key2".to_string(), "original_value".to_string()),
+                    ]),
+                ],
+                result
+            );
+        }
+    }
+}

crates/goose/src/agents/recipe_tools/sub_recipe_tools.rs

@@ -1,22 +1,35 @@
-use std::{collections::HashMap, fs};
+use std::collections::HashSet;
+use std::fs;
 
 use anyhow::Result;
 use mcp_core::tool::{Tool, ToolAnnotations};
 use serde_json::{json, Map, Value};
 
-use crate::agents::sub_recipe_execution_tool::lib::Task;
+use crate::agents::sub_recipe_execution_tool::lib::{ExecutionMode, Task};
+use crate::agents::sub_recipe_execution_tool::tasks_manager::TasksManager;
 use crate::recipe::{Recipe, RecipeParameter, RecipeParameterRequirement, SubRecipe};
 
+use super::param_utils::prepare_command_params;
+
 pub const SUB_RECIPE_TASK_TOOL_NAME_PREFIX: &str = "subrecipe__create_task";
 
 pub fn create_sub_recipe_task_tool(sub_recipe: &SubRecipe) -> Tool {
     let input_schema = get_input_schema(sub_recipe).unwrap();
     Tool::new(
         format!("{}_{}", SUB_RECIPE_TASK_TOOL_NAME_PREFIX, sub_recipe.name),
-        "Before running this sub recipe, you should first create a task with this tool and then pass the task to the task executor".to_string(),
+        format!(
+            "Create one or more tasks to run the '{}' sub recipe. \
+            Provide an array of parameter sets in the 'task_parameters' field:\n\
+            - For a single task: provide an array with one parameter set\n\
+            - For multiple tasks: provide an array with multiple parameter sets, each with different values\n\n\
+            Each task will run the same sub recipe but with different parameter values. \
+            This is useful when you need to execute the same sub recipe multiple times with varying inputs. \
+            After creating the tasks and execution_mode is provided, pass them to the task executor to run these tasks",
+            sub_recipe.name
+        ),
         input_schema,
         Some(ToolAnnotations {
-            title: Some(format!("create sub recipe task {}", sub_recipe.name)),
+            title: Some(format!("create multiple sub recipe tasks for {}", sub_recipe.name)),
             read_only_hint: false,
             destructive_hint: true,
             idempotent_hint: false,
@@ -25,6 +38,64 @@ pub fn create_sub_recipe_task_tool(sub_recipe: &SubRecipe) -> Tool {
     )
 }
 
+fn extract_task_parameters(params: &Value) -> Vec<Value> {
+    params
+        .get("task_parameters")
+        .and_then(|v| v.as_array())
+        .cloned()
+        .unwrap_or_default()
+}
+
+fn create_tasks_from_params(
+    sub_recipe: &SubRecipe,
+    command_params: &[std::collections::HashMap<String, String>],
+) -> Vec<Task> {
+    let tasks: Vec<Task> = command_params
+        .iter()
+        .map(|task_command_param| {
+            let payload = json!({
+                "sub_recipe": {
+                    "name": sub_recipe.name.clone(),
+                    "command_parameters": task_command_param,
+                    "recipe_path": sub_recipe.path.clone(),
+                    "sequential_when_repeated": sub_recipe.sequential_when_repeated
+                }
+            });
+            Task {
+                id: uuid::Uuid::new_v4().to_string(),
+                task_type: "sub_recipe".to_string(),
+                payload,
+            }
+        })
+        .collect();
+
+    tasks
+}
+
+fn create_task_execution_payload(tasks: &[Task], sub_recipe: &SubRecipe) -> Value {
+    let task_ids: Vec<String> = tasks.iter().map(|task| task.id.clone()).collect();
+    json!({
+        "task_ids": task_ids,
+        "execution_mode": if sub_recipe.sequential_when_repeated { ExecutionMode::Sequential } else { ExecutionMode::Parallel },
+    })
+}
+
+pub async fn create_sub_recipe_task(
+    sub_recipe: &SubRecipe,
+    params: Value,
+    tasks_manager: &TasksManager,
+) -> Result<String> {
+    let task_params_array = extract_task_parameters(&params);
+    let command_params = prepare_command_params(sub_recipe, task_params_array.clone())?;
+    let tasks = create_tasks_from_params(sub_recipe, &command_params);
+    let task_execution_payload = create_task_execution_payload(&tasks, sub_recipe);
+
+    let tasks_json = serde_json::to_string(&task_execution_payload)
+        .map_err(|e| anyhow::anyhow!("Failed to serialize task list: {}", e))?;
+    tasks_manager.save_tasks(tasks.clone()).await;
+    Ok(tasks_json)
+}
+
 fn get_sub_recipe_parameter_definition(
     sub_recipe: &SubRecipe,
 ) -> Result<Option<Vec<RecipeParameter>>> {
@@ -34,22 +105,55 @@ fn get_sub_recipe_parameter_definition(
     Ok(recipe.parameters)
 }
 
-fn get_input_schema(sub_recipe: &SubRecipe) -> Result<Value> {
-    let mut sub_recipe_params_map = HashMap::<String, String>::new();
+fn get_params_with_values(sub_recipe: &SubRecipe) -> HashSet<String> {
+    let mut sub_recipe_params_with_values = HashSet::<String>::new();
     if let Some(params_with_value) = &sub_recipe.values {
-        for (param_name, param_value) in params_with_value {
-            sub_recipe_params_map.insert(param_name.clone(), param_value.clone());
+        for param_name in params_with_value.keys() {
+            sub_recipe_params_with_values.insert(param_name.clone());
         }
     }
+    sub_recipe_params_with_values
+}
+
+fn create_input_schema(param_properties: Map<String, Value>, param_required: Vec<String>) -> Value {
+    let mut properties = Map::new();
+    if !param_properties.is_empty() {
+        properties.insert(
+            "task_parameters".to_string(),
+            json!({
+                "type": "array",
+                "description": "Array of parameter sets for creating tasks. \
+                    For a single task, provide an array with one element. \
+                    For multiple tasks, provide an array with multiple elements, each with different parameter values. \
+                    If there is no parameter set, provide an empty array.",
+                "items": {
+                    "type": "object",
+                    "properties": param_properties,
+                    "required": param_required
+                },
+            })
+        );
+    }
+    json!({
+        "type": "object",
+        "properties": properties,
+    })
+}
+
+fn get_input_schema(sub_recipe: &SubRecipe) -> Result<Value> {
+    let sub_recipe_params_with_values = get_params_with_values(sub_recipe);
+
     let parameter_definition = get_sub_recipe_parameter_definition(sub_recipe)?;
+
+    let mut param_properties = Map::new();
+    let mut param_required = Vec::new();
+
     if let Some(parameters) = parameter_definition {
-        let mut properties = Map::new();
-        let mut required = Vec::new();
         for param in parameters {
-            if sub_recipe_params_map.contains_key(&param.key) {
+            if sub_recipe_params_with_values.contains(&param.key.clone()) {
                 continue;
             }
-            properties.insert(
+            param_properties.insert(
                 param.key.clone(),
                 json!({
                     "type": param.input_type.to_string(),
@@ -57,60 +161,11 @@ fn get_input_schema(sub_recipe: &SubRecipe) -> Result<Value> {
                 }),
             );
             if !matches!(param.requirement, RecipeParameterRequirement::Optional) {
-                required.push(param.key);
+                param_required.push(param.key);
             }
         }
-        Ok(json!({
-            "type": "object",
-            "properties": properties,
-            "required": required
-        }))
-    } else {
-        Ok(json!({
-            "type": "object",
-            "properties": {}
-        }))
     }
-}
-
-fn prepare_command_params(
-    sub_recipe: &SubRecipe,
-    params_from_tool_call: Value,
-) -> Result<HashMap<String, String>> {
-    let mut sub_recipe_params = HashMap::<String, String>::new();
-    if let Some(params_with_value) = &sub_recipe.values {
-        for (param_name, param_value) in params_with_value {
-            sub_recipe_params.insert(param_name.clone(), param_value.clone());
-        }
-    }
-    if let Some(params_map) = params_from_tool_call.as_object() {
-        for (key, value) in params_map {
-            sub_recipe_params.insert(
-                key.to_string(),
-                value.as_str().unwrap_or(&value.to_string()).to_string(),
-            );
-        }
-    }
-    Ok(sub_recipe_params)
-}
-
-pub async fn create_sub_recipe_task(sub_recipe: &SubRecipe, params: Value) -> Result<String> {
-    let command_params = prepare_command_params(sub_recipe, params)?;
-    let payload = json!({
-        "sub_recipe": {
-            "name": sub_recipe.name.clone(),
-            "command_parameters": command_params,
-            "recipe_path": sub_recipe.path.clone(),
-        }
-    });
-    let task = Task {
-        id: uuid::Uuid::new_v4().to_string(),
-        task_type: "sub_recipe".to_string(),
-        payload,
-    };
-    let task_json = serde_json::to_string(&task)
-        .map_err(|e| anyhow::anyhow!("Failed to serialize Task: {}", e))?;
-    Ok(task_json)
+    Ok(create_input_schema(param_properties, param_required))
 }
 
 #[cfg(test)]

crates/goose/src/agents/recipe_tools/sub_recipe_tools/tests.rs

@@ -3,66 +3,48 @@ mod tests {
     use std::collections::HashMap;
 
     use crate::recipe::SubRecipe;
+    use serde_json::json;
+    use serde_json::Value;
+    use tempfile::TempDir;
 
-    fn setup_sub_recipe() -> SubRecipe {
+    fn setup_default_sub_recipe() -> SubRecipe {
         let sub_recipe = SubRecipe {
             name: "test_sub_recipe".to_string(),
             path: "test_sub_recipe.yaml".to_string(),
             values: Some(HashMap::from([("key1".to_string(), "value1".to_string())])),
+            sequential_when_repeated: true,
         };
         sub_recipe
     }
-    mod prepare_command_params_tests {
-        use std::collections::HashMap;
 
-        use crate::{
-            agents::recipe_tools::sub_recipe_tools::{
-                prepare_command_params, tests::tests::setup_sub_recipe,
-            },
-            recipe::SubRecipe,
-        };
+    mod get_input_schema {
+        use super::*;
+        use crate::agents::recipe_tools::sub_recipe_tools::get_input_schema;
 
-        #[test]
-        fn test_prepare_command_params_basic() {
-            let mut params = HashMap::new();
-            params.insert("key2".to_string(), "value2".to_string());
-
-            let sub_recipe = setup_sub_recipe();
-
-            let params_value = serde_json::to_value(params).unwrap();
-            let result = prepare_command_params(&sub_recipe, params_value).unwrap();
-            assert_eq!(result.len(), 2);
-            assert_eq!(result.get("key1"), Some(&"value1".to_string()));
-            assert_eq!(result.get("key2"), Some(&"value2".to_string()));
+        fn prepare_sub_recipe(sub_recipe_file_content: &str) -> (SubRecipe, TempDir) {
+            let mut sub_recipe = setup_default_sub_recipe();
+            let temp_dir = tempfile::tempdir().unwrap();
+            let temp_file = temp_dir.path().join(sub_recipe.path.clone());
+            std::fs::write(&temp_file, sub_recipe_file_content).unwrap();
+            sub_recipe.path = temp_file.to_string_lossy().to_string();
+            (sub_recipe, temp_dir)
         }
 
-        #[test]
-        fn test_prepare_command_params_empty() {
-            let sub_recipe = SubRecipe {
-                name: "test_sub_recipe".to_string(),
-                path: "test_sub_recipe.yaml".to_string(),
-                values: None,
-            };
-            let params: HashMap<String, String> = HashMap::new();
-            let params_value = serde_json::to_value(params).unwrap();
-            let result = prepare_command_params(&sub_recipe, params_value).unwrap();
-            assert_eq!(result.len(), 0);
+        fn verify_task_parameters(result: Value, expected_task_parameters_items: Value) {
+            let task_parameters = result
+                .get("properties")
+                .unwrap()
+                .as_object()
+                .unwrap()
+                .get("task_parameters")
+                .unwrap()
+                .as_object()
+                .unwrap();
+            let task_parameters_items = task_parameters.get("items").unwrap();
+            assert_eq!(&expected_task_parameters_items, task_parameters_items);
         }
-    }
 
-    mod get_input_schema_tests {
-        use crate::{
-            agents::recipe_tools::sub_recipe_tools::{
-                get_input_schema, tests::tests::setup_sub_recipe,
-            },
-            recipe::SubRecipe,
-        };
-
-        #[test]
-        fn test_get_input_schema_with_parameters() {
-            let sub_recipe = setup_sub_recipe();
-
-            let sub_recipe_file_content = r#"{
+        const SUB_RECIPE_FILE_CONTENT_WITH_TWO_PARAMS: &str = r#"{
                 "version": "1.0.0",
                 "title": "Test Recipe",
                 "description": "A test recipe",
@@ -83,73 +65,67 @@ mod tests {
                 ]
             }"#;
 
-            let temp_dir = tempfile::tempdir().unwrap();
-            let temp_file = temp_dir.path().join("test_sub_recipe.yaml");
-            std::fs::write(&temp_file, sub_recipe_file_content).unwrap();
-
-            let mut sub_recipe = sub_recipe;
-            sub_recipe.path = temp_file.to_string_lossy().to_string();
+        #[test]
+        fn test_with_one_param_in_tool_input() {
+            let (mut sub_recipe, _temp_dir) =
+                prepare_sub_recipe(SUB_RECIPE_FILE_CONTENT_WITH_TWO_PARAMS);
+            sub_recipe.values = Some(HashMap::from([("key1".to_string(), "value1".to_string())]));
 
             let result = get_input_schema(&sub_recipe).unwrap();
 
-            // Verify the schema structure
-            assert_eq!(result["type"], "object");
-            assert!(result["properties"].is_object());
-
-            let properties = result["properties"].as_object().unwrap();
-            assert_eq!(properties.len(), 1);
-
-            let key2_prop = &properties["key2"];
-            assert_eq!(key2_prop["type"], "number");
-            assert_eq!(key2_prop["description"], "An optional parameter");
-
-            let required = result["required"].as_array().unwrap();
-            assert_eq!(required.len(), 0);
+            verify_task_parameters(
+                result,
+                json!({
+                    "type": "object",
+                    "properties": {
+                        "key2": { "type": "number", "description": "An optional parameter" }
+                    },
+                    "required": []
+                }),
+            );
         }
 
         #[test]
-        fn test_get_input_schema_no_parameters_values() {
-            let sub_recipe = SubRecipe {
-                name: "test_sub_recipe".to_string(),
-                path: "test_sub_recipe.yaml".to_string(),
-                values: None,
-            };
-
-            let sub_recipe_file_content = r#"{
-                "version": "1.0.0",
-                "title": "Test Recipe",
-                "description": "A test recipe",
-                "prompt": "Test prompt",
-                "parameters": [
-                    {
-                        "key": "key1",
-                        "input_type": "string",
-                        "requirement": "required",
-                        "description": "A test parameter"
-                    }
-                ]
-            }"#;
-
-            let temp_dir = tempfile::tempdir().unwrap();
-            let temp_file = temp_dir.path().join("test_sub_recipe.yaml");
-            std::fs::write(&temp_file, sub_recipe_file_content).unwrap();
-
-            let mut sub_recipe = sub_recipe;
-            sub_recipe.path = temp_file.to_string_lossy().to_string();
+        fn test_without_param_in_tool_input() {
+            let (mut sub_recipe, _temp_dir) =
+                prepare_sub_recipe(SUB_RECIPE_FILE_CONTENT_WITH_TWO_PARAMS);
+            sub_recipe.values = Some(HashMap::from([
+                ("key1".to_string(), "value1".to_string()),
+                ("key2".to_string(), "value2".to_string()),
+            ]));
 
             let result = get_input_schema(&sub_recipe).unwrap();
 
-            assert_eq!(result["type"], "object");
-            assert!(result["properties"].is_object());
+            assert_eq!(
+                None,
+                result
+                    .get("properties")
+                    .unwrap()
+                    .as_object()
+                    .unwrap()
+                    .get("task_parameters")
+            );
+        }
 
-            let properties = result["properties"].as_object().unwrap();
-            assert_eq!(properties.len(), 1);
+        #[test]
+        fn test_with_all_params_in_tool_input() {
+            let (mut sub_recipe, _temp_dir) =
+                prepare_sub_recipe(SUB_RECIPE_FILE_CONTENT_WITH_TWO_PARAMS);
+            sub_recipe.values = None;
 
-            let key1_prop = &properties["key1"];
-            assert_eq!(key1_prop["type"], "string");
-            assert_eq!(key1_prop["description"], "A test parameter");
-            assert_eq!(result["required"].as_array().unwrap().len(), 1);
-            assert_eq!(result["required"][0], "key1");
+            let result = get_input_schema(&sub_recipe).unwrap();
+
+            verify_task_parameters(
+                result,
+                json!({
+                    "type": "object",
+                    "properties": {
+                        "key1": { "type": "string", "description": "A test parameter" },
+                        "key2": { "type": "number", "description": "An optional parameter" }
+                    },
+                    "required": ["key1"]
+                }),
+            );
         }
     }
 }