From 4fabbe47e9da9f23cfd1de34043ed356bbbb2cd4 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Tue, 30 Jun 2026 10:46:33 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BAMCP=E9=9B=86=E6=88=90?= =?UTF-8?q?=EF=BC=8C=E6=94=AF=E6=8C=81=E9=85=8D=E7=BD=AE=E8=B7=AF=E5=BE=84?= =?UTF-8?q?=E5=92=8C=E8=BF=9E=E6=8E=A5=E9=87=8D=E8=AF=95=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/config/mod.rs | 6 ++- src/gateway/mod.rs | 13 ++++- src/gateway/runtime.rs | 10 ++-- src/gateway/session.rs | 2 +- src/mcp/client.rs | 112 ++++++++++++++++++++++++++++++++++------- 5 files changed, 120 insertions(+), 23 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index de44528..2159d6c 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -873,7 +873,9 @@ impl Config { } pub fn load_default() -> Result> { - let path = get_default_config_path(); + let path = std::env::var("CONFIG_PATH") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| get_default_config_path()); Self::load_from(&path) } @@ -907,6 +909,8 @@ impl Config { mcp_servers = config.mcp_servers.len(), "MCP servers loaded from config" ); + } else { + tracing::info!("No mcpServers found in config"); } Ok(config) diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index afb574b..80ce060 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -62,6 +62,7 @@ pub struct GatewayState { pub task_repository: Arc, pub cancel_manager: CancelManager, pub restart_tx: watch::Sender, + pub mcp_manager: Option>, } impl GatewayState { @@ -86,7 +87,7 @@ impl GatewayState { mcp_servers: config.mcp_servers.clone(), }; - let (session_manager, task_repository) = build_session_manager_with_sender( + let (session_manager, task_repository, mcp_manager) = build_session_manager_with_sender( agent_prompt_reinject_every, show_tool_results, config.time.timezone.clone(), @@ -113,6 +114,7 @@ impl GatewayState { task_repository, cancel_manager, restart_tx, + mcp_manager, }) } @@ -241,6 +243,7 @@ pub async fn run( let (result_tx, result_rx) = tokio::sync::oneshot::channel::(); let channel_manager = state.channel_manager.clone(); let cancel_manager = state.cancel_manager.clone(); + let mcp_manager = state.mcp_manager.clone(); // Spawn ctrl_c / restart handler tokio::spawn(async move { @@ -249,6 +252,10 @@ pub async fn run( tracing::info!("Shutdown signal received"); cancel_manager.cancel_all().await; let _ = scheduler_shutdown_tx.send(true); + if let Some(ref mgr) = mcp_manager { + tracing::info!("Disconnecting MCP servers before shutdown"); + let _ = mgr.disconnect_all().await; + } let _ = channel_manager.stop_all().await; let _ = result_tx.send(false); let _ = shutdown_tx.send(()); @@ -258,6 +265,10 @@ pub async fn run( tracing::info!("Restart signal received"); cancel_manager.cancel_all().await; let _ = scheduler_shutdown_tx.send(true); + if let Some(ref mgr) = mcp_manager { + tracing::info!("Disconnecting MCP servers before restart"); + let _ = mgr.disconnect_all().await; + } let _ = channel_manager.stop_all().await; let _ = result_tx.send(true); let _ = shutdown_tx.send(()); diff --git a/src/gateway/runtime.rs b/src/gateway/runtime.rs index a7ca2be..cbd3638 100644 --- a/src/gateway/runtime.rs +++ b/src/gateway/runtime.rs @@ -10,6 +10,7 @@ use crate::bus::MessageBus; use crate::config::{LLMProviderConfig, MemoryMaintenanceConfig, SubagentsConfig, TaskConfig}; use crate::gateway::tool_registry_factory::ToolRegistryFactory; use crate::mcp::McpInitializer; +use crate::mcp::client::McpClientManager; use crate::skills::SkillRuntime; use crate::storage::{ ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository, @@ -49,7 +50,7 @@ pub(crate) fn build_session_manager( session_ttl_hours: Option, mcp_config: crate::mcp::McpConfig, bus: Option>, -) -> Result<(SessionManager, Arc), AgentError> { +) -> Result<(SessionManager, Arc, Option>), AgentError> { build_session_manager_with_sender( agent_prompt_reinject_every, show_tool_results, @@ -84,7 +85,7 @@ pub(crate) fn build_session_manager_with_sender( session_ttl_hours: Option, mcp_config: crate::mcp::McpConfig, bus: Option>, -) -> Result<(SessionManager, Arc), AgentError> { +) -> Result<(SessionManager, Arc, Option>), AgentError> { let store = Arc::new( SessionStore::new() .map_err(|err| AgentError::Other(format!("session store init error: {}", err)))?, @@ -265,6 +266,9 @@ pub(crate) fn build_session_manager_with_sender( let memory_maintenance = MemoryMaintenanceCoordinator::new(store.clone(), provider_configs.clone()); + // Extract MCP manager for lifecycle management (e.g., disconnect on restart) + let mcp_manager = mcp_initializer.manager(); + Ok((SessionManager::from_services(SessionManagerServices { tools: tools as Arc, skills, @@ -276,5 +280,5 @@ pub(crate) fn build_session_manager_with_sender( scheduled_tasks, memory_maintenance, task_repository: task_repository.clone(), - }), task_repository)) + }), task_repository, mcp_manager)) } diff --git a/src/gateway/session.rs b/src/gateway/session.rs index 77beaab..8cdb2a6 100644 --- a/src/gateway/session.rs +++ b/src/gateway/session.rs @@ -680,7 +680,7 @@ impl SessionManager { mcp_config, None, ) - .map(|(session_manager, _)| session_manager) + .map(|(session_manager, _, _)| session_manager) } pub fn tools(&self) -> Arc { diff --git a/src/mcp/client.rs b/src/mcp/client.rs index 9cb4aeb..6d27db1 100644 --- a/src/mcp/client.rs +++ b/src/mcp/client.rs @@ -77,33 +77,90 @@ impl McpClientManager { /// This method is designed to be called asynchronously without /// blocking the main gateway startup flow. /// Takes a list of (key, config) pairs from the mcpServers map. + /// + /// Each server connection is retried up to 3 times with exponential backoff + /// (1s, 2s, 4s) to handle transient failures during gateway restart. pub async fn connect_all(&self, servers: Vec<(String, McpServerConfig)>) -> anyhow::Result<()> { + let total = servers.iter().filter(|(_, c)| c.is_active).count(); + if total == 0 { + tracing::info!("No active MCP servers to connect"); + return Ok(()); + } + + tracing::info!(server_count = total, "Connecting to MCP servers"); + + let mut connected = 0usize; + let mut failed = 0usize; + const MAX_RETRIES: u32 = 3; + for (key, config) in servers { if !config.is_active { tracing::info!(key = %key, "Skipping inactive MCP server"); continue; } - // Each server connection is independent - match self.connect_server(&key, &config).await { - Ok(info) => { - tracing::info!( - key = %key, - name = %info.name, - tools_count = info.tools.len(), - "Connected to MCP server" - ); - } - Err(e) => { - // Log error but continue with other servers - tracing::error!( - key = %key, - error = %e, - "Failed to connect to MCP server" - ); + let mut last_error = None; + for attempt in 1..=MAX_RETRIES { + match self.connect_server(&key, &config).await { + Ok(info) => { + tracing::info!( + key = %key, + name = %info.name, + tools_count = info.tools.len(), + attempt, + "Connected to MCP server" + ); + connected += 1; + last_error = None; + break; + } + Err(e) => { + last_error = Some(e); + if attempt < MAX_RETRIES { + let delay_secs = 1u64 << (attempt - 1); // 1s, 2s, 4s + tracing::warn!( + key = %key, + attempt, + max_retries = MAX_RETRIES, + retry_delay_secs = delay_secs, + error = %last_error.as_ref().unwrap(), + "MCP connection failed, retrying" + ); + tokio::time::sleep(std::time::Duration::from_secs(delay_secs)).await; + } + } } } + + if let Some(e) = last_error { + tracing::error!( + key = %key, + error = %e, + attempts = MAX_RETRIES, + "Failed to connect to MCP server after all retries" + ); + failed += 1; + } } + + tracing::info!( + total, + connected, + failed, + "MCP connection summary: {}/{} connected, {} failed", + connected, + total, + failed + ); + + // Only return error if all connections failed + if connected == 0 && failed > 0 { + return Err(anyhow::anyhow!( + "All {} MCP server connection(s) failed", + failed + )); + } + Ok(()) } @@ -343,10 +400,31 @@ impl McpInitializer { /// This spawns a background task to connect to MCP servers, /// allowing the gateway to start immediately. pub fn with_config(config: crate::mcp::McpConfig) -> Self { + let server_count = config.mcp_servers.len(); + let active_count = config.active_servers().len(); + if !config.has_active_servers() { + if server_count > 0 { + tracing::info!( + server_count, + active_count, + "MCP disabled: {} server(s) configured but none active", + server_count + ); + } else { + tracing::info!("MCP disabled: no mcpServers configured in config"); + } return Self::disabled(); } + tracing::info!( + server_count, + active_count, + "MCP enabled: {} active server(s) out of {} configured", + active_count, + server_count + ); + let manager = Arc::new(McpClientManager::new()); let servers = config.active_servers();