From 1d6f7d50850e30bc41ea85bb055dbd0af0655a29 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Fri, 20 Mar 2026 23:57:19 -0700 Subject: [PATCH] fix: persist startup-loaded MCP clients in ExtensionManager (#1509) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: persist startup-loaded MCP clients in ExtensionManager MCP servers loaded at startup had their tools registered in the ToolRegistry but the client references were dropped. This caused the ExtensionManager to report them as disconnected and broke reconnection/session management. Collect startup MCP clients from the JoinSet and inject them into the ExtensionManager via a new inject_mcp_client() method. Also fix missing extension_manager field in fire_webhook EngineContext. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 (1M context) * fix: address PR review — pub(crate) visibility and JoinError diagnostics - Narrow inject_mcp_client to pub(crate) and guard against empty names - Distinguish panic vs cancellation in MCP task JoinError logging [skip-regression-check] Co-Authored-By: Claude Opus 4.6 (1M context) * merge: sync with staging, fix duplicate extension_manager field [skip-regression-check] Co-Authored-By: Claude Opus 4.6 (1M context) * fix: validate extension name in inject_mcp_client Add validate_extension_name() check to reject path traversal characters in MCP client names, consistent with other entry points. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/app.rs | 38 ++++++++++++++++++++++++++++++++++---- src/extensions/manager.rs | 25 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/src/app.rs b/src/app.rs index bca0f110..28e7ada5 100644 --- a/src/app.rs +++ b/src/app.rs @@ -536,7 +536,7 @@ impl AppBuilder { server_name, e ); - return; + return None; } }; @@ -553,6 +553,10 @@ impl AppBuilder { tool_count, server_name ); + return Some(( + server_name, + Arc::new(client), + )); } Err(e) => { tracing::warn!( @@ -583,14 +587,27 @@ impl AppBuilder { } } } + None }); } + let mut startup_clients = Vec::new(); while let Some(result) = join_set.join_next().await { - if let Err(e) = result { - tracing::warn!("MCP server loading task panicked: {}", e); + match result { + Ok(Some(client_pair)) => { + startup_clients.push(client_pair); + } + Ok(None) => {} + Err(e) => { + if e.is_panic() { + tracing::error!("MCP server loading task panicked: {}", e); + } else { + tracing::warn!("MCP server loading task failed: {}", e); + } + } } } + return startup_clients; } Err(e) => { if matches!( @@ -608,10 +625,12 @@ impl AppBuilder { } } } + Vec::new() } }; - let (dev_loaded_tool_names, _) = tokio::join!(wasm_tools_future, mcp_servers_future); + let (dev_loaded_tool_names, startup_mcp_clients) = + tokio::join!(wasm_tools_future, mcp_servers_future); // Load registry catalog entries for extension discovery let mut catalog_entries = match crate::registry::RegistryCatalog::load_or_embedded() { @@ -673,6 +692,17 @@ impl AppBuilder { )); tools.register_extension_tools(Arc::clone(&manager)); tracing::debug!("Extension manager initialized with in-chat discovery tools"); + + if !startup_mcp_clients.is_empty() { + tracing::info!( + count = startup_mcp_clients.len(), + "Injecting startup MCP clients into extension manager" + ); + for (name, client) in startup_mcp_clients { + manager.inject_mcp_client(name, client).await; + } + } + Some(manager) }; diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index f06def20..b8af4c68 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -937,6 +937,31 @@ impl ExtensionManager { &self.secrets } + /// Inject a pre-created MCP client (from startup loading) into the manager. + /// + /// Startup-loaded MCP clients register their tools in `ToolRegistry` but are + /// otherwise dropped. This method stores the client so that `list()` reports + /// accurate "connected" status and reconnection/session management works. + pub(crate) async fn inject_mcp_client( + &self, + name: String, + client: Arc, + ) { + if name.is_empty() { + tracing::warn!("inject_mcp_client called with empty name; ignoring"); + return; + } + if let Err(e) = Self::validate_extension_name(&name) { + tracing::warn!( + error = %e, + name = %name, + "inject_mcp_client called with invalid name; ignoring" + ); + return; + } + self.mcp_clients.write().await.insert(name, client); + } + /// Register channel names that were loaded at startup. /// Called after WASM channels are loaded so `list()` reports accurate active status. pub async fn set_active_channels(&self, names: Vec) {