diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index b020fe7a5125..c79e6e372acf 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -102,6 +102,38 @@ pub struct SessionOptions { pub max_turns: Option, } +#[derive(Debug, Clone)] +pub struct StreamableHttpOptions { + pub url: String, + pub timeout: u64, +} + +fn parse_streamable_http_extension(input: &str) -> Result { + 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::() { + 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, + pub streamable_http_extensions: Vec, #[arg( long = "with-builtin", diff --git a/crates/goose-cli/src/commands/bench.rs b/crates/goose-cli/src/commands/bench.rs index 7d4522ff118f..caa63835e4ee 100644 --- a/crates/goose-cli/src/commands/bench.rs +++ b/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 = 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, diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index a6c132355a54..c8aaa9708bec 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/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, /// List of streamable HTTP extension commands to add - pub streamable_http_extensions: Vec, + pub streamable_http_extensions: Vec, /// List of builtin extension commands to add pub builtins: Vec, /// 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()), diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index a7678327b731..19efdbc38c01 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/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(), };