600 lines
20 KiB
Rust
600 lines
20 KiB
Rust
//! 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 tokio::sync::RwLock;
|
|
|
|
use rmcp::{
|
|
model::{CallToolRequestParams, CallToolResult, ServerInfo, Tool},
|
|
RoleClient, ServiceExt,
|
|
service::RunningService,
|
|
transport::TokioChildProcess,
|
|
transport::streamable_http_client::{StreamableHttpClientTransport, StreamableHttpClientTransportConfig},
|
|
};
|
|
use http::{HeaderName, HeaderValue};
|
|
use tokio::process::Command;
|
|
|
|
use crate::mcp::config::{McpServerConfig, McpTransportConfig};
|
|
use std::env;
|
|
|
|
/// 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()
|
|
}
|
|
|
|
/// Type alias for the MCP client service
|
|
pub type McpClient = RunningService<RoleClient, ()>;
|
|
|
|
/// 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<ServerInfo>,
|
|
/// Available tools
|
|
pub tools: Vec<Tool>,
|
|
}
|
|
|
|
/// 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<HashMap<String, Arc<McpClient>>>,
|
|
/// Server information cache keyed by server key
|
|
server_info: RwLock<HashMap<String, McpServerInfo>>,
|
|
/// 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<HashMap<String, String>>,
|
|
}
|
|
|
|
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<McpServerInfo> {
|
|
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 } => {
|
|
self.connect_stdio(&command, &args, &env).await?
|
|
}
|
|
McpTransportConfig::Http { url, headers } => {
|
|
self.connect_http(&url, &headers).await?
|
|
}
|
|
};
|
|
|
|
// Get server info (returns Option<&ServerInfo>)
|
|
let info = client.peer_info().cloned();
|
|
|
|
// 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,
|
|
command: &str,
|
|
args: &[String],
|
|
env: &HashMap<String, String>,
|
|
) -> anyhow::Result<McpClient> {
|
|
let mut cmd = Command::new(command);
|
|
cmd.args(args);
|
|
|
|
// Set environment variables
|
|
for (key, value) in env {
|
|
cmd.env(key, value);
|
|
}
|
|
|
|
let transport = TokioChildProcess::new(cmd)?;
|
|
|
|
// Use default client handler (empty tuple)
|
|
let client = ().serve(transport).await?;
|
|
|
|
// 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<String, String>,
|
|
) -> anyhow::Result<McpClient> {
|
|
// Resolve env placeholders in headers
|
|
let resolved_headers: HashMap<String, String> = 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<HeaderName, HeaderValue> = 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<Arc<McpClient>> {
|
|
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<McpServerInfo> {
|
|
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<String> {
|
|
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<String>,
|
|
tool_name: impl Into<String>,
|
|
args: serde_json::Value,
|
|
) -> anyhow::Result<CallToolResult> {
|
|
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<String>) -> 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<String, crate::mcp::McpServerConfig>,
|
|
) -> 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<String>,
|
|
}
|
|
|
|
/// 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<McpServerStatus>,
|
|
}
|
|
|
|
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<Arc<McpClientManager>>,
|
|
/// Connection task handle (for background initialization)
|
|
connection_task: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
|
|
}
|
|
|
|
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<Arc<McpClientManager>> {
|
|
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) -> 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).await?;
|
|
}
|
|
Ok(())
|
|
}
|
|
} |