mod connection;
pub mod descriptors;
pub mod lease;
use std::collections::HashMap;
use std::sync::Arc;
pub use lease::{McpLeaseCapability, McpRuntimeManager, McpSessionEndGuard};
#[derive(Debug, Clone, serde::Deserialize)]
pub struct McpServerConfig {
pub command: String,
#[serde(default)]
pub args: Vec<String>,
#[serde(default)]
pub env: HashMap<String, String>,
#[serde(default)]
pub shared: bool,
}
#[derive(Debug, Clone, serde::Deserialize)]
pub struct McpConfig {
#[serde(rename = "mcpServers", default)]
pub mcp_servers: HashMap<String, McpServerConfig>,
}
#[derive(Debug, Clone)]
pub(crate) struct McpToolDef {
pub(crate) name: String,
pub(crate) description: String,
pub(crate) input_schema: serde_json::Value,
}
pub const MCP_CONFIG_MAX_BYTES: u64 = 1024 * 1024;
pub fn load_mcp_config() -> Option<McpConfig> {
let path = crate::config::resolve_read_path("mcp.json");
let content = match descriptors::read_bounded_regular_file(&path, MCP_CONFIG_MAX_BYTES) {
Ok(content) => content,
Err(descriptors::DescriptorCacheError::NotFound) => return None,
Err(err) => {
let category = match err {
descriptors::DescriptorCacheError::NotFound => "not_found",
descriptors::DescriptorCacheError::NotRegularFile => "not_regular_file",
descriptors::DescriptorCacheError::Oversize { .. } => "oversize",
descriptors::DescriptorCacheError::Io(_) => "io",
descriptors::DescriptorCacheError::Parse(_) => "parse",
descriptors::DescriptorCacheError::Version(_) => "version",
};
tracing::warn!(
file = "mcp.json",
category,
"Refusing unsafe MCP config read"
);
return None;
}
};
let config = match serde_json::from_str::<McpConfig>(&content) {
Ok(config) => config,
Err(e) => {
tracing::warn!(
file = "mcp.json",
class = ?e.classify(),
line = e.line(),
column = e.column(),
"Failed to parse MCP config"
);
return None;
}
};
if let Err(reason) = validate_mcp_config(&config) {
tracing::warn!(
file = "mcp.json",
reason,
"Refusing structurally unsafe MCP config"
);
return None;
}
Some(config)
}
pub const MCP_CONFIG_MAX_SERVERS: usize = 32;
pub const MCP_CONFIG_MAX_ARGS: usize = 64;
pub const MCP_CONFIG_MAX_ENV_ENTRIES: usize = 64;
pub const MCP_CONFIG_MAX_SERVER_NAME_BYTES: usize = 64;
pub const MCP_CONFIG_MAX_COMMAND_BYTES: usize = 512;
pub const MCP_CONFIG_MAX_ARG_BYTES: usize = 4096;
pub const MCP_CONFIG_MAX_ENV_KEY_BYTES: usize = 128;
pub const MCP_CONFIG_MAX_ENV_VALUE_BYTES: usize = 4096;
fn safe_str(value: &str, max: usize) -> bool {
!value.is_empty() && value.len() <= max && !value.chars().any(char::is_control)
}
pub fn validate_mcp_config(config: &McpConfig) -> Result<(), &'static str> {
if config.mcp_servers.len() > MCP_CONFIG_MAX_SERVERS {
return Err("too_many_servers");
}
for (name, server) in &config.mcp_servers {
if !safe_str(name, MCP_CONFIG_MAX_SERVER_NAME_BYTES) {
return Err("invalid_server_name");
}
if !safe_str(&server.command, MCP_CONFIG_MAX_COMMAND_BYTES) {
return Err("invalid_command");
}
if server.args.len() > MCP_CONFIG_MAX_ARGS {
return Err("too_many_args");
}
for arg in &server.args {
if arg.len() > MCP_CONFIG_MAX_ARG_BYTES || arg.chars().any(char::is_control) {
return Err("invalid_arg");
}
}
if server.env.len() > MCP_CONFIG_MAX_ENV_ENTRIES {
return Err("too_many_env_entries");
}
for (key, value) in &server.env {
if !safe_str(key, MCP_CONFIG_MAX_ENV_KEY_BYTES) {
return Err("invalid_env_key");
}
if value.len() > MCP_CONFIG_MAX_ENV_VALUE_BYTES || value.chars().any(char::is_control) {
return Err("invalid_env_value");
}
}
}
Ok(())
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum McpDiscovery {
#[default]
Inline,
Background,
}
const SEED_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
pub const DISCOVERY_CONCURRENCY: usize = 8;
#[doc(hidden)]
#[derive(Clone, Copy, Debug)]
pub struct McpDiscoveryOpts {
pub timeout: std::time::Duration,
pub now_unix: Option<u64>,
}
impl Default for McpDiscoveryOpts {
fn default() -> Self {
Self {
timeout: SEED_TIMEOUT,
now_unix: None,
}
}
}
impl McpDiscoveryOpts {
fn now_unix(&self) -> u64 {
self.now_unix.unwrap_or_else(|| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_secs())
})
}
}
type SharedRegistry = Arc<tokio::sync::RwLock<crate::ToolRegistry>>;
pub async fn setup_lazy_mcp(registry: &SharedRegistry) -> usize {
setup_lazy_mcp_with(registry, McpDiscoveryOpts::default()).await
}
#[doc(hidden)]
pub async fn setup_lazy_mcp_with(registry: &SharedRegistry, opts: McpDiscoveryOpts) -> usize {
let Some(config) = load_nonempty_config() else {
return 0;
};
let mut cache = load_cache_or_empty();
if seed_missing_descriptors(&config, &cache, &opts).await > 0 {
cache = load_cache_or_empty();
}
register_dormant(registry, &config, &cache, "boot").await
}
pub async fn setup_cached_mcp(registry: &SharedRegistry) -> (usize, Option<PendingMcpDiscovery>) {
let Some(config) = load_nonempty_config() else {
return (0, None);
};
let cache = load_cache_or_empty();
let mut missing = HashMap::new();
for (name, server) in &config.mcp_servers {
let fingerprint = descriptors::server_config_fingerprint(server);
let cached = cache
.servers
.get(name)
.is_some_and(|entry| entry.fingerprint == fingerprint);
if !cached {
missing.insert(name.clone(), server.clone());
}
}
let count = register_dormant(registry, &config, &cache, "cached").await;
let pending = (!missing.is_empty() && descriptors::cache_writeback_enabled()).then_some(
PendingMcpDiscovery {
config: McpConfig {
mcp_servers: missing,
},
},
);
(count, pending)
}
#[derive(Debug)]
pub struct PendingMcpDiscovery {
config: McpConfig,
}
impl PendingMcpDiscovery {
pub fn server_names(&self) -> Vec<String> {
let mut names: Vec<String> = self.config.mcp_servers.keys().cloned().collect();
names.sort();
names
}
pub async fn run(self, registry: &SharedRegistry) -> usize {
self.run_with(registry, McpDiscoveryOpts::default()).await
}
#[doc(hidden)]
pub async fn run_with(self, registry: &SharedRegistry, opts: McpDiscoveryOpts) -> usize {
let mut cache = load_cache_or_empty();
if seed_missing_descriptors(&self.config, &cache, &opts).await > 0 {
cache = load_cache_or_empty();
}
register_dormant(registry, &self.config, &cache, "background").await
}
}
fn load_nonempty_config() -> Option<McpConfig> {
load_mcp_config().filter(|config| !config.mcp_servers.is_empty())
}
fn load_cache_or_empty() -> descriptors::McpDescriptorCache {
match descriptors::load_default_cache() {
Ok(cache) => cache,
Err(descriptors::DescriptorCacheError::NotFound) => descriptors::McpDescriptorCache::empty(),
Err(err) => {
tracing::warn!(error = %err,
"Refusing unsafe MCP descriptor cache; MCP capabilities stay undiscoverable this run");
descriptors::McpDescriptorCache::empty()
}
}
}
async fn register_dormant(
registry: &SharedRegistry,
config: &McpConfig,
cache: &descriptors::McpDescriptorCache,
phase: &'static str,
) -> usize {
let server_count = config.mcp_servers.len();
let dormant = descriptors::dormant_tools_for_config(config, cache);
if dormant.is_empty() {
tracing::info!(
servers = server_count,
phase,
"MCP: no valid cached descriptors; MCP tools are not discoverable this run"
);
return 0;
}
match registry.write().await.try_register_batch(dormant) {
Ok(count) => {
tracing::info!(
tools = count,
servers = server_count,
phase,
"MCP: descriptor-backed tools registered (no process started)"
);
count
}
Err(e) => {
tracing::warn!(error = %e, phase,
"MCP descriptor batch rejected; failing closed with zero MCP capabilities");
0
}
}
}
fn format_unix(secs: u64) -> String {
i64::try_from(secs)
.ok()
.and_then(|s| chrono::DateTime::<chrono::Utc>::from_timestamp(s, 0))
.map_or_else(|| secs.to_string(), |dt| dt.to_rfc3339())
}
async fn seed_missing_descriptors(
config: &McpConfig,
cache: &descriptors::McpDescriptorCache,
opts: &McpDiscoveryOpts,
) -> usize {
use futures::StreamExt;
if !descriptors::cache_writeback_enabled() {
return 0;
}
let now = opts.now_unix();
let mut names: Vec<&String> = config.mcp_servers.keys().collect();
names.sort();
let mut pending = Vec::new();
for name in names {
let server = &config.mcp_servers[name];
let fingerprint = descriptors::server_config_fingerprint(server);
if cache
.servers
.get(name)
.is_some_and(|entry| entry.fingerprint == fingerprint)
{
continue;
}
if let Some(failure) = cache.discovery_failures.get(name) {
if failure.suppresses(&fingerprint, now) {
let when = format_unix(failure.retry_after_unix());
tracing::info!(
server = %name,
attempts = failure.attempts,
"MCP discovery for '{}' failed recently; retrying after {}",
name,
when
);
continue;
}
}
pending.push((name.clone(), server.clone(), fingerprint));
}
if pending.is_empty() {
return 0;
}
let opts = *opts;
let outcomes: Vec<bool> = futures::stream::iter(
pending
.into_iter()
.map(|(name, server, fingerprint)| discover_one(name, server, fingerprint, opts)),
)
.buffer_unordered(DISCOVERY_CONCURRENCY)
.collect()
.await;
outcomes.into_iter().filter(|recorded| *recorded).count()
}
async fn discover_one(
name: String,
server: McpServerConfig,
fingerprint: String,
opts: McpDiscoveryOpts,
) -> bool {
let listing = tokio::time::timeout(opts.timeout, async {
let mut conn = connection::McpConnection::start(&server).await?;
let result = conn.list_tools().await;
conn.start_kill();
result
})
.await;
let path = descriptors::default_cache_path();
let defs = match listing {
Ok(Ok(defs)) => defs,
Ok(Err(err)) => {
tracing::warn!(server = %name, error = %err, "MCP discovery failed; its tools stay undiscoverable until a retry");
record_failure(path, name, fingerprint, opts.now_unix()).await;
return false;
}
Err(_) => {
tracing::warn!(server = %name, "MCP discovery timed out; its tools stay undiscoverable until a retry");
record_failure(path, name, fingerprint, opts.now_unix()).await;
return false;
}
};
let tools: Vec<descriptors::CachedToolDescriptor> = defs
.into_iter()
.map(|def| descriptors::CachedToolDescriptor {
name: def.name,
description: def.description,
input_schema: def.input_schema,
})
.collect();
let server_name = name.clone();
let written = tokio::task::spawn_blocking(move || {
descriptors::record_server_listing(&path, &server_name, &fingerprint, &tools)
})
.await;
match written {
Ok(Ok(count)) => {
tracing::info!(server = %name, tools = count, "MCP server discovered and cached");
true
}
Ok(Err(err)) => {
tracing::warn!(server = %name, error = %err, "MCP discovery could not be cached");
false
}
Err(err) => {
tracing::warn!(server = %name, error = %err, "MCP discovery cache write panicked");
false
}
}
}
async fn record_failure(path: std::path::PathBuf, name: String, fingerprint: String, now: u64) {
let server_name = name.clone();
let written = tokio::task::spawn_blocking(move || {
descriptors::record_discovery_failure(&path, &server_name, &fingerprint, now)
})
.await;
match written {
Ok(Ok(failure)) => tracing::info!(
server = %name,
attempts = failure.attempts,
retry_after = %format_unix(failure.retry_after_unix()),
"MCP discovery failure recorded"
),
Ok(Err(err)) => {
tracing::warn!(server = %name, error = %err, "MCP discovery failure could not be cached")
}
Err(err) => {
tracing::warn!(server = %name, error = %err, "MCP discovery failure write panicked")
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_mcp_config_deserialize() {
let json_str = r#"{"mcpServers": {"test": {"command": "echo", "args": ["hi"]}}}"#;
let config: McpConfig = serde_json::from_str(json_str).unwrap();
assert_eq!(config.mcp_servers.len(), 1);
assert!(config.mcp_servers.contains_key("test"));
let server = &config.mcp_servers["test"];
assert_eq!(server.command, "echo");
assert_eq!(server.args, vec!["hi"]);
}
#[test]
fn test_mcp_config_empty_servers() {
let json_str = r#"{"mcpServers": {}}"#;
let config: McpConfig = serde_json::from_str(json_str).unwrap();
assert_eq!(config.mcp_servers.len(), 0);
assert!(config.mcp_servers.is_empty());
}
#[test]
fn test_mcp_server_config_defaults() {
let json_str = r#"{"command": "echo"}"#;
let server_config: McpServerConfig = serde_json::from_str(json_str).unwrap();
assert_eq!(server_config.command, "echo");
assert_eq!(server_config.args, Vec::<String>::new());
assert_eq!(server_config.env, HashMap::new());
}
#[test]
fn test_mcp_config_deserialize_from_value() {
let json_value = json!({
"mcpServers": {
"test": {
"command": "echo",
"args": ["hi"]
}
}
});
let config: McpConfig = serde_json::from_value(json_value).unwrap();
assert_eq!(config.mcp_servers.len(), 1);
assert!(config.mcp_servers.contains_key("test"));
let server = &config.mcp_servers["test"];
assert_eq!(server.command, "echo");
assert_eq!(server.args, vec!["hi"]);
}
fn base_server() -> McpServerConfig {
McpServerConfig {
command: "/bin/true".to_string(),
args: vec!["--flag".to_string()],
env: HashMap::from([("KEY".to_string(), "value".to_string())]),
shared: false,
}
}
fn config_of(servers: Vec<(&str, McpServerConfig)>) -> McpConfig {
McpConfig {
mcp_servers: servers
.into_iter()
.map(|(k, v)| (k.to_string(), v))
.collect(),
}
}
#[test]
fn validate_mcp_config_accepts_sane_configs() {
assert!(validate_mcp_config(&config_of(vec![("srv", base_server())])).is_ok());
assert!(validate_mcp_config(&config_of(vec![])).is_ok());
}
#[test]
fn validate_mcp_config_fails_whole_config_closed_on_structural_violations() {
let many: Vec<(String, McpServerConfig)> = (0..=MCP_CONFIG_MAX_SERVERS)
.map(|i| (format!("s{i}"), base_server()))
.collect();
let config = McpConfig {
mcp_servers: many.into_iter().collect(),
};
assert_eq!(validate_mcp_config(&config), Err("too_many_servers"));
let mut c = base_server();
assert_eq!(
validate_mcp_config(&config_of(vec![("evil\u{7}", c.clone())])),
Err("invalid_server_name")
);
let long_name = "n".repeat(MCP_CONFIG_MAX_SERVER_NAME_BYTES + 1);
assert_eq!(
validate_mcp_config(&config_of(vec![(long_name.as_str(), c.clone())])),
Err("invalid_server_name")
);
c.command = String::new();
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_command")
);
c.command = "x".repeat(MCP_CONFIG_MAX_COMMAND_BYTES + 1);
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_command")
);
c.command = "bad\u{0}cmd".to_string();
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_command")
);
let mut c = base_server();
c.args = vec!["a".to_string(); MCP_CONFIG_MAX_ARGS + 1];
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("too_many_args")
);
c.args = vec!["ok".to_string(), "bad\u{1b}arg".to_string()];
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_arg")
);
let mut c = base_server();
c.env = (0..=MCP_CONFIG_MAX_ENV_ENTRIES)
.map(|i| (format!("K{i}"), "v".to_string()))
.collect();
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("too_many_env_entries")
);
c.env = HashMap::from([(String::new(), "v".to_string())]);
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_env_key")
);
c.env = HashMap::from([(
"K".to_string(),
"v".repeat(MCP_CONFIG_MAX_ENV_VALUE_BYTES + 1),
)]);
assert_eq!(
validate_mcp_config(&config_of(vec![("srv", c.clone())])),
Err("invalid_env_value")
);
}
#[test]
fn test_load_mcp_config_returns_some_or_none() {
let result = load_mcp_config();
match result {
Some(_config) => {
}
None => {
}
}
}
}