diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index 9c49e54d..824b20f6 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -38,6 +38,7 @@ pub async fn handle_configure() -> Result<(), Box> { 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> { .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::() { + 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> { }) .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::() { + 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> { cmd, args, envs: Envs::new(envs), + timeout: Some(timeout), }, })?; @@ -539,6 +558,14 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { }) .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::() { + 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> { name: name.clone(), uri, envs: Envs::new(envs), + timeout: Some(timeout), }, })?; diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 7a0720b8..cd13f838 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/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) diff --git a/crates/goose-server/src/routes/extension.rs b/crates/goose-server/src/routes/extension.rs index c327ea2b..21ab9f5b 100644 --- a/crates/goose-server/src/routes/extension.rs +++ b/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, + timeout: Option, }, /// 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, + timeout: Option, }, /// 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, }, } @@ -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. diff --git a/crates/goose/examples/agent.rs b/crates/goose/examples/agent.rs index 6ecd4477..5165b39d 100644 --- a/crates/goose/examples/agent.rs +++ b/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:"); diff --git a/crates/goose/src/agents/capabilities.rs b/crates/goose/src/agents/capabilities.rs index 2e95ec70..6896dfd0 100644 --- a/crates/goose/src/agents/capabilities.rs +++ b/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 = 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)) } }; diff --git a/crates/goose/src/agents/extension.rs b/crates/goose/src/agents/extension.rs index db5d5f3b..6aeafea0 100644 --- a/crates/goose/src/agents/extension.rs +++ b/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, }, /// Standard I/O client with command and arguments #[serde(rename = "stdio")] @@ -62,38 +67,43 @@ pub enum ExtensionConfig { args: Vec, #[serde(default)] envs: Envs, + timeout: Option, }, /// 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, }, } 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>(name: S, uri: S) -> Self { + pub fn sse, T: Into>(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>(name: S, cmd: S) -> Self { + pub fn stdio, T: Into>(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), } } } diff --git a/crates/goose/src/config/extensions.rs b/crates/goose/src/config/extensions.rs index ffbc955a..3abc3260 100644 --- a/crates/goose/src/config/extensions.rs +++ b/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), }, }, )]); diff --git a/crates/goose/src/config/mod.rs b/crates/goose/src/config/mod.rs index 3c809121..94191d90 100644 --- a/crates/goose/src/config/mod.rs +++ b/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; diff --git a/ui/desktop/src/components/settings/extensions/ManualExtensionModal.tsx b/ui/desktop/src/components/settings/extensions/ManualExtensionModal.tsx index 60d8b8b5..97259c98 100644 --- a/ui/desktop/src/components/settings/extensions/ManualExtensionModal.tsx +++ b/ui/desktop/src/components/settings/extensions/ManualExtensionModal.tsx @@ -2,7 +2,7 @@ import React, { useState } from 'react'; import { Card } from '../../ui/card'; import { Button } from '../../ui/button'; import { Input } from '../../ui/input'; -import { FullExtensionConfig } from '../../../extensions'; +import { FullExtensionConfig, DEFAULT_EXTENSION_TIMEOUT } from '../../../extensions'; import { toast } from 'react-toastify'; import Select from 'react-select'; import { createDarkSelectStyles, darkSelectTheme } from '../../ui/select-styles'; @@ -22,6 +22,7 @@ export function ManualExtensionModal({ isOpen, onClose, onSubmit }: ManualExtens enabled: true, args: [], commandInput: '', + timeout: DEFAULT_EXTENSION_TIMEOUT, }); const [envKey, setEnvKey] = useState(''); const [envValue, setEnvValue] = useState(''); @@ -267,8 +268,20 @@ export function ManualExtensionModal({ isOpen, onClose, onSubmit }: ManualExtens )} - +
+ + setFormData({ ...formData, timeout: parseInt(e.target.value) })} + className="w-full" + required + /> +
+