From 776d92cbdab27377090079fb02b663cbfaf65ede Mon Sep 17 00:00:00 2001 From: Robert Mcgregor <38837341+exitcode0@users.noreply.github.com> Date: Fri, 3 Apr 2026 00:53:08 +1100 Subject: [PATCH] perf(acp): parallelize extension loading in ACP server (#8098) Signed-off-by: Douwe Osinga Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Douwe Osinga --- crates/goose-acp/src/server.rs | 60 ++++++++++++++++++++--------- crates/goose/src/agents/agent.rs | 65 ++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 18 deletions(-) diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index 281e678f..0b776364 100644 --- a/crates/goose-acp/src/server.rs +++ b/crates/goose-acp/src/server.rs @@ -479,21 +479,30 @@ impl GooseAcpAgent { let skip_developer = acp_developer.is_some(); let sid_str = session_id.map(|s| s.0.to_string()); - for ext in extensions { - if skip_developer && ext.name() == "developer" { - continue; - } - let name = ext.name().to_string(); - match agent - .extension_manager - .add_extension(ext, None, None, sid_str.as_deref()) - .await - { - Ok(_) => info!(extension = %name, "extension loaded"), - Err(e) => warn!(extension = %name, error = %e, "extension load failed"), - } + if skip_developer { + extensions.retain(|ext| ext.name() != "developer"); } + let ext_manager = &agent.extension_manager; + let extension_futures = extensions + .into_iter() + .map(|ext| { + let ext_manager = Arc::clone(ext_manager); + let sid = sid_str.clone(); + async move { + let name = ext.name().to_string(); + match ext_manager + .add_extension(ext, None, None, sid.as_deref()) + .await + { + Ok(_) => info!(extension = %name, "extension loaded"), + Err(e) => warn!(extension = %name, error = %e, "extension load failed"), + } + } + }) + .collect::>(); + futures::future::join_all(extension_futures).await; + if let Some((client, config)) = acp_developer { let info = client.get_info().cloned(); agent @@ -932,10 +941,11 @@ impl GooseAcpAgent { } async fn add_mcp_extensions( - agent: &Agent, + agent: &Arc, mcp_servers: Vec, session_id: &str, ) -> Result<(), sacp::Error> { + let mut configs = Vec::with_capacity(mcp_servers.len()); for mcp_server in mcp_servers { let config = match mcp_server_to_extension_config(mcp_server) { Ok(c) => c, @@ -943,10 +953,24 @@ impl GooseAcpAgent { return Err(sacp::Error::invalid_params().data(msg)); } }; - let name = config.name().to_string(); - if let Err(e) = agent.add_extension(config, session_id).await { - return Err(sacp::Error::internal_error() - .data(format!("Failed to add MCP server '{}': {}", name, e))); + configs.push(config); + } + + if configs.is_empty() { + return Ok(()); + } + + let results = agent + .add_extensions_bulk(configs, session_id) + .await + .map_err(|e| sacp::Error::internal_error().data(e.to_string()))?; + for result in &results { + if !result.success { + let error_msg = result.error.as_deref().unwrap_or("unknown error"); + return Err(sacp::Error::internal_error().data(format!( + "Failed to add MCP server '{}': {}", + result.name, error_msg + ))); } } Ok(()) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index ac8c6295..92de31bc 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -761,6 +761,71 @@ impl Agent { Ok(()) } + /// Load multiple extensions in parallel, persisting state once at the end. + /// + /// Unlike `add_extension`, this avoids per-extension persistence and acquires + /// the container lock once upfront to prevent serialisation of the parallel futures. + pub async fn add_extensions_bulk( + self: &Arc, + extensions: Vec, + session_id: &str, + ) -> anyhow::Result> { + let working_dir = match self + .config + .session_manager + .get_session(session_id, false) + .await + { + Ok(session) => Some(session.working_dir), + Err(e) => { + warn!("Failed to get session for bulk load: {}", e); + None + } + }; + let container = self.container.lock().await.clone(); + + let extension_futures = extensions + .into_iter() + .map(|config| { + let ext_manager = Arc::clone(&self.extension_manager); + let working_dir = working_dir.clone(); + let container = container.clone(); + let sid = session_id.to_string(); + + async move { + let name = config.name().to_string(); + match ext_manager + .add_extension(config, working_dir, container.as_ref(), Some(&sid)) + .await + { + Ok(_) => ExtensionLoadResult { + name, + success: true, + error: None, + }, + Err(e) => { + let error_msg = e.to_string(); + warn!("Failed to load extension {}: {}", name, error_msg); + ExtensionLoadResult { + name, + success: false, + error: Some(error_msg), + } + } + } + } + }) + .collect::>(); + + let results = futures::future::join_all(extension_futures).await; + + if results.iter().any(|r| r.success) { + self.persist_extension_state(session_id).await?; + } + + Ok(results) + } + async fn add_extension_inner( &self, extension: ExtensionConfig,