ricecoder_storage/markdown_config/
registry.rs

1//! Configuration registry for storing and retrieving parsed configurations
2
3use crate::markdown_config::error::{MarkdownConfigError, MarkdownConfigResult};
4use crate::markdown_config::types::{AgentConfig, CommandConfig, ModeConfig};
5use std::collections::HashMap;
6use std::sync::RwLock;
7
8/// Central registry for all loaded configurations
9///
10/// Provides thread-safe storage and retrieval of agent, mode, and command configurations.
11/// Uses RwLock for concurrent access patterns where reads are more frequent than writes.
12#[derive(Debug)]
13pub struct ConfigRegistry {
14    /// Registered agent configurations
15    agents: RwLock<HashMap<String, AgentConfig>>,
16    /// Registered mode configurations
17    modes: RwLock<HashMap<String, ModeConfig>>,
18    /// Registered command configurations
19    commands: RwLock<HashMap<String, CommandConfig>>,
20}
21
22impl ConfigRegistry {
23    /// Create a new configuration registry
24    pub fn new() -> Self {
25        Self {
26            agents: RwLock::new(HashMap::new()),
27            modes: RwLock::new(HashMap::new()),
28            commands: RwLock::new(HashMap::new()),
29        }
30    }
31
32    // ============ Agent Registration ============
33
34    /// Register an agent configuration
35    ///
36    /// # Arguments
37    /// * `config` - The agent configuration to register
38    ///
39    /// # Errors
40    /// Returns an error if:
41    /// - An agent with the same name already exists
42    /// - The configuration is invalid
43    pub fn register_agent(&self, config: AgentConfig) -> MarkdownConfigResult<()> {
44        // Validate configuration
45        config.validate().map_err(|e| {
46            MarkdownConfigError::validation_error(format!("Invalid agent configuration: {}", e))
47        })?;
48
49        let mut agents = self.agents.write().map_err(|e| {
50            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
51        })?;
52
53        // Check for duplicate registration
54        if agents.contains_key(&config.name) {
55            return Err(MarkdownConfigError::registration_error(format!(
56                "Agent '{}' is already registered",
57                config.name
58            )));
59        }
60
61        agents.insert(config.name.clone(), config);
62        Ok(())
63    }
64
65    /// Get an agent configuration by name
66    ///
67    /// # Arguments
68    /// * `name` - The name of the agent to retrieve
69    ///
70    /// # Returns
71    /// Returns the agent configuration if found, None otherwise
72    pub fn get_agent(&self, name: &str) -> MarkdownConfigResult<Option<AgentConfig>> {
73        let agents = self.agents.read().map_err(|e| {
74            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
75        })?;
76
77        Ok(agents.get(name).cloned())
78    }
79
80    /// Get all registered agents
81    pub fn get_all_agents(&self) -> MarkdownConfigResult<Vec<AgentConfig>> {
82        let agents = self.agents.read().map_err(|e| {
83            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
84        })?;
85
86        Ok(agents.values().cloned().collect())
87    }
88
89    /// Check if an agent is registered
90    pub fn has_agent(&self, name: &str) -> MarkdownConfigResult<bool> {
91        let agents = self.agents.read().map_err(|e| {
92            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
93        })?;
94
95        Ok(agents.contains_key(name))
96    }
97
98    /// Remove an agent configuration
99    pub fn remove_agent(&self, name: &str) -> MarkdownConfigResult<Option<AgentConfig>> {
100        let mut agents = self.agents.write().map_err(|e| {
101            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
102        })?;
103
104        Ok(agents.remove(name))
105    }
106
107    // ============ Mode Registration ============
108
109    /// Register a mode configuration
110    ///
111    /// # Arguments
112    /// * `config` - The mode configuration to register
113    ///
114    /// # Errors
115    /// Returns an error if:
116    /// - A mode with the same name already exists
117    /// - The configuration is invalid
118    pub fn register_mode(&self, config: ModeConfig) -> MarkdownConfigResult<()> {
119        // Validate configuration
120        config.validate().map_err(|e| {
121            MarkdownConfigError::validation_error(format!("Invalid mode configuration: {}", e))
122        })?;
123
124        let mut modes = self.modes.write().map_err(|e| {
125            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
126        })?;
127
128        // Check for duplicate registration
129        if modes.contains_key(&config.name) {
130            return Err(MarkdownConfigError::registration_error(format!(
131                "Mode '{}' is already registered",
132                config.name
133            )));
134        }
135
136        modes.insert(config.name.clone(), config);
137        Ok(())
138    }
139
140    /// Get a mode configuration by name
141    ///
142    /// # Arguments
143    /// * `name` - The name of the mode to retrieve
144    ///
145    /// # Returns
146    /// Returns the mode configuration if found, None otherwise
147    pub fn get_mode(&self, name: &str) -> MarkdownConfigResult<Option<ModeConfig>> {
148        let modes = self.modes.read().map_err(|e| {
149            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
150        })?;
151
152        Ok(modes.get(name).cloned())
153    }
154
155    /// Get all registered modes
156    pub fn get_all_modes(&self) -> MarkdownConfigResult<Vec<ModeConfig>> {
157        let modes = self.modes.read().map_err(|e| {
158            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
159        })?;
160
161        Ok(modes.values().cloned().collect())
162    }
163
164    /// Check if a mode is registered
165    pub fn has_mode(&self, name: &str) -> MarkdownConfigResult<bool> {
166        let modes = self.modes.read().map_err(|e| {
167            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
168        })?;
169
170        Ok(modes.contains_key(name))
171    }
172
173    /// Remove a mode configuration
174    pub fn remove_mode(&self, name: &str) -> MarkdownConfigResult<Option<ModeConfig>> {
175        let mut modes = self.modes.write().map_err(|e| {
176            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
177        })?;
178
179        Ok(modes.remove(name))
180    }
181
182    // ============ Command Registration ============
183
184    /// Register a command configuration
185    ///
186    /// # Arguments
187    /// * `config` - The command configuration to register
188    ///
189    /// # Errors
190    /// Returns an error if:
191    /// - A command with the same name already exists
192    /// - The configuration is invalid
193    pub fn register_command(&self, config: CommandConfig) -> MarkdownConfigResult<()> {
194        // Validate configuration
195        config.validate().map_err(|e| {
196            MarkdownConfigError::validation_error(format!("Invalid command configuration: {}", e))
197        })?;
198
199        let mut commands = self.commands.write().map_err(|e| {
200            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
201        })?;
202
203        // Check for duplicate registration
204        if commands.contains_key(&config.name) {
205            return Err(MarkdownConfigError::registration_error(format!(
206                "Command '{}' is already registered",
207                config.name
208            )));
209        }
210
211        commands.insert(config.name.clone(), config);
212        Ok(())
213    }
214
215    /// Get a command configuration by name
216    ///
217    /// # Arguments
218    /// * `name` - The name of the command to retrieve
219    ///
220    /// # Returns
221    /// Returns the command configuration if found, None otherwise
222    pub fn get_command(&self, name: &str) -> MarkdownConfigResult<Option<CommandConfig>> {
223        let commands = self.commands.read().map_err(|e| {
224            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
225        })?;
226
227        Ok(commands.get(name).cloned())
228    }
229
230    /// Get all registered commands
231    pub fn get_all_commands(&self) -> MarkdownConfigResult<Vec<CommandConfig>> {
232        let commands = self.commands.read().map_err(|e| {
233            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
234        })?;
235
236        Ok(commands.values().cloned().collect())
237    }
238
239    /// Check if a command is registered
240    pub fn has_command(&self, name: &str) -> MarkdownConfigResult<bool> {
241        let commands = self.commands.read().map_err(|e| {
242            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
243        })?;
244
245        Ok(commands.contains_key(name))
246    }
247
248    /// Remove a command configuration
249    pub fn remove_command(&self, name: &str) -> MarkdownConfigResult<Option<CommandConfig>> {
250        let mut commands = self.commands.write().map_err(|e| {
251            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
252        })?;
253
254        Ok(commands.remove(name))
255    }
256
257    // ============ Registry Management ============
258
259    /// Clear all registered configurations
260    pub fn clear(&self) -> MarkdownConfigResult<()> {
261        let mut agents = self.agents.write().map_err(|e| {
262            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
263        })?;
264        let mut modes = self.modes.write().map_err(|e| {
265            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
266        })?;
267        let mut commands = self.commands.write().map_err(|e| {
268            MarkdownConfigError::registration_error(format!("Failed to acquire write lock: {}", e))
269        })?;
270
271        agents.clear();
272        modes.clear();
273        commands.clear();
274
275        Ok(())
276    }
277
278    /// Get the total number of registered configurations
279    pub fn count(&self) -> MarkdownConfigResult<(usize, usize, usize)> {
280        let agents = self.agents.read().map_err(|e| {
281            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
282        })?;
283        let modes = self.modes.read().map_err(|e| {
284            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
285        })?;
286        let commands = self.commands.read().map_err(|e| {
287            MarkdownConfigError::registration_error(format!("Failed to acquire read lock: {}", e))
288        })?;
289
290        Ok((agents.len(), modes.len(), commands.len()))
291    }
292}
293
294impl Default for ConfigRegistry {
295    fn default() -> Self {
296        Self::new()
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303
304    fn create_test_agent(name: &str) -> AgentConfig {
305        AgentConfig {
306            name: name.to_string(),
307            description: Some("Test agent".to_string()),
308            prompt: "You are a helpful assistant".to_string(),
309            model: Some("gpt-4".to_string()),
310            temperature: Some(0.7),
311            max_tokens: Some(2000),
312            tools: vec![],
313        }
314    }
315
316    fn create_test_mode(name: &str) -> ModeConfig {
317        ModeConfig {
318            name: name.to_string(),
319            description: Some("Test mode".to_string()),
320            prompt: "Focus on the task".to_string(),
321            keybinding: Some("C-f".to_string()),
322            enabled: true,
323        }
324    }
325
326    fn create_test_command(name: &str) -> CommandConfig {
327        CommandConfig {
328            name: name.to_string(),
329            description: Some("Test command".to_string()),
330            template: "echo {{message}}".to_string(),
331            parameters: vec![],
332            keybinding: Some("C-t".to_string()),
333        }
334    }
335
336    #[test]
337    fn test_registry_creation() {
338        let registry = ConfigRegistry::new();
339        let (agents, modes, commands) = registry.count().unwrap();
340        assert_eq!(agents, 0);
341        assert_eq!(modes, 0);
342        assert_eq!(commands, 0);
343    }
344
345    #[test]
346    fn test_register_agent() {
347        let registry = ConfigRegistry::new();
348        let agent = create_test_agent("test-agent");
349
350        let result = registry.register_agent(agent.clone());
351        assert!(result.is_ok());
352
353        let retrieved = registry.get_agent("test-agent").unwrap();
354        assert_eq!(retrieved, Some(agent));
355    }
356
357    #[test]
358    fn test_register_duplicate_agent() {
359        let registry = ConfigRegistry::new();
360        let agent = create_test_agent("test-agent");
361
362        registry.register_agent(agent.clone()).unwrap();
363        let result = registry.register_agent(agent);
364        assert!(result.is_err());
365    }
366
367    #[test]
368    fn test_get_nonexistent_agent() {
369        let registry = ConfigRegistry::new();
370        let result = registry.get_agent("nonexistent").unwrap();
371        assert_eq!(result, None);
372    }
373
374    #[test]
375    fn test_has_agent() {
376        let registry = ConfigRegistry::new();
377        let agent = create_test_agent("test-agent");
378
379        registry.register_agent(agent).unwrap();
380        assert!(registry.has_agent("test-agent").unwrap());
381        assert!(!registry.has_agent("nonexistent").unwrap());
382    }
383
384    #[test]
385    fn test_remove_agent() {
386        let registry = ConfigRegistry::new();
387        let agent = create_test_agent("test-agent");
388
389        registry.register_agent(agent.clone()).unwrap();
390        let removed = registry.remove_agent("test-agent").unwrap();
391        assert_eq!(removed, Some(agent));
392        assert!(!registry.has_agent("test-agent").unwrap());
393    }
394
395    #[test]
396    fn test_get_all_agents() {
397        let registry = ConfigRegistry::new();
398        let agent1 = create_test_agent("agent1");
399        let agent2 = create_test_agent("agent2");
400
401        registry.register_agent(agent1).unwrap();
402        registry.register_agent(agent2).unwrap();
403
404        let all = registry.get_all_agents().unwrap();
405        assert_eq!(all.len(), 2);
406    }
407
408    #[test]
409    fn test_register_mode() {
410        let registry = ConfigRegistry::new();
411        let mode = create_test_mode("test-mode");
412
413        let result = registry.register_mode(mode.clone());
414        assert!(result.is_ok());
415
416        let retrieved = registry.get_mode("test-mode").unwrap();
417        assert_eq!(retrieved, Some(mode));
418    }
419
420    #[test]
421    fn test_register_duplicate_mode() {
422        let registry = ConfigRegistry::new();
423        let mode = create_test_mode("test-mode");
424
425        registry.register_mode(mode.clone()).unwrap();
426        let result = registry.register_mode(mode);
427        assert!(result.is_err());
428    }
429
430    #[test]
431    fn test_register_command() {
432        let registry = ConfigRegistry::new();
433        let command = create_test_command("test-command");
434
435        let result = registry.register_command(command.clone());
436        assert!(result.is_ok());
437
438        let retrieved = registry.get_command("test-command").unwrap();
439        assert_eq!(retrieved, Some(command));
440    }
441
442    #[test]
443    fn test_register_duplicate_command() {
444        let registry = ConfigRegistry::new();
445        let command = create_test_command("test-command");
446
447        registry.register_command(command.clone()).unwrap();
448        let result = registry.register_command(command);
449        assert!(result.is_err());
450    }
451
452    #[test]
453    fn test_clear_registry() {
454        let registry = ConfigRegistry::new();
455        registry.register_agent(create_test_agent("agent1")).unwrap();
456        registry.register_mode(create_test_mode("mode1")).unwrap();
457        registry.register_command(create_test_command("command1")).unwrap();
458
459        registry.clear().unwrap();
460
461        let (agents, modes, commands) = registry.count().unwrap();
462        assert_eq!(agents, 0);
463        assert_eq!(modes, 0);
464        assert_eq!(commands, 0);
465    }
466
467    #[test]
468    fn test_count() {
469        let registry = ConfigRegistry::new();
470        registry.register_agent(create_test_agent("agent1")).unwrap();
471        registry.register_agent(create_test_agent("agent2")).unwrap();
472        registry.register_mode(create_test_mode("mode1")).unwrap();
473        registry.register_command(create_test_command("command1")).unwrap();
474
475        let (agents, modes, commands) = registry.count().unwrap();
476        assert_eq!(agents, 2);
477        assert_eq!(modes, 1);
478        assert_eq!(commands, 1);
479    }
480}