3e97ed903c
在 McpToolWrapper 适配层用 tokio::time::timeout 包裹 call_tool，防止外部 MCP server 挂起导致 agent loop 无限阻塞。超时时间通过 config.mcp_tool_timeout_secs 配置（默认 300 秒，0=不超时），前端 McpTab 设置页提供输入框。
868 lines
29 KiB
Rust
868 lines
29 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 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<PathBuf> {
|
|
#[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<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,
|
|
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<Arc<ServerInfo>> 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<String, String>,
|
|
cwd: &Option<PathBuf>,
|
|
) -> anyhow::Result<McpClient> {
|
|
// 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(|| "<inherit>".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<Mutex<Vec<String>>> = 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<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,
|
|
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()
|
|
);
|
|
}
|
|
}
|
|
}
|