1use crate::config::CodeConfig;
62use crate::permissions::PermissionPolicy;
63use serde::{Deserialize, Serialize};
64use std::collections::HashMap;
65use std::path::Path;
66use std::sync::RwLock;
67
68use crate::error::{read_or_recover, write_or_recover};
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
72#[serde(rename_all = "snake_case")]
73pub enum AgentMode {
74 #[default]
76 Primary,
77 Subagent,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct ModelConfig {
84 pub model: String,
86 pub provider: Option<String>,
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct AgentDefinition {
95 pub name: String,
97 pub description: String,
99 #[serde(default)]
101 pub mode: AgentMode,
102 #[serde(default)]
104 pub native: bool,
105 #[serde(default)]
107 pub hidden: bool,
108 #[serde(default)]
110 pub permissions: PermissionPolicy,
111 #[serde(skip_serializing_if = "Option::is_none")]
113 pub model: Option<ModelConfig>,
114 #[serde(skip_serializing_if = "Option::is_none")]
116 pub prompt: Option<String>,
117 #[serde(skip_serializing_if = "Option::is_none")]
119 pub max_steps: Option<usize>,
120 #[serde(default)]
122 pub can_spawn_subagents: bool,
123}
124
125impl AgentDefinition {
126 pub fn new(name: &str, description: &str) -> Self {
128 Self {
129 name: name.to_string(),
130 description: description.to_string(),
131 mode: AgentMode::Subagent,
132 native: false,
133 hidden: false,
134 permissions: PermissionPolicy::default(),
135 model: None,
136 prompt: None,
137 max_steps: None,
138 can_spawn_subagents: false,
139 }
140 }
141
142 pub fn with_mode(mut self, mode: AgentMode) -> Self {
144 self.mode = mode;
145 self
146 }
147
148 pub fn native(mut self) -> Self {
150 self.native = true;
151 self
152 }
153
154 pub fn hidden(mut self) -> Self {
156 self.hidden = true;
157 self
158 }
159
160 pub fn with_permissions(mut self, permissions: PermissionPolicy) -> Self {
162 self.permissions = permissions;
163 self
164 }
165
166 pub fn with_model(mut self, model: ModelConfig) -> Self {
168 self.model = Some(model);
169 self
170 }
171
172 pub fn with_prompt(mut self, prompt: &str) -> Self {
174 self.prompt = Some(prompt.to_string());
175 self
176 }
177
178 pub fn with_max_steps(mut self, max_steps: usize) -> Self {
180 self.max_steps = Some(max_steps);
181 self
182 }
183
184 pub fn allow_subagents(mut self) -> Self {
186 self.can_spawn_subagents = true;
187 self
188 }
189}
190
191pub struct AgentRegistry {
196 agents: RwLock<HashMap<String, AgentDefinition>>,
197}
198
199impl Default for AgentRegistry {
200 fn default() -> Self {
201 Self::new()
202 }
203}
204
205impl AgentRegistry {
206 pub fn new() -> Self {
208 let registry = Self {
209 agents: RwLock::new(HashMap::new()),
210 };
211
212 for agent in builtin_agents() {
214 registry.register(agent);
215 }
216
217 registry
218 }
219
220 pub fn with_config(config: &CodeConfig) -> Self {
224 let registry = Self::new();
225
226 for dir in &config.agent_dirs {
228 let agents = load_agents_from_dir(dir);
229 for agent in agents {
230 tracing::info!("Loaded agent '{}' from {}", agent.name, dir.display());
231 registry.register(agent);
232 }
233 }
234
235 registry
236 }
237
238 pub fn register(&self, agent: AgentDefinition) {
240 let mut agents = write_or_recover(&self.agents);
241 tracing::debug!("Registering agent: {}", agent.name);
242 agents.insert(agent.name.clone(), agent);
243 }
244
245 pub fn unregister(&self, name: &str) -> bool {
249 let mut agents = write_or_recover(&self.agents);
250 agents.remove(name).is_some()
251 }
252
253 pub fn get(&self, name: &str) -> Option<AgentDefinition> {
255 let agents = read_or_recover(&self.agents);
256 agents.get(name).cloned()
257 }
258
259 pub fn list(&self) -> Vec<AgentDefinition> {
261 let agents = read_or_recover(&self.agents);
262 agents.values().cloned().collect()
263 }
264
265 pub fn list_visible(&self) -> Vec<AgentDefinition> {
267 let agents = read_or_recover(&self.agents);
268 agents.values().filter(|a| !a.hidden).cloned().collect()
269 }
270
271 pub fn exists(&self, name: &str) -> bool {
273 let agents = read_or_recover(&self.agents);
274 agents.contains_key(name)
275 }
276
277 pub fn len(&self) -> usize {
279 let agents = read_or_recover(&self.agents);
280 agents.len()
281 }
282
283 pub fn is_empty(&self) -> bool {
285 self.len() == 0
286 }
287}
288
289pub fn parse_agent_yaml(content: &str) -> anyhow::Result<AgentDefinition> {
297 let agent: AgentDefinition = serde_yaml::from_str(content)
298 .map_err(|e| anyhow::anyhow!("Failed to parse agent YAML: {}", e))?;
299
300 if agent.name.is_empty() {
301 return Err(anyhow::anyhow!("Agent name is required"));
302 }
303
304 Ok(agent)
305}
306
307pub fn parse_agent_md(content: &str) -> anyhow::Result<AgentDefinition> {
311 let parts: Vec<&str> = content.splitn(3, "---").collect();
313
314 if parts.len() < 3 {
315 return Err(anyhow::anyhow!(
316 "Invalid markdown format: missing YAML frontmatter"
317 ));
318 }
319
320 let frontmatter = parts[1].trim();
321 let body = parts[2].trim();
322
323 let mut agent: AgentDefinition = serde_yaml::from_str(frontmatter)
325 .map_err(|e| anyhow::anyhow!("Failed to parse agent frontmatter: {}", e))?;
326
327 if agent.name.is_empty() {
328 return Err(anyhow::anyhow!("Agent name is required"));
329 }
330
331 if agent.prompt.is_none() && !body.is_empty() {
333 agent.prompt = Some(body.to_string());
334 }
335
336 Ok(agent)
337}
338
339pub fn load_agents_from_dir(dir: &Path) -> Vec<AgentDefinition> {
344 let mut agents = Vec::new();
345
346 let Ok(entries) = std::fs::read_dir(dir) else {
347 tracing::warn!("Failed to read agent directory: {}", dir.display());
348 return agents;
349 };
350
351 for entry in entries.flatten() {
352 let path = entry.path();
353
354 if !path.is_file() {
356 continue;
357 }
358
359 let Some(ext) = path.extension().and_then(|e| e.to_str()) else {
360 continue;
361 };
362
363 let Ok(content) = std::fs::read_to_string(&path) else {
365 tracing::warn!("Failed to read agent file: {}", path.display());
366 continue;
367 };
368
369 let result = match ext {
371 "yaml" | "yml" => parse_agent_yaml(&content),
372 "md" => parse_agent_md(&content),
373 _ => continue,
374 };
375
376 match result {
377 Ok(agent) => {
378 tracing::debug!("Loaded agent '{}' from {}", agent.name, path.display());
379 agents.push(agent);
380 }
381 Err(e) => {
382 tracing::warn!("Failed to parse agent file {}: {}", path.display(), e);
383 }
384 }
385 }
386
387 agents
388}
389
390pub fn builtin_agents() -> Vec<AgentDefinition> {
392 vec![
393 AgentDefinition::new(
395 "explore",
396 "Fast codebase exploration agent. Use for searching files, reading code, \
397 and understanding codebase structure. Read-only operations only.",
398 )
399 .native()
400 .with_permissions(explore_permissions())
401 .with_max_steps(20)
402 .with_prompt(EXPLORE_PROMPT),
403 AgentDefinition::new(
405 "general",
406 "General-purpose agent for multi-step task execution. Can read, write, \
407 and execute commands. Cannot spawn subagents.",
408 )
409 .native()
410 .with_permissions(general_permissions())
411 .with_max_steps(50),
412 AgentDefinition::new(
414 "plan",
415 "Planning agent for designing implementation approaches. Read-only access \
416 to explore codebase and create plans.",
417 )
418 .native()
419 .with_mode(AgentMode::Primary)
420 .with_permissions(plan_permissions())
421 .with_max_steps(30)
422 .with_prompt(PLAN_PROMPT),
423 AgentDefinition::new(
425 "verification",
426 "Verification agent for adversarial validation. Prefer real checks, \
427 reproductions, and regression testing over code reading alone.",
428 )
429 .native()
430 .with_mode(AgentMode::Primary)
431 .with_permissions(verification_permissions())
432 .with_max_steps(30)
433 .with_prompt(VERIFICATION_PROMPT),
434 AgentDefinition::new(
436 "review",
437 "Code review agent focused on correctness, regressions, security, \
438 maintainability, and clear findings.",
439 )
440 .native()
441 .with_mode(AgentMode::Primary)
442 .with_permissions(review_permissions())
443 .with_max_steps(25)
444 .with_prompt(REVIEW_PROMPT),
445 AgentDefinition::new(
447 "title",
448 "Generate a concise title for the session based on conversation content.",
449 )
450 .native()
451 .hidden()
452 .with_mode(AgentMode::Primary)
453 .with_permissions(PermissionPolicy::new())
454 .with_max_steps(1)
455 .with_prompt(TITLE_PROMPT),
456 AgentDefinition::new(
458 "summary",
459 "Summarize the session conversation for context compaction.",
460 )
461 .native()
462 .hidden()
463 .with_mode(AgentMode::Primary)
464 .with_permissions(summary_permissions())
465 .with_max_steps(5)
466 .with_prompt(SUMMARY_PROMPT),
467 ]
468}
469
470fn explore_permissions() -> PermissionPolicy {
476 PermissionPolicy::new()
477 .allow_all(&["read", "grep", "glob", "ls"])
478 .deny_all(&["write", "edit", "task"])
479 .allow("Bash(ls:*)")
480 .allow("Bash(cat:*)")
481 .allow("Bash(head:*)")
482 .allow("Bash(tail:*)")
483 .allow("Bash(find:*)")
484 .allow("Bash(wc:*)")
485 .deny("Bash(rm:*)")
486 .deny("Bash(mv:*)")
487 .deny("Bash(cp:*)")
488}
489
490fn general_permissions() -> PermissionPolicy {
492 PermissionPolicy::new()
493 .allow_all(&["read", "write", "edit", "grep", "glob", "ls", "bash"])
494 .deny("task")
495}
496
497fn plan_permissions() -> PermissionPolicy {
499 PermissionPolicy::new()
500 .allow_all(&["read", "grep", "glob", "ls"])
501 .deny_all(&["write", "edit", "bash", "task"])
502}
503
504fn summary_permissions() -> PermissionPolicy {
506 PermissionPolicy::new()
507 .allow("read")
508 .deny_all(&["write", "edit", "bash", "grep", "glob", "ls", "task"])
509}
510
511fn verification_permissions() -> PermissionPolicy {
513 PermissionPolicy::new()
514 .allow_all(&["read", "grep", "glob", "ls", "bash"])
515 .deny_all(&["write", "edit", "task"])
516}
517
518fn review_permissions() -> PermissionPolicy {
520 PermissionPolicy::new()
521 .allow_all(&["read", "grep", "glob", "ls", "bash"])
522 .deny_all(&["write", "edit", "task"])
523}
524
525const EXPLORE_PROMPT: &str = crate::prompts::SUBAGENT_EXPLORE;
530
531const PLAN_PROMPT: &str = crate::prompts::SUBAGENT_PLAN;
532
533const VERIFICATION_PROMPT: &str = crate::prompts::AGENT_VERIFICATION;
534
535const REVIEW_PROMPT: &str = crate::prompts::SUBAGENT_CODE_REVIEW;
536
537const TITLE_PROMPT: &str = crate::prompts::SUBAGENT_TITLE;
538
539const SUMMARY_PROMPT: &str = crate::prompts::SUBAGENT_SUMMARY;
540
541#[cfg(test)]
546mod tests {
547 use super::*;
548
549 #[test]
550 fn test_agent_definition_builder() {
551 let agent = AgentDefinition::new("test", "Test agent")
552 .native()
553 .hidden()
554 .with_max_steps(10);
555
556 assert_eq!(agent.name, "test");
557 assert_eq!(agent.description, "Test agent");
558 assert!(agent.native);
559 assert!(agent.hidden);
560 assert_eq!(agent.max_steps, Some(10));
561 assert!(!agent.can_spawn_subagents);
562 }
563
564 #[test]
565 fn test_agent_registry_new() {
566 let registry = AgentRegistry::new();
567
568 assert!(registry.exists("explore"));
570 assert!(registry.exists("general"));
571 assert!(registry.exists("plan"));
572 assert!(registry.exists("verification"));
573 assert!(registry.exists("review"));
574 assert!(registry.exists("title"));
575 assert!(registry.exists("summary"));
576 assert_eq!(registry.len(), 7);
577 }
578
579 #[test]
580 fn test_agent_registry_get() {
581 let registry = AgentRegistry::new();
582
583 let explore = registry.get("explore").unwrap();
584 assert_eq!(explore.name, "explore");
585 assert!(explore.native);
586 assert!(!explore.hidden);
587
588 let title = registry.get("title").unwrap();
589 assert!(title.hidden);
590
591 assert!(registry.get("nonexistent").is_none());
592 }
593
594 #[test]
595 fn test_agent_registry_register_unregister() {
596 let registry = AgentRegistry::new();
597 let initial_count = registry.len();
598
599 let custom = AgentDefinition::new("custom", "Custom agent");
601 registry.register(custom);
602 assert_eq!(registry.len(), initial_count + 1);
603 assert!(registry.exists("custom"));
604
605 assert!(registry.unregister("custom"));
607 assert_eq!(registry.len(), initial_count);
608 assert!(!registry.exists("custom"));
609
610 assert!(!registry.unregister("nonexistent"));
612 }
613
614 #[test]
615 fn test_agent_registry_list_visible() {
616 let registry = AgentRegistry::new();
617
618 let visible = registry.list_visible();
619 let all = registry.list();
620
621 assert!(visible.len() < all.len());
623 assert!(visible.iter().all(|a| !a.hidden));
624 }
625
626 #[test]
627 fn test_builtin_agents() {
628 let agents = builtin_agents();
629
630 let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
632 assert!(names.contains(&"explore"));
633 assert!(names.contains(&"general"));
634 assert!(names.contains(&"plan"));
635 assert!(names.contains(&"verification"));
636 assert!(names.contains(&"review"));
637 assert!(names.contains(&"title"));
638 assert!(names.contains(&"summary"));
639
640 let explore = agents.iter().find(|a| a.name == "explore").unwrap();
642 assert!(!explore.permissions.deny.is_empty());
643
644 let general = agents.iter().find(|a| a.name == "general").unwrap();
646 assert!(!general.can_spawn_subagents);
647 }
648
649 #[test]
650 fn test_agent_mode_default() {
651 let mode = AgentMode::default();
652 assert_eq!(mode, AgentMode::Primary);
653 }
654
655 #[test]
660 fn test_parse_agent_yaml() {
661 let yaml = r#"
662name: test-agent
663description: A test agent
664mode: subagent
665hidden: false
666max_steps: 20
667"#;
668 let agent = parse_agent_yaml(yaml).unwrap();
669 assert_eq!(agent.name, "test-agent");
670 assert_eq!(agent.description, "A test agent");
671 assert_eq!(agent.mode, AgentMode::Subagent);
672 assert!(!agent.hidden);
673 assert_eq!(agent.max_steps, Some(20));
674 }
675
676 #[test]
677 fn test_parse_agent_yaml_with_permissions() {
678 let yaml = r#"
679name: restricted-agent
680description: Agent with permissions
681permissions:
682 allow:
683 - rule: read
684 - rule: grep
685 deny:
686 - rule: write
687"#;
688 let agent = parse_agent_yaml(yaml).unwrap();
689 assert_eq!(agent.name, "restricted-agent");
690 assert_eq!(agent.permissions.allow.len(), 2);
691 assert_eq!(agent.permissions.deny.len(), 1);
692 assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
694 assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
695 assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
696 }
697
698 #[test]
699 fn test_parse_agent_yaml_with_plain_string_permissions() {
700 let yaml = r#"
702name: plain-agent
703description: Agent with plain string permissions
704permissions:
705 allow:
706 - read
707 - grep
708 - "Bash(cargo:*)"
709 deny:
710 - write
711"#;
712 let agent = parse_agent_yaml(yaml).unwrap();
713 assert_eq!(agent.name, "plain-agent");
714 assert_eq!(agent.permissions.allow.len(), 3);
715 assert_eq!(agent.permissions.deny.len(), 1);
716 assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
718 assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
719 assert!(agent.permissions.allow[2]
720 .matches("Bash", &serde_json::json!({"command": "cargo build"})));
721 assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
722 }
723
724 #[test]
725 fn test_parse_agent_yaml_missing_name() {
726 let yaml = r#"
727description: Agent without name
728"#;
729 let result = parse_agent_yaml(yaml);
730 assert!(result.is_err());
731 }
732
733 #[test]
734 fn test_parse_agent_md() {
735 let md = r#"---
736name: md-agent
737description: Agent from markdown
738mode: subagent
739max_steps: 15
740---
741# System Prompt
742
743You are a helpful agent.
744Do your best work.
745"#;
746 let agent = parse_agent_md(md).unwrap();
747 assert_eq!(agent.name, "md-agent");
748 assert_eq!(agent.description, "Agent from markdown");
749 assert_eq!(agent.max_steps, Some(15));
750 assert!(agent.prompt.is_some());
751 assert!(agent.prompt.unwrap().contains("helpful agent"));
752 }
753
754 #[test]
755 fn test_parse_agent_md_with_prompt_in_frontmatter() {
756 let md = r#"---
757name: prompt-agent
758description: Agent with prompt in frontmatter
759prompt: "Frontmatter prompt"
760---
761Body content that should be ignored
762"#;
763 let agent = parse_agent_md(md).unwrap();
764 assert_eq!(agent.prompt.unwrap(), "Frontmatter prompt");
765 }
766
767 #[test]
768 fn test_parse_agent_md_missing_frontmatter() {
769 let md = "Just markdown without frontmatter";
770 let result = parse_agent_md(md);
771 assert!(result.is_err());
772 }
773
774 #[test]
775 fn test_load_agents_from_dir() {
776 let temp_dir = tempfile::tempdir().unwrap();
777
778 std::fs::write(
780 temp_dir.path().join("agent1.yaml"),
781 r#"
782name: yaml-agent
783description: Agent from YAML file
784"#,
785 )
786 .unwrap();
787
788 std::fs::write(
790 temp_dir.path().join("agent2.md"),
791 r#"---
792name: md-agent
793description: Agent from Markdown file
794---
795System prompt here
796"#,
797 )
798 .unwrap();
799
800 std::fs::write(temp_dir.path().join("invalid.yaml"), "not: valid: yaml: [").unwrap();
802
803 std::fs::write(temp_dir.path().join("readme.txt"), "Just a text file").unwrap();
805
806 let agents = load_agents_from_dir(temp_dir.path());
807 assert_eq!(agents.len(), 2);
808
809 let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
810 assert!(names.contains(&"yaml-agent"));
811 assert!(names.contains(&"md-agent"));
812 }
813
814 #[test]
815 fn test_load_agents_from_nonexistent_dir() {
816 let agents = load_agents_from_dir(std::path::Path::new("/nonexistent/dir"));
817 assert!(agents.is_empty());
818 }
819
820 #[test]
821 fn test_registry_with_config() {
822 let temp_dir = tempfile::tempdir().unwrap();
823
824 std::fs::write(
826 temp_dir.path().join("custom.yaml"),
827 r#"
828name: custom-agent
829description: Custom agent from config
830"#,
831 )
832 .unwrap();
833
834 let config = CodeConfig::new().add_agent_dir(temp_dir.path());
835 let registry = AgentRegistry::with_config(&config);
836
837 assert!(registry.exists("explore"));
839 assert!(registry.exists("custom-agent"));
840 assert_eq!(registry.len(), 8); }
842
843 #[test]
844 fn test_agent_definition_with_model() {
845 let model = ModelConfig {
846 model: "claude-3-5-sonnet".to_string(),
847 provider: Some("anthropic".to_string()),
848 };
849 let agent = AgentDefinition::new("test", "Test").with_model(model);
850 assert!(agent.model.is_some());
851 assert_eq!(agent.model.unwrap().provider, Some("anthropic".to_string()));
852 }
853
854 #[test]
855 fn test_agent_definition_allow_subagents() {
856 let agent = AgentDefinition::new("test", "Test").allow_subagents();
857 assert!(agent.can_spawn_subagents);
858 }
859
860 #[test]
861 fn test_agent_registry_default() {
862 let registry = AgentRegistry::default();
863 assert!(!registry.is_empty());
864 assert_eq!(registry.len(), 7);
865 }
866
867 #[test]
868 fn test_agent_registry_is_empty() {
869 let registry = AgentRegistry {
870 agents: RwLock::new(HashMap::new()),
871 };
872 assert!(registry.is_empty());
873 assert_eq!(registry.len(), 0);
874 }
875}