#[cfg(test)]
pub(crate) use crate::mcp::names::MCP_QUALIFIED_NAME_MAX_BYTES;
use crate::{
cancellation::AgentCancellation,
config::{McpServerConfig, McpServersSettings, validate_mcp_server_name},
mcp::tool_schema::provider_tool_definition,
mcp::{CallToolResult, McpClient, McpError, McpResult, QualifiedMcpToolName, Tool},
};
use serde_json::Value;
use std::{
collections::HashMap,
path::Path,
sync::{Arc, Mutex},
};
#[derive(Debug)]
pub(crate) struct McpManager {
servers: Vec<McpServerHandle>,
routes: HashMap<QualifiedMcpToolName, (usize, String)>,
statuses: HashMap<String, McpServerStatus>,
shutting_down: bool,
}
#[derive(Debug)]
pub(crate) struct McpServerHandle {
client: Arc<Mutex<McpClient>>,
pub(crate) tools: Vec<Tool>,
}
#[derive(Debug, Clone)]
pub(crate) struct ResolvedMcpToolCall {
client: Arc<Mutex<McpClient>>,
raw_tool_name: String,
}
impl ResolvedMcpToolCall {
#[cfg(test)]
pub(crate) fn call_tool(&self, arguments: Option<Value>) -> McpResult<CallToolResult> {
self.client
.lock()
.map_err(|_| McpError::Transport("MCP server client lock poisoned".to_string()))?
.call_tool(&self.raw_tool_name, arguments)
}
pub(crate) fn call_tool_cancellable(
&self,
arguments: Option<Value>,
cancellation: &AgentCancellation,
) -> McpResult<CallToolResult> {
self.client
.lock()
.map_err(|_| McpError::Transport("MCP server client lock poisoned".to_string()))?
.call_tool_cancellable(&self.raw_tool_name, arguments, cancellation)
}
}
#[derive(Debug, Clone)]
pub(crate) enum McpServerStatus {
Connected { tool_count: usize },
Failed { error: String, phase: String },
}
impl McpManager {
#[cfg(test)]
pub(crate) fn from_settings(mcp_servers: &McpServersSettings) -> Self {
Self::from_settings_with_paths(mcp_servers, None)
}
pub(crate) fn from_settings_with_paths(
mcp_servers: &McpServersSettings,
mc_home: Option<&Path>,
) -> Self {
let mut manager = Self {
servers: Vec::new(),
routes: HashMap::new(),
statuses: HashMap::new(),
shutting_down: false,
};
for (server_name, config) in mcp_servers {
if !config.enabled() {
continue;
}
if let Err(error) = validate_mcp_server_name(server_name) {
manager.record_failure(server_name, "config", error);
continue;
}
match Self::connect_server(server_name, config, mc_home, None) {
Ok((client, tools)) => {
let server_index = manager.servers.len();
let route_result = manager.validate_routes(server_index, server_name, &tools);
match route_result {
Ok(routes) => {
manager.servers.push(McpServerHandle {
client: Arc::new(Mutex::new(client)),
tools,
});
manager.routes.extend(routes);
manager.statuses.insert(
server_name.clone(),
McpServerStatus::Connected {
tool_count: manager.servers[server_index].tools.len(),
},
);
}
Err(error) => manager.record_failure(server_name, "route", error),
}
}
Err((phase, error)) => manager.record_failure(server_name, phase, error),
}
}
manager
}
#[cfg(test)]
pub(crate) fn from_settings_strict(
mcp_servers: &McpServersSettings,
mc_home: Option<&Path>,
) -> McpResult<Self> {
Self::from_settings_strict_inner(mcp_servers, mc_home, None)
}
pub(crate) fn from_settings_strict_cancellable(
mcp_servers: &McpServersSettings,
mc_home: Option<&Path>,
cancellation: &AgentCancellation,
) -> McpResult<Self> {
Self::from_settings_strict_inner(mcp_servers, mc_home, Some(cancellation))
}
fn from_settings_strict_inner(
mcp_servers: &McpServersSettings,
mc_home: Option<&Path>,
cancellation: Option<&AgentCancellation>,
) -> McpResult<Self> {
let mut manager = Self {
servers: Vec::new(),
routes: HashMap::new(),
statuses: HashMap::new(),
shutting_down: false,
};
for (server_name, config) in mcp_servers {
if let Some(cancellation) = cancellation
&& let Err(error) = cancellation.check()
{
manager.shutdown();
return Err(McpError::transport(error));
}
if !config.enabled() {
continue;
}
if let Err(error) = validate_mcp_server_name(server_name) {
manager.shutdown();
return Err(McpError::Config(format!(
"MCP server '{server_name}' failed during config validation: {error}"
)));
}
let (client, tools) =
match Self::connect_server(server_name, config, mc_home, cancellation) {
Ok(connected) => connected,
Err((phase, error)) => {
manager.shutdown();
return Err(McpError::Config(format!(
"MCP server '{server_name}' failed during {phase}: {error}"
)));
}
};
let server_index = manager.servers.len();
let routes = match manager.validate_routes(server_index, server_name, &tools) {
Ok(routes) => routes,
Err(error) => {
manager.shutdown();
return Err(McpError::Config(format!(
"MCP server '{server_name}' failed during route validation: {error}"
)));
}
};
manager.servers.push(McpServerHandle {
client: Arc::new(Mutex::new(client)),
tools,
});
manager.routes.extend(routes);
manager.statuses.insert(
server_name.clone(),
McpServerStatus::Connected {
tool_count: manager.servers[server_index].tools.len(),
},
);
}
Ok(manager)
}
fn connect_server(
server_name: &str,
config: &McpServerConfig,
mc_home: Option<&Path>,
cancellation: Option<&AgentCancellation>,
) -> Result<(McpClient, Vec<Tool>), (&'static str, McpError)> {
let connect_phase = match config {
McpServerConfig::Stdio(_) => "spawn",
McpServerConfig::Http(_) => "connect",
};
if let Some(cancellation) = cancellation {
cancellation
.check()
.map_err(|error| (connect_phase, McpError::transport(error)))?;
}
let client = McpClient::connect_named(Some(server_name), config, mc_home)
.map_err(|error| (connect_phase, error))?;
if let Some(cancellation) = cancellation {
cancellation
.check()
.map_err(|error| (connect_phase, McpError::transport(error)))?;
}
match cancellation {
Some(cancellation) => client.initialize_cancellable(cancellation),
None => client.initialize(),
}
.map_err(|error| ("initialize", error))?;
if let Some(cancellation) = cancellation {
cancellation
.check()
.map_err(|error| ("initialize", McpError::transport(error)))?;
}
let tools = match cancellation {
Some(cancellation) => client.list_tools_cancellable(cancellation),
None => client.list_tools(),
}
.map_err(|error| ("list_tools", error))?;
if let Some(cancellation) = cancellation {
cancellation
.check()
.map_err(|error| ("list_tools", McpError::transport(error)))?;
}
Ok((client, tools))
}
fn validate_routes(
&self,
server_index: usize,
server_name: &str,
tools: &[Tool],
) -> McpResult<HashMap<QualifiedMcpToolName, (usize, String)>> {
let mut routes = HashMap::new();
for tool in tools {
let qualified = QualifiedMcpToolName::new(server_name, &tool.name)?;
if self.routes.contains_key(&qualified) || routes.contains_key(&qualified) {
return Err(McpError::Config(format!(
"duplicate MCP tool route '{}'",
qualified
)));
}
routes.insert(qualified, (server_index, tool.name.clone()));
}
Ok(routes)
}
fn record_failure(
&mut self,
server_name: &str,
phase: impl Into<String>,
error: impl std::fmt::Display,
) {
self.statuses.insert(
server_name.to_string(),
McpServerStatus::Failed {
error: bounded_error(error.to_string()),
phase: phase.into(),
},
);
}
pub(crate) fn resolve(&self, qualified_name: &str) -> Option<(usize, String)> {
self.routes.get(qualified_name).cloned()
}
pub(crate) fn resolve_tool_call(&self, qualified_name: &str) -> McpResult<ResolvedMcpToolCall> {
if self.shutting_down {
return Err(McpError::Transport(
"MCP manager is shutting down".to_string(),
));
}
let (server_index, raw_tool_name) = self.resolve(qualified_name).ok_or_else(|| {
McpError::Config(format!("unknown MCP tool route '{qualified_name}'"))
})?;
let client = self
.servers
.get(server_index)
.ok_or_else(|| McpError::Config(format!("unknown MCP tool route '{qualified_name}'")))?
.client
.clone();
Ok(ResolvedMcpToolCall {
client,
raw_tool_name,
})
}
pub(crate) fn list_tool_definitions(&self) -> Vec<(QualifiedMcpToolName, Tool)> {
let mut definitions = self
.routes
.iter()
.filter_map(|(qualified, (server_index, _raw_tool_name))| {
let tool = self
.servers
.get(*server_index)?
.tools
.iter()
.find(|tool| tool.name == qualified.tool())?;
Some((qualified.clone(), tool.clone()))
})
.collect::<Vec<_>>();
definitions.sort_by(|left, right| left.0.cmp(&right.0));
definitions
}
pub(crate) fn provider_tool_definitions(&self) -> Vec<Value> {
let mut definitions = self
.list_tool_definitions()
.into_iter()
.map(|(qualified, tool)| provider_tool_definition(&qualified, &tool))
.collect::<Vec<_>>();
definitions.sort_by(|left, right| {
left.get("name")
.and_then(Value::as_str)
.cmp(&right.get("name").and_then(Value::as_str))
});
definitions
}
#[cfg(test)]
pub(crate) fn call_tool(
&self,
qualified_name: &str,
arguments: Option<Value>,
) -> McpResult<CallToolResult> {
self.resolve_tool_call(qualified_name)?.call_tool(arguments)
}
pub(crate) fn statuses(&self) -> &HashMap<String, McpServerStatus> {
&self.statuses
}
pub(crate) fn shutdown(&mut self) {
self.shutting_down = true;
for server in &self.servers {
if let Ok(mut client) = server.client.lock() {
client.shutdown();
}
}
}
}
impl Drop for McpManager {
fn drop(&mut self) {
self.shutdown();
}
}
#[cfg(test)]
fn qualified_name(server: &str, tool: &str) -> String {
QualifiedMcpToolName::new(server, tool).unwrap().to_string()
}
fn bounded_error(mut error: String) -> String {
const MAX_ERROR_BYTES: usize = 512;
if error.len() > MAX_ERROR_BYTES {
let truncate_at = error
.char_indices()
.map(|(index, _)| index)
.take_while(|index| *index <= MAX_ERROR_BYTES)
.last()
.unwrap_or(0);
error.truncate(truncate_at);
error.push_str("...");
}
error
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(unix)]
use std::{
path::Path,
process::Command,
thread,
time::{Duration, Instant},
};
#[test]
fn bounded_error_truncates_ascii() {
let error = bounded_error("x".repeat(600));
assert!(error.ends_with("..."), "{error}");
assert_eq!(error.len(), 515);
}
#[test]
fn bounded_error_truncates_on_utf8_boundary() {
let error = bounded_error(format!("{}é", "x".repeat(511)));
assert!(error.ends_with("..."), "{error}");
assert!(error.is_char_boundary(error.len()));
assert!(!error.contains('é'), "{error}");
}
#[test]
fn manager_failed_http_status_does_not_leak_url_credentials_or_headers() {
let mut headers = std::collections::BTreeMap::new();
headers.insert("X-Team".to_string(), "literal-secret-value".to_string());
let mut settings = crate::config::McpServersSettings::new();
settings.insert(
"remote".to_string(),
crate::config::McpServerConfig::Http(crate::config::McpHttpServerConfig {
url: "http://user:pass@127.0.0.1:9/mcp?token=url-secret#frag".to_string(),
headers,
oauth: None,
enabled: true,
timeout: Some(1),
}),
);
let manager = McpManager::from_settings(&settings);
let Some(McpServerStatus::Failed { error, .. }) = manager.statuses().get("remote") else {
panic!("expected failed status")
};
assert!(!error.contains("literal-secret-value"), "{error}");
assert!(!error.contains("user:pass"), "{error}");
assert!(!error.contains("url-secret"), "{error}");
assert!(!error.contains("token="), "{error}");
}
#[test]
fn qualified_names_are_prefixed() {
assert_eq!(qualified_name("mock", "echo"), "mcp__mock__echo");
}
#[test]
fn route_registration_rejects_long_names_without_stale_routes() {
let manager = McpManager {
servers: Vec::new(),
routes: HashMap::new(),
statuses: HashMap::new(),
shutting_down: false,
};
let good = Tool {
name: "echo".to_string(),
title: None,
description: None,
input_schema: serde_json::json!({"type":"object"}),
output_schema: None,
annotations: None,
};
let bad = Tool {
name: "x".repeat(80),
title: None,
description: None,
input_schema: serde_json::json!({"type":"object"}),
output_schema: None,
annotations: None,
};
let error = manager
.validate_routes(0, "server", &[good, bad])
.unwrap_err()
.to_string();
assert!(error.contains("exceeds"), "{error}");
assert!(manager.routes.is_empty());
assert!(manager.resolve("mcp__server__echo").is_none());
}
fn mock_config(mode: &str, timeout: u64) -> crate::config::McpServerConfig {
let mut env = std::collections::BTreeMap::new();
env.insert("MCP_MOCK_MODE".to_string(), mode.to_string());
crate::config::McpServerConfig::Stdio(crate::config::McpStdioServerConfig {
command: "sh".to_string(),
args: vec![format!(
"{}/tests/fixtures/mcp/mock_stdio_server.sh",
env!("CARGO_MANIFEST_DIR")
)],
env,
enabled: true,
timeout: Some(timeout),
})
}
#[cfg(unix)]
fn mock_config_with_test_files(
mode: &str,
timeout: u64,
pid_file: &Path,
request_file: Option<&Path>,
) -> crate::config::McpServerConfig {
let mut config = mock_config(mode, timeout);
let crate::config::McpServerConfig::Stdio(stdio) = &mut config else {
unreachable!();
};
stdio.env.insert(
"MCP_MOCK_TEST_PID_FILE".to_string(),
pid_file.to_string_lossy().into_owned(),
);
if let Some(request_file) = request_file {
stdio.env.insert(
"MCP_MOCK_TEST_REQUEST_FILE".to_string(),
request_file.to_string_lossy().into_owned(),
);
}
config
}
#[cfg(unix)]
fn wait_for_pid_file(path: &Path, timeout: Duration) -> u32 {
let deadline = Instant::now() + timeout;
loop {
if let Some(pid) = std::fs::read_to_string(path)
.ok()
.and_then(|contents| contents.trim().parse::<u32>().ok())
{
return pid;
}
assert!(
Instant::now() < deadline,
"timed out waiting for {}",
path.display()
);
thread::sleep(Duration::from_millis(10));
}
}
#[cfg(unix)]
fn pid_is_running(pid: u32) -> bool {
Command::new("/bin/kill")
.args(["-0", &pid.to_string()])
.status()
.map(|status| status.success())
.unwrap_or(false)
}
#[cfg(unix)]
fn assert_pid_stopped(pid: u32) {
let deadline = Instant::now() + Duration::from_secs(2);
while pid_is_running(pid) && Instant::now() < deadline {
thread::sleep(Duration::from_millis(20));
}
assert!(
!pid_is_running(pid),
"MCP fixture process {pid} is still running"
);
}
#[test]
fn strict_manager_connects_enabled_servers() {
let mut settings = crate::config::McpServersSettings::new();
settings.insert("good".to_string(), mock_config("normal", 2));
let manager = McpManager::from_settings_strict(&settings, None).unwrap();
assert!(matches!(
manager.statuses().get("good"),
Some(McpServerStatus::Connected { tool_count: 2, .. })
));
assert_eq!(manager.list_tool_definitions().len(), 2);
}
#[test]
fn strict_manager_fails_route_errors() {
let mut settings = crate::config::McpServersSettings::new();
settings.insert("bad".to_string(), mock_config("one_valid_one_long", 2));
let error = McpManager::from_settings_strict(&settings, None)
.unwrap_err()
.to_string();
assert!(error.contains("failed during route validation"), "{error}");
assert!(error.contains("exceeds"), "{error}");
}
#[test]
fn cancellable_strict_manager_aborts_hanging_stdio_initialization() {
let mut settings = crate::config::McpServersSettings::new();
settings.insert("hang".to_string(), mock_config("hang", 5));
let cancel_flag = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let cancellation =
crate::cancellation::AgentCancellation::new(std::sync::Arc::clone(&cancel_flag));
let canceler = {
let cancel_flag = std::sync::Arc::clone(&cancel_flag);
std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(150));
cancel_flag.store(true, std::sync::atomic::Ordering::SeqCst);
})
};
let started = std::time::Instant::now();
let error = McpManager::from_settings_strict_cancellable(&settings, None, &cancellation)
.unwrap_err()
.to_string();
canceler.join().unwrap();
assert!(error.contains("prompt canceled"), "{error}");
assert!(started.elapsed() < std::time::Duration::from_secs(2));
}
#[cfg(unix)]
#[test]
fn strict_manager_cancellation_stops_all_ordered_stdio_processes() {
let temp = tempfile::TempDir::new().unwrap();
let first_pid_file = temp.path().join("first.pid");
let second_pid_file = temp.path().join("second.pid");
let second_request_file = temp.path().join("second.request");
let mut settings = crate::config::McpServersSettings::new();
settings.insert(
"first".to_string(),
mock_config_with_test_files("normal", 3, &first_pid_file, None),
);
settings.insert(
"second".to_string(),
mock_config_with_test_files("hang", 3, &second_pid_file, Some(&second_request_file)),
);
let cancel_flag = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let cancellation =
crate::cancellation::AgentCancellation::new(std::sync::Arc::clone(&cancel_flag));
let canceler = {
let cancel_flag = std::sync::Arc::clone(&cancel_flag);
let second_request_file = second_request_file.clone();
std::thread::spawn(move || {
let deadline = Instant::now() + Duration::from_secs(2);
while !second_request_file.exists() && Instant::now() < deadline {
thread::sleep(Duration::from_millis(10));
}
thread::sleep(Duration::from_millis(50));
cancel_flag.store(true, std::sync::atomic::Ordering::SeqCst);
})
};
let error = McpManager::from_settings_strict_cancellable(&settings, None, &cancellation)
.unwrap_err()
.to_string();
canceler.join().unwrap();
assert!(
second_request_file.exists(),
"second stdio server never received initialize"
);
assert!(error.contains("second"), "{error}");
assert!(error.contains("prompt canceled"), "{error}");
let first_pid = wait_for_pid_file(&first_pid_file, Duration::from_secs(1));
let second_pid = wait_for_pid_file(&second_pid_file, Duration::from_secs(1));
assert_ne!(first_pid, second_pid);
assert_pid_stopped(first_pid);
assert_pid_stopped(second_pid);
}
#[test]
fn repeated_shutdown_rejects_new_resolved_calls() {
let mut settings = crate::config::McpServersSettings::new();
settings.insert("server".to_string(), mock_config("normal", 2));
let mut manager = McpManager::from_settings_strict(&settings, None).unwrap();
manager.shutdown();
manager.shutdown();
let error = manager
.resolve_tool_call("mcp__server__echo")
.unwrap_err()
.to_string();
assert!(error.contains("shutting down"), "{error}");
}
#[test]
fn manager_keeps_good_server_when_another_server_fails_routes() {
let mut settings = crate::config::McpServersSettings::new();
settings.insert("bad".to_string(), mock_config("one_valid_one_long", 2));
settings.insert("good".to_string(), mock_config("normal", 2));
let mut manager = McpManager::from_settings(&settings);
assert!(matches!(
manager.statuses().get("bad"),
Some(McpServerStatus::Failed { phase, .. }) if phase == "route"
));
assert!(matches!(
manager.statuses().get("good"),
Some(McpServerStatus::Connected { tool_count: 2, .. })
));
let result = manager
.call_tool("mcp__good__echo", Some(serde_json::json!({"text":"ok"})))
.unwrap();
assert!(matches!(
&result.content[0],
crate::mcp::ContentBlock::Text { text } if text == "ok"
));
assert!(manager.resolve("mcp__bad__echo").is_none());
manager.shutdown();
}
}