use crate::{
config::{McpServerConfig, McpServersSettings, validate_mcp_server_name},
mcp::tool_schema::{provider_tool_definition, validate_tool_name},
mcp::{CallToolResult, InitializeResult, McpClient, McpError, McpResult, Tool},
};
use serde_json::Value;
use std::{collections::HashMap, path::Path};
pub(crate) const MCP_QUALIFIED_NAME_MAX_BYTES: usize = 64;
pub(crate) struct McpManager {
servers: Vec<McpServerHandle>,
routes: HashMap<String, (usize, String)>,
statuses: HashMap<String, McpServerStatus>,
}
pub(crate) struct McpServerHandle {
pub(crate) server_name: String,
pub(crate) client: McpClient,
pub(crate) tools: Vec<Tool>,
}
#[derive(Debug, Clone)]
pub(crate) enum McpServerStatus {
Connected {
#[allow(dead_code)]
server_info: crate::mcp::protocol::Implementation,
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(),
};
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) {
Ok((client, init, 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 {
server_name: server_name.clone(),
client,
tools,
});
manager.routes.extend(routes);
manager.statuses.insert(
server_name.clone(),
McpServerStatus::Connected {
server_info: init.server_info,
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
}
fn connect_server(
server_name: &str,
config: &McpServerConfig,
mc_home: Option<&Path>,
) -> Result<(McpClient, InitializeResult, Vec<Tool>), (&'static str, McpError)> {
let connect_phase = match config {
McpServerConfig::Stdio(_) => "spawn",
McpServerConfig::Http(_) => "connect",
};
let client = McpClient::connect_named(Some(server_name), config, mc_home)
.map_err(|error| (connect_phase, error))?;
let init = client.initialize().map_err(|error| ("initialize", error))?;
let tools = client.list_tools().map_err(|error| ("list_tools", error))?;
Ok((client, init, tools))
}
fn validate_routes(
&self,
server_index: usize,
server_name: &str,
tools: &[Tool],
) -> McpResult<HashMap<String, (usize, String)>> {
let mut routes = HashMap::new();
for tool in tools {
validate_tool_name(&tool.name)?;
let qualified = qualified_name(server_name, &tool.name);
validate_qualified_name(&qualified)?;
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 list_tool_definitions(&self) -> Vec<(String, Tool)> {
let mut definitions = Vec::new();
for handle in &self.servers {
for tool in &handle.tools {
definitions.push((
qualified_name(&handle.server_name, &tool.name),
tool.clone(),
));
}
}
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()
.filter_map(|(qualified, tool)| {
let server_name = qualified.strip_prefix("mcp__")?.split("__").next()?;
provider_tool_definition(server_name, &tool).ok()
})
.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> {
let (server_index, raw_tool_name) = self.resolve(qualified_name).ok_or_else(|| {
McpError::Config(format!("unknown MCP tool route '{qualified_name}'"))
})?;
self.servers[server_index]
.client
.call_tool(&raw_tool_name, arguments)
}
pub(crate) fn call_tool_cancellable(
&self,
qualified_name: &str,
arguments: Option<Value>,
cancellation: &crate::agent::cancellation::AgentCancellation,
) -> McpResult<CallToolResult> {
let (server_index, raw_tool_name) = self.resolve(qualified_name).ok_or_else(|| {
McpError::Config(format!("unknown MCP tool route '{qualified_name}'"))
})?;
self.servers[server_index].client.call_tool_cancellable(
&raw_tool_name,
arguments,
cancellation,
)
}
pub(crate) fn statuses(&self) -> &HashMap<String, McpServerStatus> {
&self.statuses
}
pub(crate) fn shutdown(&mut self) {
for server in &mut self.servers {
server.client.shutdown();
}
}
}
impl Drop for McpManager {
fn drop(&mut self) {
self.shutdown();
}
}
pub(crate) fn qualified_name(server: &str, tool: &str) -> String {
format!("mcp__{server}__{tool}")
}
fn validate_qualified_name(name: &str) -> McpResult<()> {
if name.len() > MCP_QUALIFIED_NAME_MAX_BYTES {
return Err(McpError::Config(format!(
"MCP qualified tool name exceeds {MCP_QUALIFIED_NAME_MAX_BYTES} bytes"
)));
}
Ok(())
}
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::*;
#[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(),
};
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),
})
}
#[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();
}
}