feat(cli): allow custom timeout for streamable HTTP extensions

crates/goose-cli/src/cli.rs

@@ -102,6 +102,38 @@ pub struct SessionOptions {
     pub max_turns: Option<u32>,
 }
 
+#[derive(Debug, Clone)]
+pub struct StreamableHttpOptions {
+    pub url: String,
+    pub timeout: u64,
+}
+
+fn parse_streamable_http_extension(input: &str) -> Result<StreamableHttpOptions, String> {
+    let mut input_iter = input.split_whitespace();
+    let (mut url, mut timeout) = (String::new(), goose::config::DEFAULT_EXTENSION_TIMEOUT);
+
+    if let Some(url_str) = input_iter.next() {
+        url.push_str(url_str);
+    }
+
+    for kv_pair in input_iter {
+        if !kv_pair.contains('=') {
+            continue;
+        }
+
+        let (key, value) = kv_pair.split_once('=').unwrap();
+
+        // We Can have more keys here for setting other properties
+        if key == "timeout" {
+            if let Ok(seconds) = value.parse::<u64>() {
+                timeout = seconds;
+            }
+        }
+    }
+
+    Ok(StreamableHttpOptions { url, timeout })
+}
+
 /// Extension configuration options shared between Session and Run commands
 #[derive(Args, Debug, Clone, Default)]
 pub struct ExtensionOptions {
@@ -118,10 +150,11 @@ pub struct ExtensionOptions {
         long = "with-streamable-http-extension",
         value_name = "URL",
         help = "Add streamable HTTP extensions (can be specified multiple times)",
-        long_help = "Add streamable HTTP extensions from a URL. Can be specified multiple times. Format: 'url...'",
-        action = clap::ArgAction::Append
+        long_help = "Add streamable HTTP extensions from a URL. Can be specified multiple times. Format: 'url...' or 'url... timeout=100' to set up timeout other than default",
+        action = clap::ArgAction::Append,
+        value_parser = parse_streamable_http_extension
     )]
-    pub streamable_http_extensions: Vec<String>,
+    pub streamable_http_extensions: Vec<StreamableHttpOptions>,
 
     #[arg(
         long = "with-builtin",

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

@@ -1,3 +1,4 @@
+use crate::cli::StreamableHttpOptions;
 use crate::session::build_session;
 use crate::session::SessionBuilderConfig;
 use crate::{logging, CliSession};
@@ -34,12 +35,20 @@ pub async fn agent_generator(
     requirements: ExtensionRequirements,
     session_id: String,
 ) -> BenchAgent {
+    let streamable_http_extensions: Vec<StreamableHttpOptions> = requirements
+        .streamable_http
+        .iter()
+        .map(|s| StreamableHttpOptions {
+            url: s.clone(),
+            timeout: goose::config::DEFAULT_EXTENSION_TIMEOUT,
+        })
+        .collect();
     let base_session = build_session(SessionBuilderConfig {
         session_id: Some(session_id),
         resume: false,
         no_session: false,
         extensions: requirements.external,
-        streamable_http_extensions: requirements.streamable_http,
+        streamable_http_extensions,
         builtins: requirements.builtin,
         extensions_override: None,
         additional_system_prompt: None,

crates/goose-cli/src/session/builder.rs

@@ -1,3 +1,5 @@
+use crate::cli::StreamableHttpOptions;
+
 use super::output;
 use super::CliSession;
 use console::style;
@@ -33,7 +35,7 @@ pub struct SessionBuilderConfig {
     /// List of stdio extension commands to add
     pub extensions: Vec<String>,
     /// List of streamable HTTP extension commands to add
-    pub streamable_http_extensions: Vec<String>,
+    pub streamable_http_extensions: Vec<StreamableHttpOptions>,
     /// List of builtin extension commands to add
     pub builtins: Vec<String>,
     /// List of extensions to enable, enable only this set and ignore configured ones
@@ -584,23 +586,23 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
     }
 
     // Add streamable HTTP extensions if provided
-    for extension_str in session_config.streamable_http_extensions {
+    for extension_options in session_config.streamable_http_extensions {
         if let Err(e) = session
-            .add_streamable_http_extension(extension_str.clone())
+            .add_streamable_http_extension(extension_options.clone())
             .await
         {
             eprintln!(
                 "{}",
                 style(format!(
                     "Warning: Failed to start streamable HTTP extension '{}' ({}), continuing without it",
-                    extension_str, e
+                    extension_options.url, e
                 ))
                 .yellow()
             );
 
             // Offer debugging help
             if let Err(debug_err) = offer_extension_debugging_help(
-                &extension_str,
+                &extension_options.url,
                 &e.to_string(),
                 Arc::clone(&provider_for_display),
                 session_config.interactive,
@@ -695,7 +697,10 @@ mod tests {
             resume: false,
             no_session: false,
             extensions: vec!["echo test".to_string()],
-            streamable_http_extensions: vec!["http://localhost:8080/mcp".to_string()],
+            streamable_http_extensions: vec![StreamableHttpOptions {
+                url: "http://localhost:8080/mcp".to_string(),
+                timeout: goose::config::DEFAULT_EXTENSION_TIMEOUT,
+            }],
             builtins: vec!["developer".to_string()],
             extensions_override: None,
             additional_system_prompt: Some("Test prompt".to_string()),

crates/goose-cli/src/session/mod.rs

@@ -8,6 +8,7 @@ mod prompt;
 mod task_execution_display;
 mod thinking;
 
+use crate::cli::StreamableHttpOptions;
 use crate::session::task_execution_display::{
     format_task_execution_notification, TASK_EXECUTION_NOTIFICATION_TYPE,
 };
@@ -306,17 +307,20 @@ impl CliSession {
     /// Add a streamable HTTP extension to the session
     ///
     /// # Arguments
-    /// * `extension_url` - URL of the server
-    pub async fn add_streamable_http_extension(&mut self, extension_url: String) -> Result<()> {
+    /// * `extension_options` - Options including both URL and timeout
+    ///   See [`crate::cli::StreamableHttpOptions`] for details
+    pub async fn add_streamable_http_extension(
+        &mut self,
+        extension_options: StreamableHttpOptions,
+    ) -> Result<()> {
         let config = ExtensionConfig::StreamableHttp {
             name: String::new(),
-            uri: extension_url,
+            uri: extension_options.url,
             envs: Envs::new(HashMap::new()),
             env_keys: Vec::new(),
             headers: HashMap::new(),
             description: goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string(),
-            // TODO: should set timeout
-            timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
+            timeout: Some(extension_options.timeout),
             bundled: None,
             available_tools: Vec::new(),
         };