1use crate::error::{Error, Result};
4use ricecoder_storage::types::ConfigFormat;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::path::Path;
8use tracing::{debug, info};
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct MCPConfig {
13 pub servers: Vec<MCPServerConfig>,
14 pub custom_tools: Vec<CustomToolConfig>,
15 pub permissions: Vec<PermissionConfig>,
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct MCPServerConfig {
21 pub id: String,
22 pub name: String,
23 pub command: String,
24 pub args: Vec<String>,
25 pub env: HashMap<String, String>,
26 pub timeout_ms: u64,
27 pub auto_reconnect: bool,
28 pub max_retries: u32,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct CustomToolConfig {
34 pub id: String,
35 pub name: String,
36 pub description: String,
37 pub category: String,
38 pub parameters: Vec<ParameterConfig>,
39 pub return_type: String,
40 pub handler: String,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct ParameterConfig {
46 pub name: String,
47 pub type_: String,
48 pub description: String,
49 pub required: bool,
50 pub default: Option<serde_json::Value>,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct PermissionConfig {
56 pub pattern: String,
57 pub level: String,
58 pub agent_id: Option<String>,
59}
60
61impl MCPConfig {
62 pub fn new() -> Self {
64 Self {
65 servers: Vec::new(),
66 custom_tools: Vec::new(),
67 permissions: Vec::new(),
68 }
69 }
70
71 pub fn add_server(&mut self, server: MCPServerConfig) {
73 self.servers.push(server);
74 }
75
76 pub fn add_custom_tool(&mut self, tool: CustomToolConfig) {
78 self.custom_tools.push(tool);
79 }
80
81 pub fn add_permission(&mut self, permission: PermissionConfig) {
83 self.permissions.push(permission);
84 }
85}
86
87impl Default for MCPConfig {
88 fn default() -> Self {
89 Self::new()
90 }
91}
92
93pub struct MCPConfigLoader;
95
96impl MCPConfigLoader {
97 pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<MCPConfig> {
101 let path = path.as_ref();
102 debug!("Loading MCP configuration from: {:?}", path);
103
104 let content = std::fs::read_to_string(path)
105 .map_err(|e| Error::ConfigError(format!("Failed to read config file: {}", e)))?;
106
107 let extension = path
108 .extension()
109 .and_then(|ext| ext.to_str())
110 .ok_or_else(|| Error::ConfigError("Config file has no extension".to_string()))?;
111
112 let format = match extension {
113 "yaml" | "yml" => ConfigFormat::Yaml,
114 "json" => ConfigFormat::Json,
115 _ => {
116 return Err(Error::ConfigError(format!(
117 "Unsupported config format: {}",
118 extension
119 )))
120 }
121 };
122
123 Self::load_from_string(&content, format)
124 }
125
126 pub fn load_from_string(content: &str, format: ConfigFormat) -> Result<MCPConfig> {
128 debug!("Parsing MCP configuration from string");
129
130 match format {
131 ConfigFormat::Yaml => serde_yaml::from_str(content)
132 .map_err(|e| Error::ConfigError(format!("Failed to parse YAML: {}", e))),
133 ConfigFormat::Json => serde_json::from_str(content)
134 .map_err(|e| Error::ConfigError(format!("Failed to parse JSON: {}", e))),
135 ConfigFormat::Toml => Err(Error::ConfigError(
136 "TOML format is not supported for MCP configuration".to_string(),
137 )),
138 }
139 }
140
141 pub fn load_with_precedence(
151 project_dir: Option<&Path>,
152 user_dir: Option<&Path>,
153 ) -> Result<MCPConfig> {
154 let mut config = MCPConfig::new();
155
156 if let Some(user_dir) = user_dir {
158 if let Ok(user_config) = Self::load_from_directory(user_dir) {
159 info!("Loaded user-level MCP configuration");
160 config = Self::merge_configs(config, user_config);
161 }
162 }
163
164 if let Some(project_dir) = project_dir {
166 if let Ok(project_config) = Self::load_from_directory(project_dir) {
167 info!("Loaded project-level MCP configuration");
168 config = Self::merge_configs(config, project_config);
169 }
170 }
171
172 Ok(config)
173 }
174
175 pub fn load_from_directory<P: AsRef<Path>>(dir: P) -> Result<MCPConfig> {
182 let dir = dir.as_ref();
183 let mut config = MCPConfig::new();
184
185 let servers_yaml = dir.join("mcp-servers.yaml");
187 let servers_json = dir.join("mcp-servers.json");
188
189 if servers_yaml.exists() {
190 debug!("Loading MCP servers from: {:?}", servers_yaml);
191 if let Ok(servers_config) = Self::load_from_file(&servers_yaml) {
192 config.servers.extend(servers_config.servers);
193 }
194 } else if servers_json.exists() {
195 debug!("Loading MCP servers from: {:?}", servers_json);
196 if let Ok(servers_config) = Self::load_from_file(&servers_json) {
197 config.servers.extend(servers_config.servers);
198 }
199 }
200
201 let custom_tools_json = dir.join("custom-tools.json");
203 let custom_tools_md = dir.join("custom-tools.md");
204
205 if custom_tools_json.exists() {
206 debug!("Loading custom tools from: {:?}", custom_tools_json);
207 if let Ok(tools_config) = Self::load_from_file(&custom_tools_json) {
208 config.custom_tools.extend(tools_config.custom_tools);
209 }
210 } else if custom_tools_md.exists() {
211 debug!("Loading custom tools from markdown: {:?}", custom_tools_md);
212 if let Ok(tools_config) = Self::load_custom_tools_from_markdown(&custom_tools_md) {
213 config.custom_tools.extend(tools_config);
214 }
215 }
216
217 let permissions_yaml = dir.join("permissions.yaml");
219 let permissions_json = dir.join("permissions.json");
220
221 if permissions_yaml.exists() {
222 debug!("Loading permissions from: {:?}", permissions_yaml);
223 if let Ok(perms_config) = Self::load_from_file(&permissions_yaml) {
224 config.permissions.extend(perms_config.permissions);
225 }
226 } else if permissions_json.exists() {
227 debug!("Loading permissions from: {:?}", permissions_json);
228 if let Ok(perms_config) = Self::load_from_file(&permissions_json) {
229 config.permissions.extend(perms_config.permissions);
230 }
231 }
232
233 Ok(config)
234 }
235
236 fn load_custom_tools_from_markdown<P: AsRef<Path>>(path: P) -> Result<Vec<CustomToolConfig>> {
238 let path = path.as_ref();
239 let content = std::fs::read_to_string(path)
240 .map_err(|e| Error::ConfigError(format!("Failed to read markdown file: {}", e)))?;
241
242 if content.starts_with("---") {
244 if let Some(end_idx) = content[3..].find("---") {
245 let frontmatter = &content[3..end_idx + 3];
246 let tools: Vec<CustomToolConfig> = serde_yaml::from_str(frontmatter)
247 .map_err(|e| Error::ConfigError(format!("Failed to parse markdown frontmatter: {}", e)))?;
248 return Ok(tools);
249 }
250 }
251
252 serde_yaml::from_str(&content)
254 .map_err(|e| Error::ConfigError(format!("Failed to parse markdown content: {}", e)))
255 }
256
257 fn merge_configs(mut base: MCPConfig, override_config: MCPConfig) -> MCPConfig {
262 for server in override_config.servers {
264 if let Some(pos) = base.servers.iter().position(|s| s.id == server.id) {
265 base.servers[pos] = server;
266 } else {
267 base.servers.push(server);
268 }
269 }
270
271 for permission in override_config.permissions {
272 if let Some(pos) = base
273 .permissions
274 .iter()
275 .position(|p| p.pattern == permission.pattern && p.agent_id == permission.agent_id)
276 {
277 base.permissions[pos] = permission;
278 } else {
279 base.permissions.push(permission);
280 }
281 }
282
283 base.custom_tools.extend(override_config.custom_tools);
285
286 base
287 }
288
289 pub fn validate(config: &MCPConfig) -> Result<()> {
291 for server in &config.servers {
293 if server.id.is_empty() {
294 return Err(Error::ValidationError(
295 "Server ID cannot be empty".to_string(),
296 ));
297 }
298 if server.command.is_empty() {
299 return Err(Error::ValidationError(format!(
300 "Server '{}' has no command",
301 server.id
302 )));
303 }
304 }
305
306 for tool in &config.custom_tools {
308 if tool.id.is_empty() {
309 return Err(Error::ValidationError(
310 "Custom tool ID cannot be empty".to_string(),
311 ));
312 }
313 if tool.handler.is_empty() {
314 return Err(Error::ValidationError(format!(
315 "Custom tool '{}' has no handler",
316 tool.id
317 )));
318 }
319 }
320
321 for perm in &config.permissions {
323 if perm.pattern.is_empty() {
324 return Err(Error::ValidationError(
325 "Permission pattern cannot be empty".to_string(),
326 ));
327 }
328 }
329
330 Ok(())
331 }
332
333 pub fn save_to_file<P: AsRef<Path>>(config: &MCPConfig, path: P) -> Result<()> {
335 let path = path.as_ref();
336 debug!("Saving MCP configuration to: {:?}", path);
337
338 let extension = path
339 .extension()
340 .and_then(|ext| ext.to_str())
341 .ok_or_else(|| Error::ConfigError("Config file has no extension".to_string()))?;
342
343 let content = match extension {
344 "yaml" | "yml" => serde_yaml::to_string(config)
345 .map_err(|e| Error::ConfigError(format!("Failed to serialize to YAML: {}", e)))?,
346 "json" => serde_json::to_string_pretty(config)
347 .map_err(|e| Error::ConfigError(format!("Failed to serialize to JSON: {}", e)))?,
348 _ => {
349 return Err(Error::ConfigError(format!(
350 "Unsupported config format: {}",
351 extension
352 )))
353 }
354 };
355
356 std::fs::write(path, content)
357 .map_err(|e| Error::ConfigError(format!("Failed to write config file: {}", e)))?;
358
359 info!("MCP configuration saved to: {:?}", path);
360 Ok(())
361 }
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367 use tempfile::TempDir;
368
369 #[test]
370 fn test_create_config() {
371 let config = MCPConfig::new();
372 assert_eq!(config.servers.len(), 0);
373 assert_eq!(config.custom_tools.len(), 0);
374 assert_eq!(config.permissions.len(), 0);
375 }
376
377 #[test]
378 fn test_add_server() {
379 let mut config = MCPConfig::new();
380 let server = MCPServerConfig {
381 id: "test-server".to_string(),
382 name: "Test Server".to_string(),
383 command: "test".to_string(),
384 args: vec![],
385 env: HashMap::new(),
386 timeout_ms: 5000,
387 auto_reconnect: true,
388 max_retries: 3,
389 };
390
391 config.add_server(server);
392 assert_eq!(config.servers.len(), 1);
393 }
394
395 #[test]
396 fn test_add_custom_tool() {
397 let mut config = MCPConfig::new();
398 let tool = CustomToolConfig {
399 id: "test-tool".to_string(),
400 name: "Test Tool".to_string(),
401 description: "A test tool".to_string(),
402 category: "test".to_string(),
403 parameters: vec![],
404 return_type: "string".to_string(),
405 handler: "test::handler".to_string(),
406 };
407
408 config.add_custom_tool(tool);
409 assert_eq!(config.custom_tools.len(), 1);
410 }
411
412 #[test]
413 fn test_load_yaml_config() {
414 let yaml_content = r#"
415servers:
416 - id: test-server
417 name: Test Server
418 command: test
419 args: []
420 env: {}
421 timeout_ms: 5000
422 auto_reconnect: true
423 max_retries: 3
424custom_tools: []
425permissions: []
426"#;
427 let config = MCPConfigLoader::load_from_string(yaml_content, ConfigFormat::Yaml)
428 .expect("Failed to load YAML config");
429 assert_eq!(config.servers.len(), 1);
430 assert_eq!(config.servers[0].id, "test-server");
431 }
432
433 #[test]
434 fn test_load_json_config() {
435 let json_content = r#"{
436 "servers": [
437 {
438 "id": "test-server",
439 "name": "Test Server",
440 "command": "test",
441 "args": [],
442 "env": {},
443 "timeout_ms": 5000,
444 "auto_reconnect": true,
445 "max_retries": 3
446 }
447 ],
448 "custom_tools": [],
449 "permissions": []
450}"#;
451 let config = MCPConfigLoader::load_from_string(json_content, ConfigFormat::Json)
452 .expect("Failed to load JSON config");
453 assert_eq!(config.servers.len(), 1);
454 assert_eq!(config.servers[0].id, "test-server");
455 }
456
457 #[test]
458 fn test_validate_config_valid() {
459 let mut config = MCPConfig::new();
460 config.add_server(MCPServerConfig {
461 id: "test-server".to_string(),
462 name: "Test Server".to_string(),
463 command: "test".to_string(),
464 args: vec![],
465 env: HashMap::new(),
466 timeout_ms: 5000,
467 auto_reconnect: true,
468 max_retries: 3,
469 });
470
471 assert!(MCPConfigLoader::validate(&config).is_ok());
472 }
473
474 #[test]
475 fn test_validate_config_empty_server_id() {
476 let mut config = MCPConfig::new();
477 config.add_server(MCPServerConfig {
478 id: "".to_string(),
479 name: "Test Server".to_string(),
480 command: "test".to_string(),
481 args: vec![],
482 env: HashMap::new(),
483 timeout_ms: 5000,
484 auto_reconnect: true,
485 max_retries: 3,
486 });
487
488 assert!(MCPConfigLoader::validate(&config).is_err());
489 }
490
491 #[test]
492 fn test_validate_config_empty_command() {
493 let mut config = MCPConfig::new();
494 config.add_server(MCPServerConfig {
495 id: "test-server".to_string(),
496 name: "Test Server".to_string(),
497 command: "".to_string(),
498 args: vec![],
499 env: HashMap::new(),
500 timeout_ms: 5000,
501 auto_reconnect: true,
502 max_retries: 3,
503 });
504
505 assert!(MCPConfigLoader::validate(&config).is_err());
506 }
507
508 #[test]
509 fn test_save_and_load_yaml() {
510 let temp_dir = TempDir::new().expect("Failed to create temp dir");
511 let config_path = temp_dir.path().join("config.yaml");
512
513 let mut config = MCPConfig::new();
514 config.add_server(MCPServerConfig {
515 id: "test-server".to_string(),
516 name: "Test Server".to_string(),
517 command: "test".to_string(),
518 args: vec!["arg1".to_string()],
519 env: HashMap::new(),
520 timeout_ms: 5000,
521 auto_reconnect: true,
522 max_retries: 3,
523 });
524
525 MCPConfigLoader::save_to_file(&config, &config_path)
526 .expect("Failed to save config");
527 assert!(config_path.exists());
528
529 let loaded_config = MCPConfigLoader::load_from_file(&config_path)
530 .expect("Failed to load config");
531 assert_eq!(loaded_config.servers.len(), 1);
532 assert_eq!(loaded_config.servers[0].id, "test-server");
533 }
534
535 #[test]
536 fn test_merge_configs() {
537 let mut base = MCPConfig::new();
538 base.add_server(MCPServerConfig {
539 id: "server1".to_string(),
540 name: "Server 1".to_string(),
541 command: "cmd1".to_string(),
542 args: vec![],
543 env: HashMap::new(),
544 timeout_ms: 5000,
545 auto_reconnect: true,
546 max_retries: 3,
547 });
548
549 let mut override_config = MCPConfig::new();
550 override_config.add_server(MCPServerConfig {
551 id: "server1".to_string(),
552 name: "Server 1 Updated".to_string(),
553 command: "cmd1_updated".to_string(),
554 args: vec![],
555 env: HashMap::new(),
556 timeout_ms: 10000,
557 auto_reconnect: false,
558 max_retries: 5,
559 });
560
561 let merged = MCPConfigLoader::merge_configs(base, override_config);
562 assert_eq!(merged.servers.len(), 1);
563 assert_eq!(merged.servers[0].name, "Server 1 Updated");
564 assert_eq!(merged.servers[0].timeout_ms, 10000);
565 }
566}