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

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;