chore: refactor configure_extensions_dialog to reduce line count (#6277)

crates/goose-cli/src/commands/configure.rs

@@ -690,28 +690,109 @@ pub fn toggle_extensions_dialog() -> anyhow::Result<()> {
     Ok(())
 }
 
-pub fn configure_extensions_dialog() -> anyhow::Result<()> {
+fn prompt_extension_timeout() -> anyhow::Result<u64> {
-    let extension_type = cliclack::select("What type of extension would you like to add?")
+    Ok(
-        .item(
+        cliclack::input("Please set the timeout for this tool (in secs):")
-            "built-in",
+            .placeholder(&goose::config::DEFAULT_EXTENSION_TIMEOUT.to_string())
-            "Built-in Extension",
+            .validate(|input: &String| match input.parse::<u64>() {
-            "Use an extension that comes with goose",
+                Ok(_) => Ok(()),
+                Err(_) => Err("Please enter a valid timeout"),
+            })
+            .interact()?,
     )
-        .item(
+}
-            "stdio",
+
-            "Command-line Extension",
+fn prompt_extension_description() -> anyhow::Result<String> {
-            "Run a local command or script",
+    Ok(cliclack::input("Enter a description for this extension:")
-        )
+        .placeholder("Description")
-        .item(
+        .validate(|input: &String| {
-            "streamable_http",
+            if input.trim().is_empty() {
-            "Remote Extension (Streamable HTTP)",
+                Err("Please enter a valid description")
-            "Connect to a remote extension via MCP Streamable HTTP",
+            } else {
+                Ok(())
+            }
+        })
+        .interact()?)
+}
+
+fn prompt_extension_name(placeholder: &str) -> anyhow::Result<String> {
+    let extensions = get_all_extension_names();
+    Ok(
+        cliclack::input("What would you like to call this extension?")
+            .placeholder(placeholder)
+            .validate(move |input: &String| {
+                if input.is_empty() {
+                    Err("Please enter a name")
+                } else if extensions.contains(input) {
+                    Err("An extension with this name already exists")
+                } else {
+                    Ok(())
+                }
+            })
+            .interact()?,
     )
+}
+
+fn collect_env_vars() -> anyhow::Result<(HashMap<String, String>, Vec<String>)> {
+    let mut envs = HashMap::new();
+    let mut env_keys = Vec::new();
+    let config = Config::global();
+
+    if !cliclack::confirm("Would you like to add environment variables?").interact()? {
+        return Ok((envs, env_keys));
+    }
+
+    loop {
+        let key: String = cliclack::input("Environment variable name:")
+            .placeholder("API_KEY")
             .interact()?;
 
-    match extension_type {
+        let value: String = cliclack::password("Environment variable value:")
-        // TODO we'll want a place to collect all these options, maybe just an enum in goose-mcp
+            .mask('▪')
-        "built-in" => {
+            .interact()?;
+
+        match config.set_secret(&key, &value) {
+            Ok(_) => env_keys.push(key),
+            Err(_) => {
+                envs.insert(key, value);
+            }
+        }
+
+        if !cliclack::confirm("Add another environment variable?").interact()? {
+            break;
+        }
+    }
+
+    Ok((envs, env_keys))
+}
+
+fn collect_headers() -> anyhow::Result<HashMap<String, String>> {
+    let mut headers = HashMap::new();
+
+    if !cliclack::confirm("Would you like to add custom headers?").interact()? {
+        return Ok(headers);
+    }
+
+    loop {
+        let key: String = cliclack::input("Header name:")
+            .placeholder("Authorization")
+            .interact()?;
+
+        let value: String = cliclack::input("Header value:")
+            .placeholder("Bearer token123")
+            .interact()?;
+
+        headers.insert(key, value);
+
+        if !cliclack::confirm("Add another header?").interact()? {
+            break;
+        }
+    }
+
+    Ok(headers)
+}
+
+fn configure_builtin_extension() -> anyhow::Result<()> {
     let extensions = vec![
         (
             "autovisualiser",
@@ -745,14 +826,7 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
         select = select.item(id, name, desc);
     }
     let extension = select.interact()?.to_string();
-
+    let timeout = prompt_extension_timeout()?;
-            let timeout: u64 = cliclack::input("Please set the timeout for this tool (in secs):")
-                .placeholder(&goose::config::DEFAULT_EXTENSION_TIMEOUT.to_string())
-                .validate(|input: &String| match input.parse::<u64>() {
-                    Ok(_) => Ok(()),
-                    Err(_) => Err("Please enter a valid timeout"),
-                })
-                .interact()?;
 
     let (display_name, description) = extensions
         .iter()
@@ -773,21 +847,11 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
     });
 
     cliclack::outro(format!("Enabled {} extension", style(extension).green()))?;
-        }
-        "stdio" => {
-            let extensions = get_all_extension_names();
-            let name: String = cliclack::input("What would you like to call this extension?")
-                .placeholder("my-extension")
-                .validate(move |input: &String| {
-                    if input.is_empty() {
-                        Err("Please enter a name")
-                    } else if extensions.contains(input) {
-                        Err("An extension with this name already exists")
-                    } else {
     Ok(())
 }
-                })
+
-                .interact()?;
+fn configure_stdio_extension() -> anyhow::Result<()> {
+    let name = prompt_extension_name("my-extension")?;
 
     let command_str: String = cliclack::input("What command should be run?")
         .placeholder("npx -y @block/gdrive")
@@ -800,63 +864,14 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
         })
         .interact()?;
 
-            let timeout: u64 = cliclack::input("Please set the timeout for this tool (in secs):")
+    let timeout = prompt_extension_timeout()?;
-                .placeholder(&goose::config::DEFAULT_EXTENSION_TIMEOUT.to_string())
-                .validate(|input: &String| match input.parse::<u64>() {
-                    Ok(_) => Ok(()),
-                    Err(_) => Err("Please enter a valid timeout"),
-                })
-                .interact()?;
 
-            // Split the command string into command and args
-            // TODO: find a way to expose this to the frontend so we dont need to re-write code
     let mut parts = command_str.split_whitespace();
     let cmd = parts.next().unwrap_or("").to_string();
     let args: Vec<String> = parts.map(String::from).collect();
 
-            let description = cliclack::input("Enter a description for this extension:")
+    let description = prompt_extension_description()?;
-                .placeholder("Description")
+    let (envs, env_keys) = collect_env_vars()?;
-                .validate(|input: &String| match input.parse::<String>() {
-                    Ok(_) => Ok(()),
-                    Err(_) => Err("Please enter a valid description"),
-                })
-                .interact()?;
-
-            let add_env =
-                cliclack::confirm("Would you like to add environment variables?").interact()?;
-
-            let mut envs = HashMap::new();
-            let mut env_keys = Vec::new();
-            let config = Config::global();
-
-            if add_env {
-                loop {
-                    let key: String = cliclack::input("Environment variable name:")
-                        .placeholder("API_KEY")
-                        .interact()?;
-
-                    let value: String = cliclack::password("Environment variable value:")
-                        .mask('▪')
-                        .interact()?;
-
-                    // Try to store in keychain
-                    let keychain_key = key.to_string();
-                    match config.set_secret(&keychain_key, &value) {
-                        Ok(_) => {
-                            // Successfully stored in keychain, add to env_keys
-                            env_keys.push(keychain_key);
-                        }
-                        Err(_) => {
-                            // Failed to store in keychain, store directly in envs
-                            envs.insert(key, value);
-                        }
-                    }
-
-                    if !cliclack::confirm("Add another environment variable?").interact()? {
-                        break;
-                    }
-                }
-            }
 
     set_extension(ExtensionEntry {
         enabled: true,
@@ -874,23 +889,13 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
     });
 
     cliclack::outro(format!("Added {} extension", style(name).green()))?;
-        }
-        "streamable_http" => {
-            let extensions = get_all_extension_names();
-            let name: String = cliclack::input("What would you like to call this extension?")
-                .placeholder("my-remote-extension")
-                .validate(move |input: &String| {
-                    if input.is_empty() {
-                        Err("Please enter a name")
-                    } else if extensions.contains(input) {
-                        Err("An extension with this name already exists")
-                    } else {
     Ok(())
 }
-                })
-                .interact()?;
 
-            let uri: String = cliclack::input("What is the Streamable HTTP endpoint URI?")
+fn configure_streamable_http_extension() -> anyhow::Result<()> {
+    let name = prompt_extension_name("my-remote-extension")?;
+
+    let uri: String = cliclack::input("What is the Streaming HTTP endpoint URI?")
         .placeholder("http://localhost:8000/messages")
         .validate(|input: &String| {
             if input.is_empty() {
@@ -903,81 +908,13 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
         })
         .interact()?;
 
-            let timeout: u64 = cliclack::input("Please set the timeout for this tool (in secs):")
+    let timeout = prompt_extension_timeout()?;
-                .placeholder(&goose::config::DEFAULT_EXTENSION_TIMEOUT.to_string())
+    let description = prompt_extension_description()?;
-                .validate(|input: &String| match input.parse::<u64>() {
+    let headers = collect_headers()?;
-                    Ok(_) => Ok(()),
-                    Err(_) => Err("Please enter a valid timeout"),
-                })
-                .interact()?;
 
-            let description = cliclack::input("Enter a description for this extension:")
+    // Original behavior: no env var collection for Streamable HTTP
-                .placeholder("Description")
+    let envs = HashMap::new();
-                .validate(|input: &String| {
+    let env_keys = Vec::new();
-                    if input.trim().is_empty() {
-                        Err("Please enter a valid description")
-                    } else {
-                        Ok(())
-                    }
-                })
-                .interact()?;
-
-            let add_headers =
-                cliclack::confirm("Would you like to add custom headers?").interact()?;
-
-            let mut headers = HashMap::new();
-            if add_headers {
-                loop {
-                    let key: String = cliclack::input("Header name:")
-                        .placeholder("Authorization")
-                        .interact()?;
-
-                    let value: String = cliclack::input("Header value:")
-                        .placeholder("Bearer token123")
-                        .interact()?;
-
-                    headers.insert(key, value);
-
-                    if !cliclack::confirm("Add another header?").interact()? {
-                        break;
-                    }
-                }
-            }
-
-            let add_env = false; // No env prompt for Streamable HTTP
-
-            let mut envs = HashMap::new();
-            let mut env_keys = Vec::new();
-            let config = Config::global();
-
-            if add_env {
-                loop {
-                    let key: String = cliclack::input("Environment variable name:")
-                        .placeholder("API_KEY")
-                        .interact()?;
-
-                    let value: String = cliclack::password("Environment variable value:")
-                        .mask('▪')
-                        .interact()?;
-
-                    // Try to store in keychain
-                    let keychain_key = key.to_string();
-                    match config.set_secret(&keychain_key, &Value::String(value.clone())) {
-                        Ok(_) => {
-                            // Successfully stored in keychain, add to env_keys
-                            env_keys.push(keychain_key);
-                        }
-                        Err(_) => {
-                            // Failed to store in keychain, store directly in envs
-                            envs.insert(key, value);
-                        }
-                    }
-
-                    if !cliclack::confirm("Add another environment variable?").interact()? {
-                        break;
-                    }
-                }
-            }
 
     set_extension(ExtensionEntry {
         enabled: true,
@@ -995,12 +932,36 @@ pub fn configure_extensions_dialog() -> anyhow::Result<()> {
     });
 
     cliclack::outro(format!("Added {} extension", style(name).green()))?;
+    Ok(())
 }
+
+pub fn configure_extensions_dialog() -> anyhow::Result<()> {
+    let extension_type = cliclack::select("What type of extension would you like to add?")
+        .item(
+            "built-in",
+            "Built-in Extension",
+            "Use an extension that comes with goose",
+        )
+        .item(
+            "stdio",
+            "Command-line Extension",
+            "Run a local command or script",
+        )
+        .item(
+            "streamable_http",
+            "Remote Extension (Streamable HTTP)",
+            "Connect to a remote extension via MCP Streamable HTTP",
+        )
+        .interact()?;
+
+    match extension_type {
+        "built-in" => configure_builtin_extension()?,
+        "stdio" => configure_stdio_extension()?,
+        "streamable_http" => configure_streamable_http_extension()?,
         _ => unreachable!(),
     };
 
     print_config_file_saved()?;
-
     Ok(())
 }