1mod connection;
3pub mod descriptors;
4pub mod lease;
5
6use std::collections::HashMap;
7use std::sync::Arc;
8
9pub use lease::{McpLeaseCapability, McpRuntimeManager, McpSessionEndGuard};
10
11#[derive(Debug, Clone, serde::Deserialize)]
13pub struct McpServerConfig {
14 pub command: String,
15 #[serde(default)]
16 pub args: Vec<String>,
17 #[serde(default)]
18 pub env: HashMap<String, String>,
19 #[serde(default)]
26 pub shared: bool,
27}
28
29#[derive(Debug, Clone, serde::Deserialize)]
31pub struct McpConfig {
32 #[serde(rename = "mcpServers", default)]
33 pub mcp_servers: HashMap<String, McpServerConfig>,
34}
35
36#[derive(Debug, Clone)]
38pub(crate) struct McpToolDef {
39 pub(crate) name: String,
40 pub(crate) description: String,
41 pub(crate) input_schema: serde_json::Value,
42}
43
44pub const MCP_CONFIG_MAX_BYTES: u64 = 1024 * 1024;
46
47pub fn load_mcp_config() -> Option<McpConfig> {
55 let path = crate::config::resolve_read_path("mcp.json");
56 let content = match descriptors::read_bounded_regular_file(&path, MCP_CONFIG_MAX_BYTES) {
57 Ok(content) => content,
58 Err(descriptors::DescriptorCacheError::NotFound) => return None,
59 Err(err) => {
60 let category = match err {
62 descriptors::DescriptorCacheError::NotFound => "not_found",
63 descriptors::DescriptorCacheError::NotRegularFile => "not_regular_file",
64 descriptors::DescriptorCacheError::Oversize { .. } => "oversize",
65 descriptors::DescriptorCacheError::Io(_) => "io",
66 descriptors::DescriptorCacheError::Parse(_) => "parse",
67 descriptors::DescriptorCacheError::Version(_) => "version",
68 };
69 tracing::warn!(
70 file = "mcp.json",
71 category,
72 "Refusing unsafe MCP config read"
73 );
74 return None;
75 }
76 };
77 let config = match serde_json::from_str::<McpConfig>(&content) {
78 Ok(config) => config,
79 Err(e) => {
80 tracing::warn!(
83 file = "mcp.json",
84 class = ?e.classify(),
85 line = e.line(),
86 column = e.column(),
87 "Failed to parse MCP config"
88 );
89 return None;
90 }
91 };
92 if let Err(reason) = validate_mcp_config(&config) {
93 tracing::warn!(
95 file = "mcp.json",
96 reason,
97 "Refusing structurally unsafe MCP config"
98 );
99 return None;
100 }
101 Some(config)
102}
103
104pub const MCP_CONFIG_MAX_SERVERS: usize = 32;
108pub const MCP_CONFIG_MAX_ARGS: usize = 64;
109pub const MCP_CONFIG_MAX_ENV_ENTRIES: usize = 64;
110pub const MCP_CONFIG_MAX_SERVER_NAME_BYTES: usize = 64;
111pub const MCP_CONFIG_MAX_COMMAND_BYTES: usize = 512;
112pub const MCP_CONFIG_MAX_ARG_BYTES: usize = 4096;
113pub const MCP_CONFIG_MAX_ENV_KEY_BYTES: usize = 128;
114pub const MCP_CONFIG_MAX_ENV_VALUE_BYTES: usize = 4096;
115
116fn safe_str(value: &str, max: usize) -> bool {
117 !value.is_empty() && value.len() <= max && !value.chars().any(char::is_control)
118}
119
120pub fn validate_mcp_config(config: &McpConfig) -> Result<(), &'static str> {
123 if config.mcp_servers.len() > MCP_CONFIG_MAX_SERVERS {
124 return Err("too_many_servers");
125 }
126 for (name, server) in &config.mcp_servers {
127 if !safe_str(name, MCP_CONFIG_MAX_SERVER_NAME_BYTES) {
128 return Err("invalid_server_name");
129 }
130 if !safe_str(&server.command, MCP_CONFIG_MAX_COMMAND_BYTES) {
131 return Err("invalid_command");
132 }
133 if server.args.len() > MCP_CONFIG_MAX_ARGS {
134 return Err("too_many_args");
135 }
136 for arg in &server.args {
137 if arg.len() > MCP_CONFIG_MAX_ARG_BYTES || arg.chars().any(char::is_control) {
138 return Err("invalid_arg");
139 }
140 }
141 if server.env.len() > MCP_CONFIG_MAX_ENV_ENTRIES {
142 return Err("too_many_env_entries");
143 }
144 for (key, value) in &server.env {
145 if !safe_str(key, MCP_CONFIG_MAX_ENV_KEY_BYTES) {
146 return Err("invalid_env_key");
147 }
148 if value.len() > MCP_CONFIG_MAX_ENV_VALUE_BYTES || value.chars().any(char::is_control) {
149 return Err("invalid_env_value");
150 }
151 }
152 }
153 Ok(())
154}
155
156#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
159pub enum McpDiscovery {
160 #[default]
164 Inline,
165 Background,
171}
172
173const SEED_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
176
177pub const DISCOVERY_CONCURRENCY: usize = 8;
181
182#[doc(hidden)]
185#[derive(Clone, Copy, Debug)]
186pub struct McpDiscoveryOpts {
187 pub timeout: std::time::Duration,
189 pub now_unix: Option<u64>,
192}
193
194impl Default for McpDiscoveryOpts {
195 fn default() -> Self {
196 Self {
197 timeout: SEED_TIMEOUT,
198 now_unix: None,
199 }
200 }
201}
202
203impl McpDiscoveryOpts {
204 fn now_unix(&self) -> u64 {
205 self.now_unix.unwrap_or_else(|| {
206 std::time::SystemTime::now()
207 .duration_since(std::time::UNIX_EPOCH)
208 .map_or(0, |d| d.as_secs())
209 })
210 }
211}
212
213type SharedRegistry = Arc<tokio::sync::RwLock<crate::ToolRegistry>>;
214
215pub async fn setup_lazy_mcp(registry: &SharedRegistry) -> usize {
228 setup_lazy_mcp_with(registry, McpDiscoveryOpts::default()).await
229}
230
231#[doc(hidden)]
233pub async fn setup_lazy_mcp_with(registry: &SharedRegistry, opts: McpDiscoveryOpts) -> usize {
234 let Some(config) = load_nonempty_config() else {
235 return 0;
236 };
237 let mut cache = load_cache_or_empty();
238 if seed_missing_descriptors(&config, &cache, &opts).await > 0 {
239 cache = load_cache_or_empty();
242 }
243 register_dormant(registry, &config, &cache, "boot").await
244}
245
246pub async fn setup_cached_mcp(registry: &SharedRegistry) -> (usize, Option<PendingMcpDiscovery>) {
251 let Some(config) = load_nonempty_config() else {
252 return (0, None);
253 };
254 let cache = load_cache_or_empty();
255 let mut missing = HashMap::new();
256 for (name, server) in &config.mcp_servers {
257 let fingerprint = descriptors::server_config_fingerprint(server);
258 let cached = cache
259 .servers
260 .get(name)
261 .is_some_and(|entry| entry.fingerprint == fingerprint);
262 if !cached {
263 missing.insert(name.clone(), server.clone());
264 }
265 }
266 let count = register_dormant(registry, &config, &cache, "cached").await;
267 let pending = (!missing.is_empty() && descriptors::cache_writeback_enabled()).then_some(
268 PendingMcpDiscovery {
269 config: McpConfig {
270 mcp_servers: missing,
271 },
272 },
273 );
274 (count, pending)
275}
276
277#[derive(Debug)]
281pub struct PendingMcpDiscovery {
282 config: McpConfig,
283}
284
285impl PendingMcpDiscovery {
286 pub fn server_names(&self) -> Vec<String> {
288 let mut names: Vec<String> = self.config.mcp_servers.keys().cloned().collect();
289 names.sort();
290 names
291 }
292
293 pub async fn run(self, registry: &SharedRegistry) -> usize {
307 self.run_with(registry, McpDiscoveryOpts::default()).await
308 }
309
310 #[doc(hidden)]
312 pub async fn run_with(self, registry: &SharedRegistry, opts: McpDiscoveryOpts) -> usize {
313 let mut cache = load_cache_or_empty();
316 if seed_missing_descriptors(&self.config, &cache, &opts).await > 0 {
317 cache = load_cache_or_empty();
318 }
319 register_dormant(registry, &self.config, &cache, "background").await
320 }
321}
322
323fn load_nonempty_config() -> Option<McpConfig> {
324 load_mcp_config().filter(|config| !config.mcp_servers.is_empty())
325}
326
327fn load_cache_or_empty() -> descriptors::McpDescriptorCache {
328 match descriptors::load_default_cache() {
329 Ok(cache) => cache,
330 Err(descriptors::DescriptorCacheError::NotFound) => descriptors::McpDescriptorCache::empty(),
331 Err(err) => {
332 tracing::warn!(error = %err,
333 "Refusing unsafe MCP descriptor cache; MCP capabilities stay undiscoverable this run");
334 descriptors::McpDescriptorCache::empty()
335 }
336 }
337}
338
339async fn register_dormant(
342 registry: &SharedRegistry,
343 config: &McpConfig,
344 cache: &descriptors::McpDescriptorCache,
345 phase: &'static str,
346) -> usize {
347 let server_count = config.mcp_servers.len();
348 let dormant = descriptors::dormant_tools_for_config(config, cache);
349 if dormant.is_empty() {
350 tracing::info!(
351 servers = server_count,
352 phase,
353 "MCP: no valid cached descriptors; MCP tools are not discoverable this run"
354 );
355 return 0;
356 }
357 match registry.write().await.try_register_batch(dormant) {
358 Ok(count) => {
359 tracing::info!(
360 tools = count,
361 servers = server_count,
362 phase,
363 "MCP: descriptor-backed tools registered (no process started)"
364 );
365 count
366 }
367 Err(e) => {
368 tracing::warn!(error = %e, phase,
369 "MCP descriptor batch rejected; failing closed with zero MCP capabilities");
370 0
371 }
372 }
373}
374
375fn format_unix(secs: u64) -> String {
376 i64::try_from(secs)
377 .ok()
378 .and_then(|s| chrono::DateTime::<chrono::Utc>::from_timestamp(s, 0))
379 .map_or_else(|| secs.to_string(), |dt| dt.to_rfc3339())
380}
381
382async fn seed_missing_descriptors(
396 config: &McpConfig,
397 cache: &descriptors::McpDescriptorCache,
398 opts: &McpDiscoveryOpts,
399) -> usize {
400 use futures::StreamExt;
401
402 if !descriptors::cache_writeback_enabled() {
403 return 0;
404 }
405 let now = opts.now_unix();
406 let mut names: Vec<&String> = config.mcp_servers.keys().collect();
407 names.sort();
408 let mut pending = Vec::new();
409 for name in names {
410 let server = &config.mcp_servers[name];
411 let fingerprint = descriptors::server_config_fingerprint(server);
412 if cache
413 .servers
414 .get(name)
415 .is_some_and(|entry| entry.fingerprint == fingerprint)
416 {
417 continue;
418 }
419 if let Some(failure) = cache.discovery_failures.get(name) {
420 if failure.suppresses(&fingerprint, now) {
421 let when = format_unix(failure.retry_after_unix());
422 tracing::info!(
423 server = %name,
424 attempts = failure.attempts,
425 "MCP discovery for '{}' failed recently; retrying after {}",
426 name,
427 when
428 );
429 continue;
430 }
431 }
432 pending.push((name.clone(), server.clone(), fingerprint));
433 }
434 if pending.is_empty() {
435 return 0;
436 }
437 let opts = *opts;
438 let outcomes: Vec<bool> = futures::stream::iter(
439 pending
440 .into_iter()
441 .map(|(name, server, fingerprint)| discover_one(name, server, fingerprint, opts)),
442 )
443 .buffer_unordered(DISCOVERY_CONCURRENCY)
444 .collect()
445 .await;
446 outcomes.into_iter().filter(|recorded| *recorded).count()
447}
448
449async fn discover_one(
453 name: String,
454 server: McpServerConfig,
455 fingerprint: String,
456 opts: McpDiscoveryOpts,
457) -> bool {
458 let listing = tokio::time::timeout(opts.timeout, async {
459 let mut conn = connection::McpConnection::start(&server).await?;
460 let result = conn.list_tools().await;
461 conn.start_kill();
462 result
463 })
464 .await;
465 let path = descriptors::default_cache_path();
466 let defs = match listing {
467 Ok(Ok(defs)) => defs,
468 Ok(Err(err)) => {
469 tracing::warn!(server = %name, error = %err, "MCP discovery failed; its tools stay undiscoverable until a retry");
470 record_failure(path, name, fingerprint, opts.now_unix()).await;
471 return false;
472 }
473 Err(_) => {
474 tracing::warn!(server = %name, "MCP discovery timed out; its tools stay undiscoverable until a retry");
475 record_failure(path, name, fingerprint, opts.now_unix()).await;
476 return false;
477 }
478 };
479 let tools: Vec<descriptors::CachedToolDescriptor> = defs
480 .into_iter()
481 .map(|def| descriptors::CachedToolDescriptor {
482 name: def.name,
483 description: def.description,
484 input_schema: def.input_schema,
485 })
486 .collect();
487 let server_name = name.clone();
488 let written = tokio::task::spawn_blocking(move || {
489 descriptors::record_server_listing(&path, &server_name, &fingerprint, &tools)
490 })
491 .await;
492 match written {
493 Ok(Ok(count)) => {
494 tracing::info!(server = %name, tools = count, "MCP server discovered and cached");
495 true
496 }
497 Ok(Err(err)) => {
498 tracing::warn!(server = %name, error = %err, "MCP discovery could not be cached");
499 false
500 }
501 Err(err) => {
502 tracing::warn!(server = %name, error = %err, "MCP discovery cache write panicked");
503 false
504 }
505 }
506}
507
508async fn record_failure(path: std::path::PathBuf, name: String, fingerprint: String, now: u64) {
509 let server_name = name.clone();
510 let written = tokio::task::spawn_blocking(move || {
511 descriptors::record_discovery_failure(&path, &server_name, &fingerprint, now)
512 })
513 .await;
514 match written {
515 Ok(Ok(failure)) => tracing::info!(
516 server = %name,
517 attempts = failure.attempts,
518 retry_after = %format_unix(failure.retry_after_unix()),
519 "MCP discovery failure recorded"
520 ),
521 Ok(Err(err)) => {
522 tracing::warn!(server = %name, error = %err, "MCP discovery failure could not be cached")
523 }
524 Err(err) => {
525 tracing::warn!(server = %name, error = %err, "MCP discovery failure write panicked")
526 }
527 }
528}
529
530#[cfg(test)]
531mod tests {
532 use super::*;
533 use serde_json::json;
534
535 #[test]
536 fn test_mcp_config_deserialize() {
537 let json_str = r#"{"mcpServers": {"test": {"command": "echo", "args": ["hi"]}}}"#;
538 let config: McpConfig = serde_json::from_str(json_str).unwrap();
539
540 assert_eq!(config.mcp_servers.len(), 1);
541 assert!(config.mcp_servers.contains_key("test"));
542
543 let server = &config.mcp_servers["test"];
544 assert_eq!(server.command, "echo");
545 assert_eq!(server.args, vec!["hi"]);
546 }
547
548 #[test]
549 fn test_mcp_config_empty_servers() {
550 let json_str = r#"{"mcpServers": {}}"#;
551 let config: McpConfig = serde_json::from_str(json_str).unwrap();
552
553 assert_eq!(config.mcp_servers.len(), 0);
554 assert!(config.mcp_servers.is_empty());
555 }
556
557 #[test]
558 fn test_mcp_server_config_defaults() {
559 let json_str = r#"{"command": "echo"}"#;
560 let server_config: McpServerConfig = serde_json::from_str(json_str).unwrap();
561
562 assert_eq!(server_config.command, "echo");
563 assert_eq!(server_config.args, Vec::<String>::new());
564 assert_eq!(server_config.env, HashMap::new());
565 }
566
567 #[test]
568 fn test_mcp_config_deserialize_from_value() {
569 let json_value = json!({
570 "mcpServers": {
571 "test": {
572 "command": "echo",
573 "args": ["hi"]
574 }
575 }
576 });
577
578 let config: McpConfig = serde_json::from_value(json_value).unwrap();
579
580 assert_eq!(config.mcp_servers.len(), 1);
581 assert!(config.mcp_servers.contains_key("test"));
582
583 let server = &config.mcp_servers["test"];
584 assert_eq!(server.command, "echo");
585 assert_eq!(server.args, vec!["hi"]);
586 }
587
588 fn base_server() -> McpServerConfig {
589 McpServerConfig {
590 command: "/bin/true".to_string(),
591 args: vec!["--flag".to_string()],
592 env: HashMap::from([("KEY".to_string(), "value".to_string())]),
593 shared: false,
594 }
595 }
596
597 fn config_of(servers: Vec<(&str, McpServerConfig)>) -> McpConfig {
598 McpConfig {
599 mcp_servers: servers
600 .into_iter()
601 .map(|(k, v)| (k.to_string(), v))
602 .collect(),
603 }
604 }
605
606 #[test]
607 fn validate_mcp_config_accepts_sane_configs() {
608 assert!(validate_mcp_config(&config_of(vec![("srv", base_server())])).is_ok());
609 assert!(validate_mcp_config(&config_of(vec![])).is_ok());
610 }
611
612 #[test]
613 fn validate_mcp_config_fails_whole_config_closed_on_structural_violations() {
614 let many: Vec<(String, McpServerConfig)> = (0..=MCP_CONFIG_MAX_SERVERS)
616 .map(|i| (format!("s{i}"), base_server()))
617 .collect();
618 let config = McpConfig {
619 mcp_servers: many.into_iter().collect(),
620 };
621 assert_eq!(validate_mcp_config(&config), Err("too_many_servers"));
622
623 let mut c = base_server();
625 assert_eq!(
626 validate_mcp_config(&config_of(vec![("evil\u{7}", c.clone())])),
627 Err("invalid_server_name")
628 );
629 let long_name = "n".repeat(MCP_CONFIG_MAX_SERVER_NAME_BYTES + 1);
630 assert_eq!(
631 validate_mcp_config(&config_of(vec![(long_name.as_str(), c.clone())])),
632 Err("invalid_server_name")
633 );
634
635 c.command = String::new();
637 assert_eq!(
638 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
639 Err("invalid_command")
640 );
641 c.command = "x".repeat(MCP_CONFIG_MAX_COMMAND_BYTES + 1);
642 assert_eq!(
643 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
644 Err("invalid_command")
645 );
646 c.command = "bad\u{0}cmd".to_string();
647 assert_eq!(
648 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
649 Err("invalid_command")
650 );
651
652 let mut c = base_server();
654 c.args = vec!["a".to_string(); MCP_CONFIG_MAX_ARGS + 1];
655 assert_eq!(
656 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
657 Err("too_many_args")
658 );
659 c.args = vec!["ok".to_string(), "bad\u{1b}arg".to_string()];
660 assert_eq!(
661 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
662 Err("invalid_arg")
663 );
664
665 let mut c = base_server();
667 c.env = (0..=MCP_CONFIG_MAX_ENV_ENTRIES)
668 .map(|i| (format!("K{i}"), "v".to_string()))
669 .collect();
670 assert_eq!(
671 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
672 Err("too_many_env_entries")
673 );
674 c.env = HashMap::from([(String::new(), "v".to_string())]);
675 assert_eq!(
676 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
677 Err("invalid_env_key")
678 );
679 c.env = HashMap::from([(
680 "K".to_string(),
681 "v".repeat(MCP_CONFIG_MAX_ENV_VALUE_BYTES + 1),
682 )]);
683 assert_eq!(
684 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
685 Err("invalid_env_value")
686 );
687 }
688
689 #[test]
690 fn test_load_mcp_config_returns_some_or_none() {
691 let result = load_mcp_config();
694
695 match result {
698 Some(_config) => {
699 }
702 None => {
703 }
706 }
707 }
708}