ricecoder_mcp/
config.rs

1//! Configuration management for MCP
2
3use 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/// MCP Configuration
11#[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/// MCP Server Configuration
19#[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/// Custom Tool Configuration
32#[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/// Parameter Configuration
44#[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/// Permission Configuration
54#[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    /// Creates a new empty MCP configuration
63    pub fn new() -> Self {
64        Self {
65            servers: Vec::new(),
66            custom_tools: Vec::new(),
67            permissions: Vec::new(),
68        }
69    }
70
71    /// Adds an MCP server configuration
72    pub fn add_server(&mut self, server: MCPServerConfig) {
73        self.servers.push(server);
74    }
75
76    /// Adds a custom tool configuration
77    pub fn add_custom_tool(&mut self, tool: CustomToolConfig) {
78        self.custom_tools.push(tool);
79    }
80
81    /// Adds a permission configuration
82    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
93/// Configuration loader for MCP using ricecoder-storage
94pub struct MCPConfigLoader;
95
96impl MCPConfigLoader {
97    /// Load MCP configuration from a file
98    ///
99    /// Supports YAML and JSON formats. Automatically detects format based on file extension.
100    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    /// Load MCP configuration from a string
127    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    /// Load MCP configuration from multiple sources with precedence
142    ///
143    /// Loads configuration from project, user, and default locations.
144    /// Later sources override earlier ones.
145    ///
146    /// Precedence (highest to lowest):
147    /// 1. Project-level: `.ricecoder/mcp-servers.yaml`, `.ricecoder/custom-tools.json`, `.ricecoder/permissions.yaml`
148    /// 2. User-level: `~/.ricecoder/mcp-servers.yaml`, `~/.ricecoder/custom-tools.json`, `~/.ricecoder/permissions.yaml`
149    /// 3. Built-in defaults
150    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        // Load from user-level first (lowest priority)
157        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        // Load from project-level (highest priority)
165        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    /// Load MCP configuration from a directory
176    ///
177    /// Looks for:
178    /// - `mcp-servers.yaml` or `mcp-servers.json` for server configurations
179    /// - `custom-tools.json` or `custom-tools.md` for custom tool definitions
180    /// - `permissions.yaml` or `permissions.json` for permission configurations
181    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        // Try to load MCP servers configuration
186        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        // Try to load custom tools configuration
202        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        // Try to load permissions configuration
218        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    /// Load custom tools from a Markdown file with YAML frontmatter
237    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        // Parse YAML frontmatter if present
243        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        // If no frontmatter, try parsing the entire content as YAML
253        serde_yaml::from_str(&content)
254            .map_err(|e| Error::ConfigError(format!("Failed to parse markdown content: {}", e)))
255    }
256
257    /// Merge two MCP configurations
258    ///
259    /// Later configuration overrides earlier one for servers and permissions.
260    /// Custom tools are accumulated.
261    fn merge_configs(mut base: MCPConfig, override_config: MCPConfig) -> MCPConfig {
262        // For servers and permissions, override by ID/pattern
263        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        // Accumulate custom tools
284        base.custom_tools.extend(override_config.custom_tools);
285
286        base
287    }
288
289    /// Validate MCP configuration
290    pub fn validate(config: &MCPConfig) -> Result<()> {
291        // Validate servers
292        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        // Validate custom tools
307        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        // Validate permissions
322        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    /// Save MCP configuration to a file
334    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}