ricecoder_storage/markdown_config/
registry.rs1use crate::markdown_config::error::{MarkdownConfigError, MarkdownConfigResult};
4use crate::markdown_config::types::{AgentConfig, CommandConfig, ModeConfig};
5use std::collections::HashMap;
6use std::sync::RwLock;
7
8#[derive(Debug)]
13pub struct ConfigRegistry {
14 agents: RwLock<HashMap<String, AgentConfig>>,
16 modes: RwLock<HashMap<String, ModeConfig>>,
18 commands: RwLock<HashMap<String, CommandConfig>>,
20}
21
22impl ConfigRegistry {
23 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 pub fn register_agent(&self, config: AgentConfig) -> MarkdownConfigResult<()> {
44 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 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 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 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 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 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 pub fn register_mode(&self, config: ModeConfig) -> MarkdownConfigResult<()> {
119 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 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 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 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 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 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 pub fn register_command(&self, config: CommandConfig) -> MarkdownConfigResult<()> {
194 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 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 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 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 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 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 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 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}