agcodex_core/subagents/
registry.rs

1//! Subagent registry for loading and managing agent configurations
2//!
3//! The registry loads agent configurations from TOML files and provides
4//! hot-reload capabilities for dynamic agent management.
5
6use super::config::SubagentConfig;
7use super::config::SubagentTemplate;
8use std::collections::HashMap;
9use std::path::Path;
10use std::path::PathBuf;
11use std::sync::Arc;
12use std::sync::Mutex;
13use std::time::SystemTime;
14use thiserror::Error;
15use walkdir::WalkDir;
16
17/// Result type for registry operations
18pub type RegistryResult<T> = std::result::Result<T, SubagentRegistryError>;
19
20/// Errors specific to the subagent registry
21#[derive(Error, Debug)]
22pub enum SubagentRegistryError {
23    #[error("agent configuration file not found: {path}")]
24    ConfigNotFound { path: PathBuf },
25
26    #[error("template not found: {name}")]
27    TemplateNotFound { name: String },
28
29    #[error("invalid agent directory: {path}")]
30    InvalidDirectory { path: PathBuf },
31
32    #[error("failed to load configuration: {path}: {error}")]
33    LoadError { path: PathBuf, error: String },
34
35    #[error("agent name conflict: {name} (paths: {path1}, {path2})")]
36    NameConflict {
37        name: String,
38        path1: PathBuf,
39        path2: PathBuf,
40    },
41
42    #[error("template inheritance loop detected: {chain:?}")]
43    InheritanceLoop { chain: Vec<String> },
44
45    #[error("I/O error: {0}")]
46    Io(#[from] std::io::Error),
47
48    #[error("configuration error: {0}")]
49    Config(#[from] super::SubagentError),
50}
51
52/// Information about a loaded agent configuration
53#[derive(Debug, Clone)]
54pub struct AgentInfo {
55    /// The agent configuration
56    pub config: SubagentConfig,
57
58    /// Path to the configuration file
59    pub config_path: PathBuf,
60
61    /// Last modification time of the configuration file
62    pub last_modified: SystemTime,
63
64    /// Whether this is a global or project-specific agent
65    pub is_global: bool,
66}
67
68/// Information about a loaded template
69#[derive(Debug, Clone)]
70pub struct TemplateInfo {
71    /// The template configuration
72    pub template: SubagentTemplate,
73
74    /// Path to the template file
75    pub template_path: PathBuf,
76
77    /// Last modification time of the template file
78    pub last_modified: SystemTime,
79}
80
81/// Registry for managing subagent configurations
82#[derive(Debug)]
83pub struct SubagentRegistry {
84    /// Global agents directory (~/.agcodex/agents/global/)
85    global_agents_dir: PathBuf,
86
87    /// Project-specific agents directory (./.agcodex/agents/)
88    project_agents_dir: Option<PathBuf>,
89
90    /// Templates directory (~/.agcodex/agents/templates/)
91    templates_dir: PathBuf,
92
93    /// Loaded agent configurations
94    agents: Arc<Mutex<HashMap<String, AgentInfo>>>,
95
96    /// Executable agents (for built-in and programmatic agents)
97    executable_agents: Arc<Mutex<HashMap<String, Arc<dyn super::agents::Subagent>>>>,
98
99    /// Loaded templates
100    templates: Arc<Mutex<HashMap<String, TemplateInfo>>>,
101
102    /// Whether to watch for file changes
103    _watch_enabled: bool,
104
105    /// Last full scan time
106    last_scan: Arc<Mutex<SystemTime>>,
107}
108
109impl SubagentRegistry {
110    /// Create a new subagent registry
111    pub fn new() -> RegistryResult<Self> {
112        let home_dir = dirs::home_dir().ok_or_else(|| SubagentRegistryError::InvalidDirectory {
113            path: PathBuf::from("~"),
114        })?;
115
116        let global_agents_dir = home_dir.join(".agcodex").join("agents").join("global");
117        let templates_dir = home_dir.join(".agcodex").join("agents").join("templates");
118
119        // Try to find project-specific agents directory
120        let project_agents_dir = Self::find_project_agents_dir()?;
121
122        let registry = Self {
123            global_agents_dir,
124            project_agents_dir,
125            templates_dir,
126            agents: Arc::new(Mutex::new(HashMap::new())),
127            executable_agents: Arc::new(Mutex::new(HashMap::new())),
128            templates: Arc::new(Mutex::new(HashMap::new())),
129            _watch_enabled: true,
130            last_scan: Arc::new(Mutex::new(SystemTime::UNIX_EPOCH)),
131        };
132
133        // Create directories if they don't exist
134        registry.ensure_directories()?;
135
136        Ok(registry)
137    }
138
139    /// Find the project-specific agents directory by walking up from current directory
140    fn find_project_agents_dir() -> RegistryResult<Option<PathBuf>> {
141        let current_dir = std::env::current_dir()?;
142
143        for ancestor in current_dir.ancestors() {
144            let agents_dir = ancestor.join(".agcodex").join("agents");
145            if agents_dir.exists() && agents_dir.is_dir() {
146                return Ok(Some(agents_dir));
147            }
148        }
149
150        Ok(None)
151    }
152
153    /// Ensure that all required directories exist
154    fn ensure_directories(&self) -> RegistryResult<()> {
155        std::fs::create_dir_all(&self.global_agents_dir)?;
156        std::fs::create_dir_all(&self.templates_dir)?;
157
158        if let Some(ref project_dir) = self.project_agents_dir {
159            std::fs::create_dir_all(project_dir)?;
160        }
161
162        Ok(())
163    }
164
165    /// Load all agent configurations and templates
166    pub fn load_all(&self) -> RegistryResult<()> {
167        self.load_templates()?;
168        self.load_agents()?;
169
170        *self.last_scan.lock().unwrap() = SystemTime::now();
171
172        Ok(())
173    }
174
175    /// Load all templates from the templates directory
176    fn load_templates(&self) -> RegistryResult<()> {
177        let mut templates = self.templates.lock().unwrap();
178        templates.clear();
179
180        if !self.templates_dir.exists() {
181            return Ok(());
182        }
183
184        for entry in WalkDir::new(&self.templates_dir)
185            .follow_links(true)
186            .into_iter()
187            .filter_map(|e| e.ok())
188            .filter(|e| e.file_type().is_file())
189            .filter(|e| {
190                e.path()
191                    .extension()
192                    .map(|ext| ext == "toml")
193                    .unwrap_or(false)
194            })
195        {
196            let path = entry.path();
197
198            match self.load_template_from_file(path) {
199                Ok(template_info) => {
200                    let name = template_info.template.name.clone();
201                    templates.insert(name, template_info);
202                }
203                Err(e) => {
204                    tracing::warn!("Failed to load template from {}: {}", path.display(), e);
205                }
206            }
207        }
208
209        Ok(())
210    }
211
212    /// Load a single template from a file
213    fn load_template_from_file(&self, path: &Path) -> RegistryResult<TemplateInfo> {
214        let metadata = std::fs::metadata(path)?;
215        let last_modified = metadata.modified()?;
216
217        let template = SubagentTemplate::from_file(&path.to_path_buf()).map_err(|e| {
218            SubagentRegistryError::LoadError {
219                path: path.to_path_buf(),
220                error: e.to_string(),
221            }
222        })?;
223
224        Ok(TemplateInfo {
225            template,
226            template_path: path.to_path_buf(),
227            last_modified,
228        })
229    }
230
231    /// Load all agent configurations
232    fn load_agents(&self) -> RegistryResult<()> {
233        let mut agents = self.agents.lock().unwrap();
234        agents.clear();
235
236        // Load global agents
237        self.load_agents_from_directory(&self.global_agents_dir, true, &mut agents)?;
238
239        // Load project-specific agents
240        if let Some(ref project_dir) = self.project_agents_dir {
241            self.load_agents_from_directory(project_dir, false, &mut agents)?;
242        }
243
244        // Resolve template inheritance
245        self.resolve_template_inheritance(&mut agents)?;
246
247        Ok(())
248    }
249
250    /// Load agents from a specific directory
251    fn load_agents_from_directory(
252        &self,
253        dir: &Path,
254        is_global: bool,
255        agents: &mut HashMap<String, AgentInfo>,
256    ) -> RegistryResult<()> {
257        if !dir.exists() {
258            return Ok(());
259        }
260
261        for entry in WalkDir::new(dir)
262            .follow_links(true)
263            .into_iter()
264            .filter_map(|e| e.ok())
265            .filter(|e| e.file_type().is_file())
266            .filter(|e| {
267                e.path()
268                    .extension()
269                    .map(|ext| ext == "toml")
270                    .unwrap_or(false)
271            })
272        {
273            let path = entry.path();
274
275            match self.load_agent_from_file(path, is_global) {
276                Ok(agent_info) => {
277                    let name = agent_info.config.name.clone();
278
279                    // Check for name conflicts
280                    if let Some(existing) = agents.get(&name) {
281                        return Err(SubagentRegistryError::NameConflict {
282                            name,
283                            path1: existing.config_path.clone(),
284                            path2: path.to_path_buf(),
285                        });
286                    }
287
288                    agents.insert(name, agent_info);
289                }
290                Err(e) => {
291                    tracing::warn!("Failed to load agent from {}: {}", path.display(), e);
292                }
293            }
294        }
295
296        Ok(())
297    }
298
299    /// Load a single agent configuration from a file
300    fn load_agent_from_file(&self, path: &Path, is_global: bool) -> RegistryResult<AgentInfo> {
301        let metadata = std::fs::metadata(path)?;
302        let last_modified = metadata.modified()?;
303
304        let config = SubagentConfig::from_file(&path.to_path_buf()).map_err(|e| {
305            SubagentRegistryError::LoadError {
306                path: path.to_path_buf(),
307                error: e.to_string(),
308            }
309        })?;
310
311        Ok(AgentInfo {
312            config,
313            config_path: path.to_path_buf(),
314            last_modified,
315            is_global,
316        })
317    }
318
319    /// Resolve template inheritance for all agents
320    fn resolve_template_inheritance(
321        &self,
322        agents: &mut HashMap<String, AgentInfo>,
323    ) -> RegistryResult<()> {
324        let templates = self.templates.lock().unwrap();
325
326        for agent_info in agents.values_mut() {
327            if let Some(ref template_name) = agent_info.config.template {
328                let mut inheritance_chain = Vec::new();
329                let mut current_template = template_name.clone();
330
331                // Follow the inheritance chain
332                loop {
333                    if inheritance_chain.contains(&current_template) {
334                        return Err(SubagentRegistryError::InheritanceLoop {
335                            chain: inheritance_chain,
336                        });
337                    }
338
339                    inheritance_chain.push(current_template.clone());
340
341                    let template = templates.get(&current_template).ok_or_else(|| {
342                        SubagentRegistryError::TemplateNotFound {
343                            name: current_template.clone(),
344                        }
345                    })?;
346
347                    // Apply template to the agent config
348                    self.apply_template_to_config(&template.template, &mut agent_info.config)?;
349
350                    // Check if this template inherits from another
351                    if let Some(ref parent_template) = template.template.config.template {
352                        current_template = parent_template.clone();
353                    } else {
354                        break;
355                    }
356                }
357            }
358        }
359
360        Ok(())
361    }
362
363    /// Apply a template to an agent configuration
364    fn apply_template_to_config(
365        &self,
366        template: &SubagentTemplate,
367        config: &mut SubagentConfig,
368    ) -> RegistryResult<()> {
369        // Merge template configuration with agent configuration
370        // Agent-specific values take precedence
371
372        if config.description.is_empty() {
373            config.description = template.config.description.clone();
374        }
375
376        if config.mode_override.is_none() {
377            config.mode_override = template.config.mode_override;
378        }
379
380        if config.prompt.is_empty() {
381            config.prompt = template.config.prompt.clone();
382        }
383
384        // Merge tools (agent tools override template tools)
385        for (tool, permission) in &template.config.tools {
386            config
387                .tools
388                .entry(tool.clone())
389                .or_insert_with(|| permission.clone());
390        }
391
392        // Merge parameters (template parameters as defaults)
393        for template_param in &template.config.parameters {
394            if !config
395                .parameters
396                .iter()
397                .any(|p| p.name == template_param.name)
398            {
399                config.parameters.push(template_param.clone());
400            }
401        }
402
403        // Merge metadata
404        for (key, value) in &template.config.metadata {
405            config
406                .metadata
407                .entry(key.clone())
408                .or_insert_with(|| value.clone());
409        }
410
411        // Merge file patterns
412        for pattern in &template.config.file_patterns {
413            if !config.file_patterns.contains(pattern) {
414                config.file_patterns.push(pattern.clone());
415            }
416        }
417
418        // Merge tags
419        for tag in &template.config.tags {
420            if !config.tags.contains(tag) {
421                config.tags.push(tag.clone());
422            }
423        }
424
425        Ok(())
426    }
427
428    /// Register an agent programmatically (for built-in agents)
429    pub fn register(&self, name: &str, agent: Arc<dyn super::agents::Subagent>) {
430        // Create a minimal SubagentConfig for built-in agents
431        let config = SubagentConfig {
432            name: name.to_string(),
433            description: format!("Built-in {} agent", name),
434            mode_override: None,
435            intelligence: super::config::IntelligenceLevel::Medium,
436            tools: HashMap::new(),
437            prompt: String::new(),
438            parameters: vec![],
439            template: None,
440            timeout_seconds: 300,
441            chainable: true,
442            parallelizable: true,
443            metadata: HashMap::new(),
444            file_patterns: vec![],
445            tags: vec![],
446        };
447
448        let info = AgentInfo {
449            config,
450            config_path: PathBuf::new(),
451            last_modified: SystemTime::now(),
452            is_global: false,
453        };
454
455        // Store both configuration and executable agent
456        let mut agents = self.agents.lock().unwrap();
457        agents.insert(name.to_string(), info);
458        drop(agents);
459
460        let mut executable_agents = self.executable_agents.lock().unwrap();
461        executable_agents.insert(name.to_string(), agent);
462    }
463
464    /// Get an agent configuration by name
465    pub fn get_agent(&self, name: &str) -> Option<AgentInfo> {
466        self.agents.lock().unwrap().get(name).cloned()
467    }
468
469    /// Get an executable agent by name
470    pub fn get_executable_agent(&self, name: &str) -> Option<Arc<dyn super::agents::Subagent>> {
471        self.executable_agents.lock().unwrap().get(name).cloned()
472    }
473
474    /// Get all loaded agents
475    pub fn get_all_agents(&self) -> HashMap<String, AgentInfo> {
476        self.agents.lock().unwrap().clone()
477    }
478
479    /// Get a template by name
480    pub fn get_template(&self, name: &str) -> Option<TemplateInfo> {
481        self.templates.lock().unwrap().get(name).cloned()
482    }
483
484    /// Get all loaded templates
485    pub fn get_all_templates(&self) -> HashMap<String, TemplateInfo> {
486        self.templates.lock().unwrap().clone()
487    }
488
489    /// Check if any configurations have been modified and reload if necessary
490    pub fn check_for_updates(&self) -> RegistryResult<bool> {
491        let mut updated = false;
492
493        // Check templates
494        {
495            let templates = self.templates.lock().unwrap();
496            for template_info in templates.values() {
497                if let Ok(metadata) = std::fs::metadata(&template_info.template_path)
498                    && let Ok(modified) = metadata.modified()
499                    && modified > template_info.last_modified
500                {
501                    updated = true;
502                    break;
503                }
504            }
505        }
506
507        // Check agents
508        if !updated {
509            let agents = self.agents.lock().unwrap();
510            for agent_info in agents.values() {
511                if let Ok(metadata) = std::fs::metadata(&agent_info.config_path)
512                    && let Ok(modified) = metadata.modified()
513                    && modified > agent_info.last_modified
514                {
515                    updated = true;
516                    break;
517                }
518            }
519        }
520
521        // Reload if updates detected
522        if updated {
523            self.load_all()?;
524        }
525
526        Ok(updated)
527    }
528
529    /// Get agents that match a specific file pattern
530    pub fn get_agents_for_file(&self, file_path: &Path) -> Vec<AgentInfo> {
531        self.agents
532            .lock()
533            .unwrap()
534            .values()
535            .filter(|agent| agent.config.matches_file(file_path))
536            .cloned()
537            .collect()
538    }
539
540    /// Get agents with specific tags
541    pub fn get_agents_with_tags(&self, tags: &[String]) -> Vec<AgentInfo> {
542        self.agents
543            .lock()
544            .unwrap()
545            .values()
546            .filter(|agent| tags.iter().any(|tag| agent.config.tags.contains(tag)))
547            .cloned()
548            .collect()
549    }
550
551    /// Create default agent configurations
552    pub fn create_default_agents(&self) -> RegistryResult<()> {
553        self.ensure_directories()?;
554
555        // Create default code reviewer agent
556        let code_reviewer = SubagentConfig {
557            name: "code-reviewer".to_string(),
558            description: "Proactive code quality analysis and security review".to_string(),
559            mode_override: Some(crate::modes::OperatingMode::Review),
560            intelligence: crate::subagents::IntelligenceLevel::Hard,
561            tools: std::collections::HashMap::new(),
562            prompt: r#"You are a senior code reviewer with AST-based analysis capabilities.
563
564Focus on:
565- Syntactic correctness via tree-sitter validation
566- Security vulnerabilities (OWASP Top 10)
567- Performance bottlenecks (O(n²) or worse)
568- Memory leaks and resource management
569- Error handling completeness
570- Code quality and maintainability
571
572Use AST-powered semantic search to understand code structure and relationships."#
573                .to_string(),
574            parameters: vec![
575                super::config::ParameterDefinition {
576                    name: "files".to_string(),
577                    description: "Files or patterns to review".to_string(),
578                    required: false,
579                    default: Some("**/*.rs".to_string()),
580                    valid_values: None,
581                },
582                super::config::ParameterDefinition {
583                    name: "focus".to_string(),
584                    description: "Focus area (security, performance, quality)".to_string(),
585                    required: false,
586                    default: Some("quality".to_string()),
587                    valid_values: Some(vec![
588                        "security".to_string(),
589                        "performance".to_string(),
590                        "quality".to_string(),
591                    ]),
592                },
593            ],
594            template: None,
595            timeout_seconds: 600,
596            chainable: true,
597            parallelizable: true,
598            metadata: std::collections::HashMap::new(),
599            file_patterns: vec![
600                "*.rs".to_string(),
601                "*.py".to_string(),
602                "*.js".to_string(),
603                "*.ts".to_string(),
604            ],
605            tags: vec![
606                "review".to_string(),
607                "quality".to_string(),
608                "security".to_string(),
609            ],
610        };
611
612        let config_path = self.global_agents_dir.join("code-reviewer.toml");
613        if !config_path.exists() {
614            code_reviewer.to_file(&config_path)?;
615        }
616
617        Ok(())
618    }
619}
620
621impl Default for SubagentRegistry {
622    fn default() -> Self {
623        Self::new().expect("Failed to create default subagent registry")
624    }
625}
626
627#[cfg(test)]
628mod tests {
629    use super::*;
630    use tempfile::TempDir;
631
632    fn create_test_registry() -> (SubagentRegistry, TempDir) {
633        let temp_dir = TempDir::new().unwrap();
634        let global_dir = temp_dir.path().join("global");
635        let templates_dir = temp_dir.path().join("templates");
636
637        std::fs::create_dir_all(&global_dir).unwrap();
638        std::fs::create_dir_all(&templates_dir).unwrap();
639
640        let registry = SubagentRegistry {
641            global_agents_dir: global_dir,
642            project_agents_dir: None,
643            templates_dir,
644            agents: Arc::new(Mutex::new(HashMap::new())),
645            executable_agents: Arc::new(Mutex::new(HashMap::new())),
646            templates: Arc::new(Mutex::new(HashMap::new())),
647            _watch_enabled: false,
648            last_scan: Arc::new(Mutex::new(SystemTime::UNIX_EPOCH)),
649        };
650
651        (registry, temp_dir)
652    }
653
654    #[test]
655    fn test_registry_creation() {
656        let (registry, _temp_dir) = create_test_registry();
657        assert!(registry.global_agents_dir.exists());
658        assert!(registry.templates_dir.exists());
659    }
660
661    #[test]
662    fn test_agent_loading() {
663        let (registry, _temp_dir) = create_test_registry();
664
665        // Create a test agent configuration
666        let config = SubagentConfig {
667            name: "test-agent".to_string(),
668            description: "Test agent".to_string(),
669            mode_override: None,
670            intelligence: crate::subagents::IntelligenceLevel::Medium,
671            tools: HashMap::new(),
672            prompt: "You are a test agent.".to_string(),
673            parameters: vec![],
674            template: None,
675            timeout_seconds: 300,
676            chainable: true,
677            parallelizable: true,
678            metadata: HashMap::new(),
679            file_patterns: vec![],
680            tags: vec![],
681        };
682
683        let config_path = registry.global_agents_dir.join("test-agent.toml");
684        config.to_file(&config_path).unwrap();
685
686        // Load agents
687        registry.load_all().unwrap();
688
689        // Verify agent was loaded
690        let loaded_agent = registry.get_agent("test-agent").unwrap();
691        assert_eq!(loaded_agent.config.name, "test-agent");
692        assert!(loaded_agent.is_global);
693    }
694
695    #[test]
696    fn test_template_inheritance() {
697        let (registry, _temp_dir) = create_test_registry();
698
699        // Create a template
700        let template = SubagentTemplate {
701            name: "base-reviewer".to_string(),
702            description: "Base template for reviewers".to_string(),
703            config: SubagentConfig {
704                name: "template".to_string(),
705                description: "Template description".to_string(),
706                mode_override: Some(crate::modes::OperatingMode::Review),
707                intelligence: crate::subagents::IntelligenceLevel::Hard,
708                tools: HashMap::new(),
709                prompt: "You are a reviewer.".to_string(),
710                parameters: vec![],
711                template: None,
712                timeout_seconds: 300,
713                chainable: true,
714                parallelizable: true,
715                metadata: HashMap::new(),
716                file_patterns: vec!["*.rs".to_string()],
717                tags: vec!["review".to_string()],
718            },
719            placeholders: vec![],
720        };
721
722        let template_path = registry.templates_dir.join("base-reviewer.toml");
723        let template_content = toml::to_string(&template).unwrap();
724        std::fs::write(&template_path, template_content).unwrap();
725
726        // Create an agent that inherits from the template
727        let agent_config = SubagentConfig {
728            name: "code-reviewer".to_string(),
729            description: "".to_string(), // Will inherit from template
730            mode_override: None,
731            intelligence: crate::subagents::IntelligenceLevel::Medium,
732            tools: HashMap::new(),
733            prompt: "".to_string(), // Will inherit from template
734            parameters: vec![],
735            template: Some("base-reviewer".to_string()),
736            timeout_seconds: 300,
737            chainable: true,
738            parallelizable: true,
739            metadata: HashMap::new(),
740            file_patterns: vec![],
741            tags: vec![],
742        };
743
744        let config_path = registry.global_agents_dir.join("code-reviewer.toml");
745        agent_config.to_file(&config_path).unwrap();
746
747        // Load all configurations
748        registry.load_all().unwrap();
749
750        // Verify inheritance was applied
751        let loaded_agent = registry.get_agent("code-reviewer").unwrap();
752        assert_eq!(loaded_agent.config.description, "Template description");
753        assert_eq!(loaded_agent.config.prompt, "You are a reviewer.");
754        assert_eq!(
755            loaded_agent.config.mode_override,
756            Some(crate::modes::OperatingMode::Review)
757        );
758        assert!(
759            loaded_agent
760                .config
761                .file_patterns
762                .contains(&"*.rs".to_string())
763        );
764        assert!(loaded_agent.config.tags.contains(&"review".to_string()));
765    }
766}