Jazzcort add timeout (#6772)

Signed-off-by: Jazzcort <jason101011113@gmail.com>
Co-authored-by: Jazzcort <jason101011113@gmail.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

crates/goose-cli/src/cli.rs

@@ -112,6 +112,38 @@ pub struct SessionOptions {
     pub container: Option<String>,
 }
 
+#[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 {
@@ -128,10 +160,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,13 +35,21 @@ 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,
         fork: false,
         no_session: false,
         extensions: requirements.external,
-        streamable_http_extensions: requirements.streamable_http,
+        streamable_http_extensions,
         builtins: requirements.builtin,
         recipe: 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;
@@ -87,7 +89,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>,
     /// Recipe for the session
@@ -578,6 +580,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
     )
     .await;
 
+
     if let Err(e) = session
         .agent
         .persist_extension_state(&session_id.clone())
@@ -629,7 +632,10 @@ mod tests {
             fork: 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()],
             recipe: None,
             additional_system_prompt: Some("Test prompt".to_string()),

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

@@ -9,6 +9,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,
 };