1mod connection;
3pub mod descriptors;
4mod lazy;
5pub mod lease;
6mod tool;
7
8use crate::ToolRegistry;
9use std::collections::HashMap;
10use std::sync::Arc;
11use tokio::sync::Mutex;
12
13pub use lazy::McpConnectTool;
14pub use lease::{McpLeaseCapability, McpRuntimeManager, McpSessionEndGuard};
15pub use tool::McpTool;
16
17#[derive(Debug, Clone, serde::Deserialize)]
19pub struct McpServerConfig {
20 pub command: String,
21 #[serde(default)]
22 pub args: Vec<String>,
23 #[serde(default)]
24 pub env: HashMap<String, String>,
25 #[serde(default)]
32 pub shared: bool,
33}
34
35#[derive(Debug, Clone, serde::Deserialize)]
37pub struct McpConfig {
38 #[serde(rename = "mcpServers", default)]
39 pub mcp_servers: HashMap<String, McpServerConfig>,
40}
41
42#[derive(Debug, Clone)]
44pub(crate) struct McpToolDef {
45 pub(crate) name: String,
46 pub(crate) description: String,
47 pub(crate) input_schema: serde_json::Value,
48}
49
50pub const MCP_CONFIG_MAX_BYTES: u64 = 1024 * 1024;
52
53pub fn load_mcp_config() -> Option<McpConfig> {
61 let path = crate::config::resolve_read_path("mcp.json");
62 let content = match descriptors::read_bounded_regular_file(&path, MCP_CONFIG_MAX_BYTES) {
63 Ok(content) => content,
64 Err(descriptors::DescriptorCacheError::NotFound) => return None,
65 Err(err) => {
66 let category = match err {
68 descriptors::DescriptorCacheError::NotFound => "not_found",
69 descriptors::DescriptorCacheError::NotRegularFile => "not_regular_file",
70 descriptors::DescriptorCacheError::Oversize { .. } => "oversize",
71 descriptors::DescriptorCacheError::Io(_) => "io",
72 descriptors::DescriptorCacheError::Parse(_) => "parse",
73 descriptors::DescriptorCacheError::Version(_) => "version",
74 };
75 tracing::warn!(
76 file = "mcp.json",
77 category,
78 "Refusing unsafe MCP config read"
79 );
80 return None;
81 }
82 };
83 let config = match serde_json::from_str::<McpConfig>(&content) {
84 Ok(config) => config,
85 Err(e) => {
86 tracing::warn!(
89 file = "mcp.json",
90 class = ?e.classify(),
91 line = e.line(),
92 column = e.column(),
93 "Failed to parse MCP config"
94 );
95 return None;
96 }
97 };
98 if let Err(reason) = validate_mcp_config(&config) {
99 tracing::warn!(
101 file = "mcp.json",
102 reason,
103 "Refusing structurally unsafe MCP config"
104 );
105 return None;
106 }
107 Some(config)
108}
109
110pub const MCP_CONFIG_MAX_SERVERS: usize = 32;
114pub const MCP_CONFIG_MAX_ARGS: usize = 64;
115pub const MCP_CONFIG_MAX_ENV_ENTRIES: usize = 64;
116pub const MCP_CONFIG_MAX_SERVER_NAME_BYTES: usize = 64;
117pub const MCP_CONFIG_MAX_COMMAND_BYTES: usize = 512;
118pub const MCP_CONFIG_MAX_ARG_BYTES: usize = 4096;
119pub const MCP_CONFIG_MAX_ENV_KEY_BYTES: usize = 128;
120pub const MCP_CONFIG_MAX_ENV_VALUE_BYTES: usize = 4096;
121
122fn safe_str(value: &str, max: usize) -> bool {
123 !value.is_empty() && value.len() <= max && !value.chars().any(char::is_control)
124}
125
126pub fn validate_mcp_config(config: &McpConfig) -> Result<(), &'static str> {
129 if config.mcp_servers.len() > MCP_CONFIG_MAX_SERVERS {
130 return Err("too_many_servers");
131 }
132 for (name, server) in &config.mcp_servers {
133 if !safe_str(name, MCP_CONFIG_MAX_SERVER_NAME_BYTES) {
134 return Err("invalid_server_name");
135 }
136 if !safe_str(&server.command, MCP_CONFIG_MAX_COMMAND_BYTES) {
137 return Err("invalid_command");
138 }
139 if server.args.len() > MCP_CONFIG_MAX_ARGS {
140 return Err("too_many_args");
141 }
142 for arg in &server.args {
143 if arg.len() > MCP_CONFIG_MAX_ARG_BYTES || arg.chars().any(char::is_control) {
144 return Err("invalid_arg");
145 }
146 }
147 if server.env.len() > MCP_CONFIG_MAX_ENV_ENTRIES {
148 return Err("too_many_env_entries");
149 }
150 for (key, value) in &server.env {
151 if !safe_str(key, MCP_CONFIG_MAX_ENV_KEY_BYTES) {
152 return Err("invalid_env_key");
153 }
154 if value.len() > MCP_CONFIG_MAX_ENV_VALUE_BYTES || value.chars().any(char::is_control) {
155 return Err("invalid_env_value");
156 }
157 }
158 }
159 Ok(())
160}
161
162pub async fn connect_mcp_servers(registry: &mut ToolRegistry) -> usize {
165 let config = match load_mcp_config() {
166 Some(c) => c,
167 None => return 0,
168 };
169
170 let mut total_tools = 0;
171
172 for (server_name, server_config) in &config.mcp_servers {
173 tracing::info!(server = %server_name, command = %server_config.command, "Connecting to MCP server");
174
175 match connection::McpConnection::start(server_config).await {
176 Ok(mut conn) => {
177 match conn.list_tools().await {
178 Ok(tools) => {
179 let tool_count = tools.len();
180 let connection = Arc::new(Mutex::new(conn));
181
182 for tool_def in tools {
183 let prefixed_name = format!("ext__{}__{}", server_name, tool_def.name);
186 let tool_name_for_log = prefixed_name.clone();
187
188 let mcp_tool = McpTool {
189 tool_name: prefixed_name,
190 server_tool_name: tool_def.name.clone(),
191 server_name: server_name.clone(),
192 description: format!(
193 "[MCP:{}] {}",
194 server_name, tool_def.description
195 ),
196 input_schema: tool_def.input_schema,
197 connection: Arc::clone(&connection),
198 };
199
200 if let Err(e) = registry.try_register(Arc::new(mcp_tool)) {
201 tracing::warn!(
202 server = %server_name,
203 tool = %tool_name_for_log,
204 error = %e,
205 "Refusing to expose MCP tool the capability catalog could not record"
206 );
207 continue;
208 }
209 total_tools += 1;
210 }
211
212 tracing::info!(
213 server = %server_name,
214 tools = tool_count,
215 "MCP server connected — {} tools registered",
216 tool_count
217 );
218 }
219 Err(e) => {
220 tracing::error!(server = %server_name, error = %e, "Failed to list MCP tools");
221 }
222 }
223 }
224 Err(e) => {
225 tracing::error!(server = %server_name, error = %e, "Failed to connect to MCP server");
226 }
227 }
228 }
229
230 total_tools
231}
232
233pub async fn setup_lazy_mcp(
249 registry: &Arc<tokio::sync::RwLock<crate::ToolRegistry>>,
250 progressive_exact_only: bool,
251) -> usize {
252 let config = match load_mcp_config() {
253 Some(c) => c,
254 None => return 0,
255 };
256
257 let server_count = config.mcp_servers.len();
258 if server_count == 0 {
259 return 0;
260 }
261
262 if progressive_exact_only {
263 let cache = match descriptors::load_default_cache() {
264 Ok(cache) => cache,
265 Err(descriptors::DescriptorCacheError::NotFound) => {
266 descriptors::McpDescriptorCache::empty()
267 }
268 Err(err) => {
269 tracing::warn!(error = %err,
270 "Refusing unsafe MCP descriptor cache; MCP capabilities stay undiscoverable this run");
271 descriptors::McpDescriptorCache::empty()
272 }
273 };
274 let dormant = descriptors::dormant_tools_for_config(&config, &cache);
275 if dormant.is_empty() {
276 tracing::info!(
277 servers = server_count,
278 "MCP exact mode: no valid cached descriptors; MCP tools are not discoverable this run"
279 );
280 return 0;
281 }
282 return match registry.write().await.try_register_batch(dormant) {
283 Ok(count) => {
284 tracing::info!(
285 tools = count,
286 servers = server_count,
287 "MCP exact mode: dormant descriptor-backed tools registered (no process started)"
288 );
289 count
290 }
291 Err(e) => {
292 tracing::warn!(error = %e,
293 "MCP dormant descriptor batch rejected; failing closed with zero MCP capabilities");
294 0
295 }
296 };
297 }
298
299 let server_names: Vec<&str> = config.mcp_servers.keys().map(|s| s.as_str()).collect();
300 tracing::info!(servers = ?server_names, "MCP lazy loading: {} servers available", server_count);
301
302 let connect_tool = McpConnectTool::new(config.mcp_servers);
303
304 if let Err(e) = registry.write().await.try_register(Arc::new(connect_tool)) {
305 tracing::warn!(error = %e, "Refusing to expose MCP gateway the capability catalog could not record");
306 return 0;
307 }
308
309 server_count
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315 use serde_json::json;
316
317 #[test]
318 fn test_mcp_config_deserialize() {
319 let json_str = r#"{"mcpServers": {"test": {"command": "echo", "args": ["hi"]}}}"#;
320 let config: McpConfig = serde_json::from_str(json_str).unwrap();
321
322 assert_eq!(config.mcp_servers.len(), 1);
323 assert!(config.mcp_servers.contains_key("test"));
324
325 let server = &config.mcp_servers["test"];
326 assert_eq!(server.command, "echo");
327 assert_eq!(server.args, vec!["hi"]);
328 }
329
330 #[test]
331 fn test_mcp_config_empty_servers() {
332 let json_str = r#"{"mcpServers": {}}"#;
333 let config: McpConfig = serde_json::from_str(json_str).unwrap();
334
335 assert_eq!(config.mcp_servers.len(), 0);
336 assert!(config.mcp_servers.is_empty());
337 }
338
339 #[test]
340 fn test_mcp_server_config_defaults() {
341 let json_str = r#"{"command": "echo"}"#;
342 let server_config: McpServerConfig = serde_json::from_str(json_str).unwrap();
343
344 assert_eq!(server_config.command, "echo");
345 assert_eq!(server_config.args, Vec::<String>::new());
346 assert_eq!(server_config.env, HashMap::new());
347 }
348
349 #[test]
350 fn test_mcp_config_deserialize_from_value() {
351 let json_value = json!({
352 "mcpServers": {
353 "test": {
354 "command": "echo",
355 "args": ["hi"]
356 }
357 }
358 });
359
360 let config: McpConfig = serde_json::from_value(json_value).unwrap();
361
362 assert_eq!(config.mcp_servers.len(), 1);
363 assert!(config.mcp_servers.contains_key("test"));
364
365 let server = &config.mcp_servers["test"];
366 assert_eq!(server.command, "echo");
367 assert_eq!(server.args, vec!["hi"]);
368 }
369
370 fn base_server() -> McpServerConfig {
371 McpServerConfig {
372 command: "/bin/true".to_string(),
373 args: vec!["--flag".to_string()],
374 env: HashMap::from([("KEY".to_string(), "value".to_string())]),
375 shared: false,
376 }
377 }
378
379 fn config_of(servers: Vec<(&str, McpServerConfig)>) -> McpConfig {
380 McpConfig {
381 mcp_servers: servers
382 .into_iter()
383 .map(|(k, v)| (k.to_string(), v))
384 .collect(),
385 }
386 }
387
388 #[test]
389 fn validate_mcp_config_accepts_sane_configs() {
390 assert!(validate_mcp_config(&config_of(vec![("srv", base_server())])).is_ok());
391 assert!(validate_mcp_config(&config_of(vec![])).is_ok());
392 }
393
394 #[test]
395 fn validate_mcp_config_fails_whole_config_closed_on_structural_violations() {
396 let many: Vec<(String, McpServerConfig)> = (0..=MCP_CONFIG_MAX_SERVERS)
398 .map(|i| (format!("s{i}"), base_server()))
399 .collect();
400 let config = McpConfig {
401 mcp_servers: many.into_iter().collect(),
402 };
403 assert_eq!(validate_mcp_config(&config), Err("too_many_servers"));
404
405 let mut c = base_server();
407 assert_eq!(
408 validate_mcp_config(&config_of(vec![("evil\u{7}", c.clone())])),
409 Err("invalid_server_name")
410 );
411 let long_name = "n".repeat(MCP_CONFIG_MAX_SERVER_NAME_BYTES + 1);
412 assert_eq!(
413 validate_mcp_config(&config_of(vec![(long_name.as_str(), c.clone())])),
414 Err("invalid_server_name")
415 );
416
417 c.command = String::new();
419 assert_eq!(
420 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
421 Err("invalid_command")
422 );
423 c.command = "x".repeat(MCP_CONFIG_MAX_COMMAND_BYTES + 1);
424 assert_eq!(
425 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
426 Err("invalid_command")
427 );
428 c.command = "bad\u{0}cmd".to_string();
429 assert_eq!(
430 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
431 Err("invalid_command")
432 );
433
434 let mut c = base_server();
436 c.args = vec!["a".to_string(); MCP_CONFIG_MAX_ARGS + 1];
437 assert_eq!(
438 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
439 Err("too_many_args")
440 );
441 c.args = vec!["ok".to_string(), "bad\u{1b}arg".to_string()];
442 assert_eq!(
443 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
444 Err("invalid_arg")
445 );
446
447 let mut c = base_server();
449 c.env = (0..=MCP_CONFIG_MAX_ENV_ENTRIES)
450 .map(|i| (format!("K{i}"), "v".to_string()))
451 .collect();
452 assert_eq!(
453 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
454 Err("too_many_env_entries")
455 );
456 c.env = HashMap::from([(String::new(), "v".to_string())]);
457 assert_eq!(
458 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
459 Err("invalid_env_key")
460 );
461 c.env = HashMap::from([(
462 "K".to_string(),
463 "v".repeat(MCP_CONFIG_MAX_ENV_VALUE_BYTES + 1),
464 )]);
465 assert_eq!(
466 validate_mcp_config(&config_of(vec![("srv", c.clone())])),
467 Err("invalid_env_value")
468 );
469 }
470
471 #[test]
472 fn test_load_mcp_config_returns_some_or_none() {
473 let result = load_mcp_config();
476
477 match result {
480 Some(_config) => {
481 }
484 None => {
485 }
488 }
489 }
490}