feat: configurable extension timeouts via ACP _meta and global default (#8295)
Signed-off-by: Bradley Axen <baxen@squareup.com> Signed-off-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
@@ -31,7 +31,7 @@ use sacp::schema::{
|
||||
ConfigOptionUpdate, Content, ContentBlock, ContentChunk, CurrentModeUpdate, EmbeddedResource,
|
||||
EmbeddedResourceResource, FileSystemCapabilities, ImageContent, InitializeRequest,
|
||||
InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest,
|
||||
LoadSessionResponse, McpCapabilities, McpServer, ModelId, ModelInfo, NewSessionRequest,
|
||||
LoadSessionResponse, McpCapabilities, McpServer, Meta, ModelId, ModelInfo, NewSessionRequest,
|
||||
NewSessionResponse, PermissionOption, PermissionOptionKind, PromptCapabilities, PromptRequest,
|
||||
PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, ResourceLink,
|
||||
SessionCapabilities, SessionCloseCapabilities, SessionConfigOption,
|
||||
@@ -90,34 +90,46 @@ pub struct GooseAcpAgent {
|
||||
disable_session_naming: bool,
|
||||
}
|
||||
|
||||
fn extract_timeout_from_meta(meta: &Option<Meta>) -> Option<u64> {
|
||||
meta.as_ref()
|
||||
.and_then(|m| m.get("timeout"))
|
||||
.and_then(|v| v.as_u64())
|
||||
}
|
||||
|
||||
fn mcp_server_to_extension_config(mcp_server: McpServer) -> Result<ExtensionConfig, String> {
|
||||
match mcp_server {
|
||||
McpServer::Stdio(stdio) => Ok(ExtensionConfig::Stdio {
|
||||
name: stdio.name,
|
||||
description: String::new(),
|
||||
cmd: stdio.command.to_string_lossy().to_string(),
|
||||
args: stdio.args,
|
||||
envs: Envs::new(stdio.env.into_iter().map(|e| (e.name, e.value)).collect()),
|
||||
env_keys: vec![],
|
||||
timeout: None,
|
||||
bundled: Some(false),
|
||||
available_tools: vec![],
|
||||
}),
|
||||
McpServer::Http(http) => Ok(ExtensionConfig::StreamableHttp {
|
||||
name: http.name,
|
||||
description: String::new(),
|
||||
uri: http.url,
|
||||
envs: Envs::default(),
|
||||
env_keys: vec![],
|
||||
headers: http
|
||||
.headers
|
||||
.into_iter()
|
||||
.map(|h| (h.name, h.value))
|
||||
.collect(),
|
||||
timeout: None,
|
||||
bundled: Some(false),
|
||||
available_tools: vec![],
|
||||
}),
|
||||
McpServer::Stdio(stdio) => {
|
||||
let timeout = extract_timeout_from_meta(&stdio.meta);
|
||||
Ok(ExtensionConfig::Stdio {
|
||||
name: stdio.name,
|
||||
description: String::new(),
|
||||
cmd: stdio.command.to_string_lossy().to_string(),
|
||||
args: stdio.args,
|
||||
envs: Envs::new(stdio.env.into_iter().map(|e| (e.name, e.value)).collect()),
|
||||
env_keys: vec![],
|
||||
timeout,
|
||||
bundled: Some(false),
|
||||
available_tools: vec![],
|
||||
})
|
||||
}
|
||||
McpServer::Http(http) => {
|
||||
let timeout = extract_timeout_from_meta(&http.meta);
|
||||
Ok(ExtensionConfig::StreamableHttp {
|
||||
name: http.name,
|
||||
description: String::new(),
|
||||
uri: http.url,
|
||||
envs: Envs::default(),
|
||||
env_keys: vec![],
|
||||
headers: http
|
||||
.headers
|
||||
.into_iter()
|
||||
.map(|h| (h.name, h.value))
|
||||
.collect(),
|
||||
timeout,
|
||||
bundled: Some(false),
|
||||
available_tools: vec![],
|
||||
})
|
||||
}
|
||||
McpServer::Sse(_) => Err("SSE is unsupported, migrate to streamable_http".to_string()),
|
||||
_ => Err("Unknown MCP server type".to_string()),
|
||||
}
|
||||
|
||||
@@ -58,6 +58,14 @@ static RE_ENV_BRACES: Lazy<regex::Regex> =
|
||||
static RE_ENV_SIMPLE: Lazy<regex::Regex> =
|
||||
Lazy::new(|| regex::Regex::new(r"\$([A-Za-z_][A-Za-z0-9_]*)").expect("valid regex"));
|
||||
|
||||
fn resolve_timeout(timeout: Option<u64>) -> u64 {
|
||||
timeout.unwrap_or_else(|| {
|
||||
Config::global()
|
||||
.get_goose_default_extension_timeout()
|
||||
.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)
|
||||
})
|
||||
}
|
||||
|
||||
struct Extension {
|
||||
pub config: ExtensionConfig,
|
||||
/// Resolved config snapshot (with secrets from keyring substituted)
|
||||
@@ -275,7 +283,7 @@ async fn child_process_client(
|
||||
|
||||
let client_result = McpClient::connect_with_container(
|
||||
transport,
|
||||
Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)),
|
||||
Duration::from_secs(resolve_timeout(*timeout)),
|
||||
provider,
|
||||
docker_container,
|
||||
client_name,
|
||||
@@ -441,8 +449,7 @@ async fn create_streamable_http_client(
|
||||
},
|
||||
);
|
||||
|
||||
let timeout_duration =
|
||||
Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT));
|
||||
let timeout_duration = Duration::from_secs(resolve_timeout(timeout));
|
||||
|
||||
let client_res = McpClient::connect(
|
||||
transport,
|
||||
@@ -636,7 +643,7 @@ impl ExtensionManager {
|
||||
(def.client_factory)(context)
|
||||
} else {
|
||||
// Builtin MCP server extension
|
||||
let timeout_secs = timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT);
|
||||
let timeout_secs = resolve_timeout(timeout);
|
||||
let extension_fn =
|
||||
get_builtin_extension(normalized_name.as_str()).ok_or_else(|| {
|
||||
ExtensionError::ConfigError(format!("Unknown extension: {}", name))
|
||||
|
||||
@@ -1064,6 +1064,7 @@ config_value!(GEMINI3_THINKING_LEVEL, String);
|
||||
config_value!(CLAUDE_THINKING_TYPE, String);
|
||||
config_value!(CLAUDE_THINKING_EFFORT, String);
|
||||
config_value!(CLAUDE_THINKING_BUDGET, i32);
|
||||
config_value!(GOOSE_DEFAULT_EXTENSION_TIMEOUT, u64);
|
||||
|
||||
fn find_workspace_or_exe_root() -> Option<PathBuf> {
|
||||
let exe = std::env::current_exe().ok()?;
|
||||
|
||||
Reference in New Issue
Block a user