//! MCP Client Manager - manages connections to MCP servers //! //! This module provides a decoupled MCP integration that: //! - Doesn't block gateway startup //! - Is completely optional (disabled by default) //! - Connects to MCP servers asynchronously //! - Dynamically registers MCP tools via the Tool trait adapter use std::collections::HashMap; use std::sync::Arc; use parking_lot::Mutex; use tokio::sync::RwLock; use http::{HeaderName, HeaderValue}; use rmcp::{ RoleClient, ServiceExt, model::{CallToolRequestParams, CallToolResult, ServerInfo, Tool}, service::RunningService, transport::TokioChildProcess, transport::streamable_http_client::{ StreamableHttpClientTransport, StreamableHttpClientTransportConfig, }, }; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::Command; use crate::mcp::config::{McpServerConfig, McpTransportConfig}; use std::env; use std::path::PathBuf; use std::process::Stdio; /// Resolve ${ENV_VAR} placeholders in a value string fn resolve_env_placeholders_in_value(value: &str) -> String { let re = regex::Regex::new(r"\$\{([A-Z_][A-Z0-9_]*)\}").expect("invalid regex"); re.replace_all(value, |caps: ®ex::Captures<'_>| { let var_name = &caps[1]; env::var(var_name).unwrap_or_else(|_| caps[0].to_string()) }) .to_string() } /// Resolve a command name to an executable path. /// /// On Windows, if the command has no file extension, searches PATH for /// `.exe`, `.cmd`, `.bat` variants (Rust's `Command::new` only finds `.exe`). /// On non-Windows platforms, returns `None` (defers to OS default behavior). fn resolve_command_path(command: &str) -> Option { #[cfg(windows)] { let path = PathBuf::from(command); // If already an absolute path with extension, use directly if path.is_absolute() { if path.is_file() { return Some(path); } // Absolute path but doesn't exist — try appending extensions if no ext if path.extension().is_none() { for ext in &["exe", "cmd", "bat"] { let candidate = path.with_extension(ext); if candidate.is_file() { return Some(candidate); } } } return None; } // If it has a path separator (relative path), don't search PATH if path .parent() .map(|p| !p.as_os_str().is_empty()) .unwrap_or(false) { return None; } // Bare command name — search PATH with extensions let extensions: &[&str] = if path.extension().is_some() { &[""] // already has extension, try as-is } else { &[".exe", ".cmd", ".bat"] }; let path_env = std::env::var_os("PATH")?; for dir in std::env::split_paths(&path_env) { for ext in extensions { let candidate = dir.join(format!("{}{}", command, ext)); if candidate.is_file() { return Some(candidate); } } } None } #[cfg(not(windows))] { // On non-Windows, defer to OS default behavior let _ = command; None } } /// Type alias for the MCP client service pub type McpClient = RunningService; /// Information about a connected MCP server #[derive(Debug, Clone)] pub struct McpServerInfo { /// Server name (effective name from config) pub name: String, /// Server key (the key in mcpServers map) pub key: String, /// Server information from MCP protocol pub info: Option, /// Available tools pub tools: Vec, } /// Manager for MCP client connections /// /// This manager handles: /// - Connecting to MCP servers (stdio and HTTP transports) /// - Discovering available tools /// - Calling tools on connected servers /// - Connection lifecycle management pub struct McpClientManager { /// Connected clients keyed by server key clients: RwLock>>, /// Server information cache keyed by server key server_info: RwLock>, /// Count of active stdio (child process) connections stdio_client_count: std::sync::atomic::AtomicUsize, /// Connection errors per server key (last error message) connection_errors: RwLock>, } impl McpClientManager { /// Create a new manager (no connections yet) pub fn new() -> Self { Self { clients: RwLock::new(HashMap::new()), server_info: RwLock::new(HashMap::new()), stdio_client_count: std::sync::atomic::AtomicUsize::new(0), connection_errors: RwLock::new(HashMap::new()), } } /// Connect to all configured servers (async, non-blocking) /// /// 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; } 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" ); // Record error for status reporting self.connection_errors .write() .await .insert(key.clone(), e.to_string()); failed += 1; } else { // Clear any previous error on successful connection self.connection_errors.write().await.remove(&key); } } 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(()) } /// Connect to a single MCP server pub async fn connect_server( &self, key: &str, config: &McpServerConfig, ) -> anyhow::Result { let effective_name = config.effective_name(key); tracing::info!(key = %key, name = %effective_name, transport_type = %config.transport_type, "Connecting to MCP server"); let transport = config.transport().map_err(|e| anyhow::anyhow!("{}", e))?; let client = match transport { McpTransportConfig::Stdio { command, args, env, cwd, } => self.connect_stdio(key, &command, &args, &env, &cwd).await?, McpTransportConfig::Http { url, headers } => self.connect_http(&url, &headers).await?, }; // Get server info (returns Option> in rmcp 1.8+) let info = client.peer_info().map(|arc| arc.as_ref().clone()); // List available tools let tools = client.list_all_tools().await?; let server_info = McpServerInfo { key: key.to_string(), name: effective_name, info, tools, }; // Store the client and info { let mut clients = self.clients.write().await; clients.insert(key.to_string(), Arc::new(client)); } { let mut info_map = self.server_info.write().await; info_map.insert(key.to_string(), server_info.clone()); } Ok(server_info) } /// Connect via stdio transport (spawn child process) async fn connect_stdio( &self, server_key: &str, command: &str, args: &[String], env: &HashMap, cwd: &Option, ) -> anyhow::Result { // Resolve command path (Windows extension handling) let resolved_command = resolve_command_path(command); let effective_command: PathBuf = match &resolved_command { Some(p) => p.clone(), None => PathBuf::from(command), }; // Pre-flight check on Windows: if command is bare name and not resolved, give friendly error #[cfg(windows)] if resolved_command.is_none() { let path = std::path::Path::new(command); let is_absolute = path.is_absolute(); let has_separator = path .parent() .map(|p| !p.as_os_str().is_empty()) .unwrap_or(false); if !is_absolute && !has_separator && path.extension().is_none() { // Bare name not found on Windows let path_env = std::env::var("PATH").unwrap_or_default(); return Err(anyhow::anyhow!( "Command '{}' not found in PATH (tried .exe, .cmd, .bat). \ Current PATH: {}. \ Suggestion: use the full absolute path to the executable, \ or ensure the tool is installed and its directory is in PATH.", command, path_env )); } } let env_keys: Vec<&String> = env.keys().collect(); let cwd_display = cwd .as_ref() .map(|p| p.display().to_string()) .unwrap_or_else(|| "".to_string()); tracing::info!( server_key = %server_key, command = %effective_command.display(), args = ?args, env_keys = ?env_keys, // only keys, NOT values (avoid leaking secrets) cwd = %cwd_display, "Spawning MCP stdio child process" ); let mut cmd = Command::new(&effective_command); cmd.args(args); // Set environment variables for (key, value) in env { cmd.env(key, value); } // Set working directory if specified if let Some(cwd_path) = cwd { cmd.current_dir(cwd_path); } // Shared buffer to collect stderr lines for error reporting let stderr_lines: Arc>> = Arc::new(Mutex::new(Vec::new())); let stderr_lines_for_task = stderr_lines.clone(); // Use builder to capture stderr (default is inherit) let (transport, stderr_opt) = TokioChildProcess::builder(cmd) .stderr(Stdio::piped()) .spawn() .map_err(|e| { anyhow::anyhow!( "Failed to spawn MCP stdio child process '{}': {}", effective_command.display(), e ) })?; // Spawn a background task to read stderr and log it if let Some(child_stderr) = stderr_opt { let server_key_owned = server_key.to_string(); tokio::spawn(async move { let reader = BufReader::new(child_stderr); let mut lines = reader.lines(); while let Ok(Some(line)) = lines.next_line().await { // Escalate real warnings/errors; demote normal diagnostics to debug let lower = line.to_lowercase(); let is_warning = lower.contains("error") || lower.contains("warn") || lower.contains("panic") || lower.contains("fatal"); if is_warning { tracing::warn!( server_key = %server_key_owned, stderr = %line, "MCP child process stderr" ); } else { tracing::debug!( server_key = %server_key_owned, stderr = %line, "MCP child process stderr" ); } // Also collect into the shared buffer (cap at 50 lines) { let mut buf = stderr_lines_for_task.lock(); if buf.len() < 50 { buf.push(line); } } } }); } // Use default client handler (empty tuple) let client = ().serve(transport).await.map_err(|e| { // Include stderr summary in error if available let stderr_summary = { let buf = stderr_lines.lock(); if buf.is_empty() { String::new() } else { format!("\nstderr:\n {}", buf.join("\n ")) } }; anyhow::anyhow!( "Failed to establish MCP stdio connection '{}': {}{}", effective_command.display(), e, stderr_summary ) })?; // Track that we have a stdio (child process) connection self.stdio_client_count .fetch_add(1, std::sync::atomic::Ordering::SeqCst); Ok(client) } /// Connect via HTTP transport (Streamable HTTP) async fn connect_http( &self, url: &str, headers: &HashMap, ) -> anyhow::Result { // Resolve env placeholders in headers let resolved_headers: HashMap = headers .iter() .map(|(key, value)| { // Resolve ${ENV_VAR} placeholders let resolved = if value.contains("${") { resolve_env_placeholders_in_value(value) } else { value.clone() }; (key.clone(), resolved) }) .collect(); // Build custom headers let custom_headers: HashMap = resolved_headers .iter() .filter_map(|(key, value)| { // Try to parse header name and value HeaderName::try_from(key.clone()).ok().and_then(|name| { HeaderValue::try_from(value.clone()) .ok() .map(|val| (name, val)) }) }) .collect(); // Create transport config with custom headers let config = StreamableHttpClientTransportConfig::with_uri(url).custom_headers(custom_headers); // Create transport using reqwest client (default) let transport = StreamableHttpClientTransport::with_client(reqwest::Client::default(), config); // Connect let client = ().serve(transport).await?; Ok(client) } /// Get a client by server key pub async fn get_client(&self, key: &str) -> Option> { let clients = self.clients.read().await; clients.get(key).cloned() } /// Get server info by key pub async fn get_server_info(&self, key: &str) -> Option { let info_map = self.server_info.read().await; info_map.get(key).cloned() } /// Get all connected server keys pub async fn connected_servers(&self) -> Vec { let clients = self.clients.read().await; clients.keys().cloned().collect() } /// Get all tools from all connected servers /// Returns (server_key, tool) pairs for tool registration pub async fn all_tools(&self) -> Vec<(String, Tool)> { let info_map = self.server_info.read().await; info_map .values() .flat_map(|info| { info.tools .iter() .map(|tool| (info.key.clone(), tool.clone())) }) .collect() } /// Call a tool on a specific server by key pub async fn call_tool( &self, server_key: impl Into, tool_name: impl Into, args: serde_json::Value, ) -> anyhow::Result { let server_key = server_key.into(); let tool_name = tool_name.into(); let client = self .get_client(&server_key) .await .ok_or_else(|| anyhow::anyhow!("MCP server '{}' not connected", server_key))?; // Convert Value to JsonObject if it's an object let arguments = if args.is_object() { args.as_object().unwrap().clone() } else { // If not an object, use empty object serde_json::Map::new() }; // Create params with owned String (converted to Cow<'static, str>) let params = CallToolRequestParams::new(tool_name).with_arguments(arguments); let result = client.call_tool(params).await?; Ok(result) } /// Disconnect from a server by key pub async fn disconnect(&self, key: impl Into) -> anyhow::Result<()> { let key = key.into(); let mut clients = self.clients.write().await; if clients.remove(&key).is_some() { tracing::info!(key = %key, "Disconnected MCP server"); } self.server_info.write().await.remove(&key); Ok(()) } /// Disconnect from all servers pub async fn disconnect_all(&self) -> anyhow::Result<()> { let mut clients = self.clients.write().await; for (key, _client) in clients.drain() { tracing::info!(key = %key, "Disconnected MCP server"); } self.server_info.write().await.clear(); Ok(()) } /// Shut down all MCP connections with proper cleanup /// /// Drops all client connections and waits for child processes to terminate /// if stdio transports were in use. This prevents race conditions during /// gateway restart where old MCP processes may still be running when /// new ones start. pub async fn shutdown_all(&self) -> anyhow::Result<()> { let stdio_count = self .stdio_client_count .load(std::sync::atomic::Ordering::SeqCst); // Drop all clients (triggers cancellation + graceful shutdown in rmcp) self.disconnect_all().await?; // If stdio connections were active, wait for child processes to be killed. // rmcp's RunningService::drop() triggers async cancellation with: // - 2 second graceful drain period // - 3 second process kill timeout // Total: ~5 seconds. We add 1 second buffer. if stdio_count > 0 { tracing::info!( stdio_count, "Waiting for MCP child processes to terminate (up to 6s)..." ); tokio::time::sleep(std::time::Duration::from_secs(6)).await; tracing::info!("MCP child process cleanup wait complete"); self.stdio_client_count .store(0, std::sync::atomic::Ordering::SeqCst); } Ok(()) } /// Check if any servers are connected pub async fn has_connections(&self) -> bool { !self.clients.read().await.is_empty() } /// Get current MCP connection status for all configured servers pub async fn get_status( &self, configured_servers: &HashMap, ) -> McpStatusResponse { let info_map = self.server_info.read().await; let errors = self.connection_errors.read().await; let clients = self.clients.read().await; let mut total_servers = 0usize; let mut connected_servers = 0usize; let mut failed_servers = 0usize; let mut total_tools = 0usize; let mut servers = Vec::new(); for (key, config) in configured_servers { total_servers += 1; let name = config.effective_name(key); let transport_type = config.transport_type.clone(); let is_active = config.is_active; let connected = clients.contains_key(key); let tool_count = info_map.get(key).map(|info| info.tools.len()).unwrap_or(0); let error = errors.get(key).cloned(); if connected { connected_servers += 1; } else if is_active && error.is_some() { failed_servers += 1; } total_tools += tool_count; servers.push(McpServerStatus { key: key.clone(), name, transport_type, is_active, connected, tool_count, error, }); } McpStatusResponse { enabled: !configured_servers.is_empty(), total_servers, connected_servers, failed_servers, total_tools, servers, } } } /// Status of a single MCP server connection #[derive(Debug, Clone, serde::Serialize)] pub struct McpServerStatus { pub key: String, pub name: String, pub transport_type: String, pub is_active: bool, pub connected: bool, pub tool_count: usize, pub error: Option, } /// Overall MCP status response #[derive(Debug, Clone, Default, serde::Serialize)] pub struct McpStatusResponse { pub enabled: bool, pub total_servers: usize, pub connected_servers: usize, pub failed_servers: usize, pub total_tools: usize, pub servers: Vec, } impl Default for McpClientManager { fn default() -> Self { Self::new() } } /// MCP Initializer - handles asynchronous MCP initialization /// /// This struct provides a decoupled way to initialize MCP: /// - Doesn't block gateway startup /// - Can be initialized in a background task /// - Tools are registered after connection is established pub struct McpInitializer { /// The MCP client manager (None if MCP is disabled) manager: Option>, /// Connection task handle (for background initialization) connection_task: Option>>, } impl McpInitializer { /// Create a disabled initializer (MCP not configured) pub fn disabled() -> Self { Self { manager: None, connection_task: None, } } /// Create an initializer with MCP configuration /// /// 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(); // Spawn background connection task let manager_clone = manager.clone(); let connection_task = tokio::spawn(async move { tracing::info!("Starting MCP connection task..."); manager_clone.connect_all(servers).await }); Self { manager: Some(manager), connection_task: Some(connection_task), } } /// Get the manager (if MCP is enabled) pub fn manager(&self) -> Option> { self.manager.clone() } /// Check if MCP is enabled pub fn is_enabled(&self) -> bool { self.manager.is_some() } /// Wait for connections to complete (optional) /// /// This can be called if you want to ensure MCP servers are connected /// before proceeding, but it's not required. pub async fn wait_for_connections(&mut self) -> anyhow::Result<()> { if let Some(task) = self.connection_task.take() { // Handle JoinError and inner Result task.await??; } Ok(()) } /// Register MCP tools to the tool registry /// /// This should be called after the gateway is ready to accept tools. /// Waits for connections to complete before registering tools. pub async fn register_tools( &mut self, registry: &mut crate::tools::ToolRegistry, timeout_secs: u64, ) -> anyhow::Result<()> { if let Some(manager) = self.manager.clone() { // Wait for connections to complete first self.wait_for_connections().await?; tracing::info!("Registering MCP tools after connections completed"); crate::mcp::register_mcp_tools(manager, registry, timeout_secs).await?; } Ok(()) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_resolve_command_path_returns_none_for_nonexistent_bare_name() { // A bare name that definitely doesn't exist let result = resolve_command_path("this_binary_definitely_does_not_exist_xyz123"); // On all platforms, this should return None (either because it's not found, // or because non-Windows always returns None) #[cfg(windows)] assert!( result.is_none(), "Expected None for nonexistent command on Windows" ); #[cfg(not(windows))] assert!( result.is_none(), "Expected None on non-Windows (always returns None)" ); } #[test] fn test_resolve_command_path_absolute_path_existing() { // Use an absolute path to a known existing executable #[cfg(windows)] let cmd = "C:\\Windows\\System32\\cmd.exe"; #[cfg(not(windows))] let cmd = "/bin/ls"; let result = resolve_command_path(cmd); #[cfg(windows)] assert!(result.is_some(), "Expected to find {} on Windows", cmd); #[cfg(not(windows))] assert!( result.is_none(), "Expected None on non-Windows for absolute path" ); } #[test] fn test_resolve_command_path_absolute_path_nonexistent() { let result = resolve_command_path("/definitely/not/a/real/path/binary123"); assert!(result.is_none()); } #[test] #[cfg(windows)] fn test_resolve_command_path_finds_known_windows_binary() { // cmd.exe should always be in C:\Windows\System32 which is in PATH let result = resolve_command_path("cmd"); assert!( result.is_some(), "Expected to find cmd.exe via PATH on Windows" ); if let Some(p) = result { assert!( p.to_string_lossy().to_lowercase().ends_with("cmd.exe"), "Expected resolved path to end with cmd.exe, got: {}", p.display() ); } } }