telos_agent/orchestration/subagent/
registry.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3
4use crate::error::AgentError;
5use crate::orchestration::subagent::builtins::builtin_agents;
6use crate::orchestration::subagent::definition::{AgentDefinition, AgentSource};
7
8/// Registry of available subagent definitions.
9#[derive(Debug, Clone, Default)]
10pub struct SubagentRegistry {
11    agents: BTreeMap<String, AgentDefinition>,
12}
13
14impl SubagentRegistry {
15    pub fn new() -> Self {
16        Self::default()
17    }
18
19    pub fn with_builtin_agents() -> Self {
20        let mut registry = Self::new();
21        for agent in builtin_agents() {
22            registry.register(agent);
23        }
24        registry
25    }
26
27    pub fn register(&mut self, definition: AgentDefinition) {
28        self.agents.insert(definition.name.clone(), definition);
29    }
30
31    pub fn get(&self, name: &str) -> Option<&AgentDefinition> {
32        self.agents.get(name)
33    }
34
35    pub fn definitions(&self) -> Vec<&AgentDefinition> {
36        self.agents.values().collect()
37    }
38
39    pub fn render_listing(&self) -> String {
40        self.agents
41            .values()
42            .map(|agent| {
43                let tools = if agent.allowed_tools.is_empty() {
44                    "All tools".to_string()
45                } else {
46                    agent.allowed_tools.join(", ")
47                };
48                format!("- {}: {} (Tools: {tools})", agent.name, agent.description)
49            })
50            .collect::<Vec<_>>()
51            .join("\n")
52    }
53
54    pub fn load_markdown_file(
55        &mut self,
56        path: impl AsRef<Path>,
57        source: AgentSource,
58    ) -> Result<(), AgentError> {
59        let path = path.as_ref();
60        let content = std::fs::read_to_string(path).map_err(|err| {
61            AgentError::Config(format!("failed to read agent file {}: {err}", path.display()))
62        })?;
63        let definition = AgentDefinition::from_markdown(&content, source)?;
64        self.register(definition);
65        Ok(())
66    }
67
68    pub fn load_markdown_dir(
69        &mut self,
70        dir: impl AsRef<Path>,
71        source: AgentSource,
72    ) -> Result<(), AgentError> {
73        let dir = dir.as_ref();
74        let entries = std::fs::read_dir(dir).map_err(|err| {
75            AgentError::Config(format!("failed to read agent dir {}: {err}", dir.display()))
76        })?;
77        let mut paths = Vec::new();
78        for entry in entries {
79            let entry = entry.map_err(|err| {
80                AgentError::Config(format!("failed to read agent dir entry: {err}"))
81            })?;
82            let path = entry.path();
83            if path.extension().and_then(|ext| ext.to_str()).is_some_and(|ext| ext == "md") {
84                paths.push(path);
85            }
86        }
87        paths.sort();
88        for path in paths {
89            self.load_markdown_file(&path, source.clone())?;
90        }
91        Ok(())
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::orchestration::subagent::{AgentDefinition, AgentSource};
99
100    #[test]
101    fn later_agent_registration_overrides_earlier_definition() {
102        let mut registry = SubagentRegistry::new();
103        registry.register(AgentDefinition::new(
104            "reviewer",
105            "old",
106            "old prompt",
107            AgentSource::BuiltIn,
108        ));
109        registry.register(AgentDefinition::new(
110            "reviewer",
111            "new",
112            "new prompt",
113            AgentSource::Project { path: "agents/reviewer.md".into() },
114        ));
115
116        let definition = registry.get("reviewer").unwrap();
117        assert_eq!(definition.description, "new");
118        assert_eq!(registry.definitions().len(), 1);
119    }
120
121    #[test]
122    fn builtins_include_core_agent_types() {
123        let registry = SubagentRegistry::with_builtin_agents();
124        assert!(registry.get("general-purpose").is_some());
125        assert!(registry.get("Explore").is_some());
126        assert!(registry.get("Plan").is_some());
127        assert!(registry.get("Verification").is_some());
128        assert!(registry.render_listing().contains("Explore"));
129    }
130
131    #[test]
132    fn loads_agent_markdown_files_from_directory() {
133        let dir = std::env::temp_dir().join("tiny_agent_subagent_registry_test");
134        let _ = std::fs::remove_dir_all(&dir);
135        std::fs::create_dir_all(&dir).unwrap();
136        std::fs::write(
137            dir.join("auditor.md"),
138            r#"---
139name: auditor
140description: Audit changes.
141tools: [Read]
142---
143You audit changes.
144"#,
145        )
146        .unwrap();
147        std::fs::write(dir.join("notes.txt"), "ignored").unwrap();
148
149        let mut registry = SubagentRegistry::new();
150        registry
151            .load_markdown_dir(&dir, AgentSource::Project { path: dir.display().to_string() })
152            .unwrap();
153
154        let auditor = registry.get("auditor").unwrap();
155        assert_eq!(auditor.system_prompt, "You audit changes.");
156
157        let _ = std::fs::remove_dir_all(&dir);
158    }
159}
telos_agent/orchestration/subagent/registry.rs

telos_agent/orchestration/subagent/
registry.rs
