diff --git a/crates/goose-cli/src/session/elicitation.rs b/crates/goose-cli/src/session/elicitation.rs index 012295ff968e..8ddaa496f12f 100644 --- a/crates/goose-cli/src/session/elicitation.rs +++ b/crates/goose-cli/src/session/elicitation.rs @@ -9,6 +9,19 @@ pub struct ElicitationInput { pub user_data: HashMap, } +#[derive(Clone, Debug, Eq, PartialEq)] +enum SelectChoice { + Value(String), + Skip, +} + +struct SingleSelect<'a> { + field_name: &'a str, + description: Option<&'a str>, + options: Vec<(SelectChoice, String)>, + initial_value: Option, +} + pub fn collect_elicitation_input(message: &str, schema: &Value) -> io::Result { if !message.is_empty() { println!("\n{}", style(message).cyan()); @@ -16,8 +29,12 @@ pub fn collect_elicitation_input(message: &str, schema: &Value) -> io::Result props, _ => { @@ -62,7 +79,6 @@ pub fn collect_elicitation_input(message: &str, schema: &Value) -> io::Result format!("{} ({})", name, desc), @@ -108,7 +124,6 @@ pub fn collect_elicitation_input(message: &str, schema: &Value) -> io::Result io::Result Option> { + let properties = schema.get("properties")?.as_object()?; + if properties.len() != 1 { + return None; + } + + let (field_name, field_schema) = properties.iter().next()?; + let description = field_schema.get("description").and_then(Value::as_str); + let mut options: Vec<(SelectChoice, String)> = + if let Some(one_of) = field_schema.get("oneOf").and_then(Value::as_array) { + one_of + .iter() + .map(|option| { + let value = option.get("const")?.as_str()?; + let label = option.get("title").and_then(Value::as_str).unwrap_or(value); + Some((SelectChoice::Value(value.to_string()), label.to_string())) + }) + .collect::>()? + } else { + field_schema + .get("enum")? + .as_array()? + .iter() + .map(|value| { + let value = value.as_str()?; + Some((SelectChoice::Value(value.to_string()), value.to_string())) + }) + .collect::>()? + }; + + if options.is_empty() { + return None; + } + + let is_required = schema + .get("required") + .and_then(Value::as_array) + .is_some_and(|required| { + required + .iter() + .any(|value| value.as_str() == Some(field_name)) + }); + let default_value = field_schema + .get("default") + .and_then(Value::as_str) + .map(|value| SelectChoice::Value(value.to_string())) + .filter(|value| options.iter().any(|(option, _)| option == value)); + + let initial_value = if !is_required && default_value.is_none() { + options.push((SelectChoice::Skip, "Skip".to_string())); + Some(SelectChoice::Skip) + } else { + default_value + }; + + Some(SingleSelect { + field_name, + description, + options, + initial_value, + }) +} + +fn prompt_single_select(select: SingleSelect<'_>) -> io::Result { + let items: Vec<_> = select + .options + .iter() + .map(|(value, label)| (value.clone(), label, "")) + .collect(); + let label = match select.description { + Some(desc) => format!("{} ({})", select.field_name, desc), + None => select.field_name.to_string(), + }; + let mut prompt = cliclack::select(label).items(&items); + if let Some(initial_value) = select.initial_value { + prompt = prompt.initial_value(initial_value); + } + + match prompt.interact() { + Ok(SelectChoice::Value(value)) => Ok(ElicitationInput { + action: ElicitationAction::Accept, + user_data: HashMap::from([(select.field_name.to_string(), Value::String(value))]), + }), + Ok(SelectChoice::Skip) => Ok(ElicitationInput { + action: ElicitationAction::Accept, + user_data: HashMap::new(), + }), + Err(error) if error.kind() == io::ErrorKind::Interrupted => Ok(ElicitationInput { + action: ElicitationAction::Cancel, + user_data: HashMap::new(), + }), + Err(error) => Err(error), + } +} + fn read_line() -> io::Result> { if !std::io::stdin().is_terminal() { let mut line = String::new(); @@ -157,7 +267,7 @@ fn read_line() -> io::Result> { let mut line = String::new(); match io::stdin().lock().read_line(&mut line) { - Ok(0) => Ok(None), // EOF + Ok(0) => Ok(None), Ok(_) => Ok(Some(line.trim().to_string())), Err(e) if e.kind() == io::ErrorKind::Interrupted => Ok(None), Err(e) => Err(e), @@ -204,3 +314,105 @@ fn parse_value(input: &str, field_type: &str, enum_values: Option<&Vec>) _ => Value::String(input.to_string()), } } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use test_case::test_case; + + #[test] + fn builds_required_enum_select_with_default() { + let schema = json!({ + "type": "object", + "properties": { + "color": { + "type": "string", + "enum": ["red", "green"], + "default": "green" + } + }, + "required": ["color"] + }); + + let select = single_select(&schema).unwrap(); + + assert_eq!(select.field_name, "color"); + assert_eq!( + select.options, + vec![ + (SelectChoice::Value("red".to_string()), "red".to_string()), + ( + SelectChoice::Value("green".to_string()), + "green".to_string() + ) + ] + ); + assert_eq!( + select.initial_value, + Some(SelectChoice::Value("green".to_string())) + ); + } + + #[test] + fn optional_select_defaults_to_skip() { + let schema = json!({ + "type": "object", + "properties": { + "color": { "type": "string", "enum": ["", "green"] } + } + }); + + let select = single_select(&schema).unwrap(); + + assert_eq!(select.initial_value, Some(SelectChoice::Skip)); + assert_eq!(select.options.last().unwrap().0, SelectChoice::Skip); + assert_eq!( + select.options.first().unwrap().0, + SelectChoice::Value(String::new()) + ); + } + + #[test] + fn uses_one_of_titles_as_labels() { + let schema = json!({ + "type": "object", + "properties": { + "size": { + "oneOf": [ + { "const": "s", "title": "Small" }, + { "const": "l", "title": "Large" } + ] + } + }, + "required": ["size"] + }); + + let select = single_select(&schema).unwrap(); + + assert_eq!(select.options[0].0, SelectChoice::Value("s".to_string())); + assert_eq!(select.options[0].1, "Small"); + assert_eq!(select.options[1].0, SelectChoice::Value("l".to_string())); + assert_eq!(select.options[1].1, "Large"); + } + + #[test_case(json!({}); "missing properties")] + #[test_case(json!({ "properties": {} }); "empty properties")] + #[test_case(json!({ + "properties": { + "first": { "enum": ["a"] }, + "second": { "enum": ["b"] } + } + }); "multiple properties")] + #[test_case(json!({ + "properties": { "choice": { "enum": ["a", 2] } } + }); "non-string enum value")] + #[test_case(json!({ + "properties": { + "choice": { "oneOf": [{ "const": "a" }, { "type": "string" }] } + } + }); "oneOf branch without const")] + fn unsupported_schema_does_not_build_select(schema: Value) { + assert!(single_select(&schema).is_none()); + } +}