feat(goose): support customizing extension timeout (#1428)

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

@@ -38,6 +38,7 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
                     enabled: true,
                     config: ExtensionConfig::Builtin {
                         name: "developer".to_string(),
+                        timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
                     },
                 })?;
             }
@@ -437,10 +438,19 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
                 .interact()?
                 .to_string();
 
+            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 valide timeout"),
+                })
+                .interact()?;
+
             ExtensionManager::set(ExtensionEntry {
                 enabled: true,
                 config: ExtensionConfig::Builtin {
                     name: extension.clone(),
+                    timeout: Some(timeout),
                 },
             })?;
 
@@ -472,6 +482,14 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
                 })
                 .interact()?;
 
+            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 valide timeout"),
+                })
+                .interact()?;
+
             // Split the command string into command and args
             let mut parts = command_str.split_whitespace();
             let cmd = parts.next().unwrap_or("").to_string();
@@ -506,6 +524,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
                     cmd,
                     args,
                     envs: Envs::new(envs),
+                    timeout: Some(timeout),
                 },
             })?;
 
@@ -539,6 +558,14 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
                 })
                 .interact()?;
 
+            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 valide timeout"),
+                })
+                .interact()?;
+
             let add_env =
                 cliclack::confirm("Would you like to add environment variables?").interact()?;
 
@@ -567,6 +594,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
                     name: name.clone(),
                     uri,
                     envs: Envs::new(envs),
+                    timeout: Some(timeout),
                 },
             })?;
 

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

@@ -107,6 +107,8 @@ impl Session {
             cmd,
             args: parts.iter().map(|s| s.to_string()).collect(),
             envs: Envs::new(envs),
+            // TODO: should set timeout
+            timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
         };
 
         self.agent
@@ -128,6 +130,8 @@ impl Session {
         for name in builtin_name.split(',') {
             let config = ExtensionConfig::Builtin {
                 name: name.trim().to_string(),
+                // TODO: should set a timeout
+                timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
             };
             self.agent
                 .add_extension(config)

crates/goose-server/src/routes/extension.rs

@@ -23,6 +23,7 @@ enum ExtensionConfigRequest {
         /// List of environment variable keys. The server will fetch their values from the keyring.
         #[serde(default)]
         env_keys: Vec<String>,
+        timeout: Option<u64>,
     },
     /// Standard I/O (stdio) extension.
     #[serde(rename = "stdio")]
@@ -37,12 +38,14 @@ enum ExtensionConfigRequest {
         /// List of environment variable keys. The server will fetch their values from the keyring.
         #[serde(default)]
         env_keys: Vec<String>,
+        timeout: Option<u64>,
     },
     /// Built-in extension that is part of the goose binary.
     #[serde(rename = "builtin")]
     Builtin {
         /// The name of the built-in extension.
         name: String,
+        timeout: Option<u64>,
     },
 }
 
@@ -84,6 +87,7 @@ async fn add_extension(
             name,
             uri,
             env_keys,
+            timeout,
         } => {
             let mut env_map = HashMap::new();
             for key in env_keys {
@@ -111,6 +115,7 @@ async fn add_extension(
                 name,
                 uri,
                 envs: Envs::new(env_map),
+                timeout,
             }
         }
         ExtensionConfigRequest::Stdio {
@@ -118,6 +123,7 @@ async fn add_extension(
             cmd,
             args,
             env_keys,
+            timeout,
         } => {
             let mut env_map = HashMap::new();
             for key in env_keys {
@@ -146,9 +152,12 @@ async fn add_extension(
                 cmd,
                 args,
                 envs: Envs::new(env_map),
+                timeout,
             }
         }
-        ExtensionConfigRequest::Builtin { name } => ExtensionConfig::Builtin { name },
+        ExtensionConfigRequest::Builtin { name, timeout } => {
+            ExtensionConfig::Builtin { name, timeout }
+        }
     };
 
     // Acquire a lock on the agent and attempt to add the extension.

crates/goose/examples/agent.rs

@@ -1,6 +1,7 @@
 use dotenv::dotenv;
 use futures::StreamExt;
 use goose::agents::{AgentFactory, ExtensionConfig};
+use goose::config::DEFAULT_EXTENSION_TIMEOUT;
 use goose::message::Message;
 use goose::providers::databricks::DatabricksProvider;
 
@@ -14,7 +15,11 @@ async fn main() {
     // Setup an agent with the developer extension
     let mut agent = AgentFactory::create("reference", provider).expect("default should exist");
 
-    let config = ExtensionConfig::stdio("developer", "./target/debug/developer");
+    let config = ExtensionConfig::stdio(
+        "developer",
+        "./target/debug/developer",
+        DEFAULT_EXTENSION_TIMEOUT,
+    );
     agent.add_extension(config).await.unwrap();
 
     println!("Extensions:");

crates/goose/src/agents/capabilities.rs

@@ -105,21 +105,37 @@ impl Capabilities {
     // TODO IMPORTANT need to ensure this times out if the extension command is broken!
     pub async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()> {
         let mut client: Box<dyn McpClientTrait> = match &config {
-            ExtensionConfig::Sse { uri, envs, .. } => {
+            ExtensionConfig::Sse {
+                uri, envs, timeout, ..
+            } => {
                 let transport = SseTransport::new(uri, envs.get_env());
                 let handle = transport.start().await?;
-                let service = McpService::with_timeout(handle, Duration::from_secs(300));
+                let service = McpService::with_timeout(
+                    handle,
+                    Duration::from_secs(
+                        timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
+                    ),
+                );
                 Box::new(McpClient::new(service))
             }
             ExtensionConfig::Stdio {
-                cmd, args, envs, ..
+                cmd,
+                args,
+                envs,
+                timeout,
+                ..
             } => {
                 let transport = StdioTransport::new(cmd, args.to_vec(), envs.get_env());
                 let handle = transport.start().await?;
-                let service = McpService::with_timeout(handle, Duration::from_secs(300));
+                let service = McpService::with_timeout(
+                    handle,
+                    Duration::from_secs(
+                        timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
+                    ),
+                );
                 Box::new(McpClient::new(service))
             }
-            ExtensionConfig::Builtin { name } => {
+            ExtensionConfig::Builtin { name, timeout } => {
                 // For builtin extensions, we run the current executable with mcp and extension name
                 let cmd = std::env::current_exe()
                     .expect("should find the current executable")
@@ -132,7 +148,12 @@ impl Capabilities {
                     HashMap::new(),
                 );
                 let handle = transport.start().await?;
-                let service = McpService::with_timeout(handle, Duration::from_secs(300));
+                let service = McpService::with_timeout(
+                    handle,
+                    Duration::from_secs(
+                        timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
+                    ),
+                );
                 Box::new(McpClient::new(service))
             }
         };

crates/goose/src/agents/extension.rs

@@ -4,6 +4,8 @@ use mcp_client::client::Error as ClientError;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
+use crate::config;
+
 /// Errors from Extension operation
 #[derive(Error, Debug)]
 pub enum ExtensionError {
@@ -52,6 +54,9 @@ pub enum ExtensionConfig {
         uri: String,
         #[serde(default)]
         envs: Envs,
+        // NOTE: set timeout to be optional for compatibility.
+        // However, new configurations should include this field.
+        timeout: Option<u64>,
     },
     /// Standard I/O client with command and arguments
     #[serde(rename = "stdio")]
@@ -62,38 +67,43 @@ pub enum ExtensionConfig {
         args: Vec<String>,
         #[serde(default)]
         envs: Envs,
+        timeout: Option<u64>,
     },
     /// Built-in extension that is part of the goose binary
     #[serde(rename = "builtin")]
     Builtin {
         /// The name used to identify this extension
         name: String,
+        timeout: Option<u64>,
     },
 }
 
 impl Default for ExtensionConfig {
     fn default() -> Self {
         Self::Builtin {
-            name: String::from("default"),
+            name: config::DEFAULT_EXTENSION.to_string(),
+            timeout: Some(config::DEFAULT_EXTENSION_TIMEOUT),
         }
     }
 }
 
 impl ExtensionConfig {
-    pub fn sse<S: Into<String>>(name: S, uri: S) -> Self {
+    pub fn sse<S: Into<String>, T: Into<u64>>(name: S, uri: S, timeout: T) -> Self {
         Self::Sse {
             name: name.into(),
             uri: uri.into(),
             envs: Envs::default(),
+            timeout: Some(timeout.into()),
         }
     }
 
-    pub fn stdio<S: Into<String>>(name: S, cmd: S) -> Self {
+    pub fn stdio<S: Into<String>, T: Into<u64>>(name: S, cmd: S, timeout: T) -> Self {
         Self::Stdio {
             name: name.into(),
             cmd: cmd.into(),
             args: vec![],
             envs: Envs::default(),
+            timeout: Some(timeout.into()),
         }
     }
 
@@ -104,12 +114,17 @@ impl ExtensionConfig {
     {
         match self {
             Self::Stdio {
-                name, cmd, envs, ..
+                name,
+                cmd,
+                envs,
+                timeout,
+                ..
             } => Self::Stdio {
                 name,
                 cmd,
                 envs,
                 args: args.into_iter().map(Into::into).collect(),
+                timeout,
             },
             other => other,
         }
@@ -120,7 +135,7 @@ impl ExtensionConfig {
         match self {
             Self::Sse { name, .. } => name,
             Self::Stdio { name, .. } => name,
-            Self::Builtin { name } => name,
+            Self::Builtin { name, .. } => name,
         }
     }
 }
@@ -134,7 +149,7 @@ impl std::fmt::Display for ExtensionConfig {
             } => {
                 write!(f, "Stdio({}: {} {})", name, cmd, args.join(" "))
             }
-            ExtensionConfig::Builtin { name } => write!(f, "Builtin({})", name),
+            ExtensionConfig::Builtin { name, .. } => write!(f, "Builtin({})", name),
         }
     }
 }

crates/goose/src/config/extensions.rs

@@ -4,7 +4,8 @@ use anyhow::Result;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
-const DEFAULT_EXTENSION: &str = "developer";
+pub const DEFAULT_EXTENSION: &str = "developer";
+pub const DEFAULT_EXTENSION_TIMEOUT: u64 = 300;
 
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct ExtensionEntry {
@@ -32,6 +33,7 @@ impl ExtensionManager {
                         enabled: true,
                         config: ExtensionConfig::Builtin {
                             name: DEFAULT_EXTENSION.to_string(),
+                            timeout: Some(DEFAULT_EXTENSION_TIMEOUT),
                         },
                     },
                 )]);

crates/goose/src/config/mod.rs

@@ -6,3 +6,6 @@ pub use crate::agents::ExtensionConfig;
 pub use base::{Config, ConfigError, APP_STRATEGY};
 pub use experiments::ExperimentManager;
 pub use extensions::{ExtensionEntry, ExtensionManager};
+
+pub use extensions::DEFAULT_EXTENSION;
+pub use extensions::DEFAULT_EXTENSION_TIMEOUT;