Skip to main content

a3s_code_core/
subagent.rs

1//! Subagent System
2//!
3//! Provides a system for delegating specialized tasks to focused child agents.
4//! Each subagent runs in an isolated child session with restricted permissions.
5//!
6//! ## Architecture
7//!
8//! ```text
9//! Parent Session
10//!   └── Task Tool
11//!         ├── AgentRegistry (lookup agent definitions)
12//!         └── Child Session (isolated execution)
13//!               ├── Restricted permissions
14//!               ├── Optional model override
15//!               └── Event forwarding to parent
16//! ```
17//!
18//! ## Built-in Agents
19//!
20//! - `explore`: Fast codebase exploration (read-only)
21//! - `general`: Multi-step task execution
22//! - `plan`: Read-only planning mode
23//! - `verification`: Adversarial verification specialist
24//! - `review`: Code review specialist
25//! - `title`: Session title generation (hidden)
26//! - `summary`: Session summarization (hidden)
27//!
28//! ## Loading Agents from Files
29//!
30//! Agents can be loaded from YAML or Markdown files:
31//!
32//! ### YAML Format
33//! ```yaml
34//! name: my-agent
35//! description: Custom agent for specific tasks
36//! mode: subagent
37//! hidden: false
38//! max_steps: 30
39//! permissions:
40//!   allow:
41//!     - read
42//!     - grep
43//!   deny:
44//!     - write
45//! prompt: |
46//!   You are a specialized agent...
47//! ```
48//!
49//! ### Markdown Format
50//! ```markdown
51//! ---
52//! name: my-agent
53//! description: Custom agent
54//! mode: subagent
55//! max_steps: 30
56//! ---
57//! # System Prompt
58//! You are a specialized agent...
59//! ```
60
61use 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/// Agent execution mode
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
72#[serde(rename_all = "snake_case")]
73pub enum AgentMode {
74    /// Primary agent (main conversation)
75    #[default]
76    Primary,
77    /// Subagent (child session for delegated tasks)
78    Subagent,
79}
80
81/// Model configuration for agent
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct ModelConfig {
84    /// Model identifier (e.g., "claude-3-5-sonnet-20241022")
85    pub model: String,
86    /// Optional provider override
87    pub provider: Option<String>,
88}
89
90/// Agent definition
91///
92/// Defines the configuration and capabilities of an agent type.
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct AgentDefinition {
95    /// Agent identifier (e.g., "explore", "plan", "general")
96    pub name: String,
97    /// Description of what the agent does
98    pub description: String,
99    /// Agent mode: "subagent" or "primary"
100    #[serde(default)]
101    pub mode: AgentMode,
102    /// Whether this is a built-in agent
103    #[serde(default)]
104    pub native: bool,
105    /// Whether to hide from UI
106    #[serde(default)]
107    pub hidden: bool,
108    /// Permission rules for this agent
109    #[serde(default)]
110    pub permissions: PermissionPolicy,
111    /// Optional model override
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub model: Option<ModelConfig>,
114    /// System prompt for this agent
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub prompt: Option<String>,
117    /// Maximum execution steps (tool rounds)
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub max_steps: Option<usize>,
120    /// Whether this agent can spawn subagents (default: false)
121    #[serde(default)]
122    pub can_spawn_subagents: bool,
123}
124
125impl AgentDefinition {
126    /// Create a new agent definition
127    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    /// Set agent mode
143    pub fn with_mode(mut self, mode: AgentMode) -> Self {
144        self.mode = mode;
145        self
146    }
147
148    /// Mark as native (built-in)
149    pub fn native(mut self) -> Self {
150        self.native = true;
151        self
152    }
153
154    /// Mark as hidden from UI
155    pub fn hidden(mut self) -> Self {
156        self.hidden = true;
157        self
158    }
159
160    /// Set permission policy
161    pub fn with_permissions(mut self, permissions: PermissionPolicy) -> Self {
162        self.permissions = permissions;
163        self
164    }
165
166    /// Set model override
167    pub fn with_model(mut self, model: ModelConfig) -> Self {
168        self.model = Some(model);
169        self
170    }
171
172    /// Set system prompt
173    pub fn with_prompt(mut self, prompt: &str) -> Self {
174        self.prompt = Some(prompt.to_string());
175        self
176    }
177
178    /// Set maximum execution steps
179    pub fn with_max_steps(mut self, max_steps: usize) -> Self {
180        self.max_steps = Some(max_steps);
181        self
182    }
183
184    /// Allow spawning subagents
185    pub fn allow_subagents(mut self) -> Self {
186        self.can_spawn_subagents = true;
187        self
188    }
189}
190
191/// Agent registry for managing agent definitions
192///
193/// Thread-safe registry that stores agent definitions and provides
194/// lookup functionality.
195pub 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    /// Create a new agent registry with built-in agents
207    pub fn new() -> Self {
208        let registry = Self {
209            agents: RwLock::new(HashMap::new()),
210        };
211
212        // Register built-in agents
213        for agent in builtin_agents() {
214            registry.register(agent);
215        }
216
217        registry
218    }
219
220    /// Create a new agent registry with configuration
221    ///
222    /// Loads built-in agents first, then loads agents from configured directories.
223    pub fn with_config(config: &CodeConfig) -> Self {
224        let registry = Self::new();
225
226        // Load agents from configured directories
227        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    /// Register an agent definition
239    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    /// Unregister an agent by name
246    ///
247    /// Returns true if the agent was removed, false if not found.
248    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    /// Get an agent definition by name
254    pub fn get(&self, name: &str) -> Option<AgentDefinition> {
255        let agents = read_or_recover(&self.agents);
256        agents.get(name).cloned()
257    }
258
259    /// List all registered agents
260    pub fn list(&self) -> Vec<AgentDefinition> {
261        let agents = read_or_recover(&self.agents);
262        agents.values().cloned().collect()
263    }
264
265    /// List visible agents (not hidden)
266    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    /// Check if an agent exists
272    pub fn exists(&self, name: &str) -> bool {
273        let agents = read_or_recover(&self.agents);
274        agents.contains_key(name)
275    }
276
277    /// Get the number of registered agents
278    pub fn len(&self) -> usize {
279        let agents = read_or_recover(&self.agents);
280        agents.len()
281    }
282
283    /// Check if the registry is empty
284    pub fn is_empty(&self) -> bool {
285        self.len() == 0
286    }
287}
288
289// ============================================================================
290// Agent File Loading
291// ============================================================================
292
293/// Parse an agent definition from YAML content
294///
295/// The YAML should contain fields matching AgentDefinition structure.
296pub 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
307/// Parse an agent definition from Markdown with YAML frontmatter
308///
309/// The frontmatter contains agent metadata, and the body becomes the prompt.
310pub fn parse_agent_md(content: &str) -> anyhow::Result<AgentDefinition> {
311    // Parse frontmatter (YAML between --- markers)
312    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    // Parse the frontmatter as YAML
324    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    // Use body as prompt if not already set in frontmatter
332    if agent.prompt.is_none() && !body.is_empty() {
333        agent.prompt = Some(body.to_string());
334    }
335
336    Ok(agent)
337}
338
339/// Load all agent definitions from a directory
340///
341/// Scans for *.yaml and *.md files and parses them as agent definitions.
342/// Invalid files are logged and skipped.
343pub 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        // Skip non-files
355        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        // Read file content
364        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        // Parse based on extension
370        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
390/// Create built-in agent definitions
391pub fn builtin_agents() -> Vec<AgentDefinition> {
392    vec![
393        // Explore agent: Fast codebase exploration (read-only)
394        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        // General agent: Multi-step task execution
404        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        // Plan agent: Read-only planning mode
413        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        // Verification agent: adversarial validation and repro
424        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        // Review agent: review-focused analysis
435        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        // Title agent: Session title generation (hidden)
446        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        // Summary agent: Session summarization (hidden)
457        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
470// ============================================================================
471// Permission Policies for Built-in Agents
472// ============================================================================
473
474/// Permission policy for explore agent (read-only)
475fn 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
490/// Permission policy for general agent (full access except task)
491fn general_permissions() -> PermissionPolicy {
492    PermissionPolicy::new()
493        .allow_all(&["read", "write", "edit", "grep", "glob", "ls", "bash"])
494        .deny("task")
495}
496
497/// Permission policy for plan agent (read-only)
498fn plan_permissions() -> PermissionPolicy {
499    PermissionPolicy::new()
500        .allow_all(&["read", "grep", "glob", "ls"])
501        .deny_all(&["write", "edit", "bash", "task"])
502}
503
504/// Permission policy for summary agent (read-only)
505fn summary_permissions() -> PermissionPolicy {
506    PermissionPolicy::new()
507        .allow("read")
508        .deny_all(&["write", "edit", "bash", "grep", "glob", "ls", "task"])
509}
510
511/// Permission policy for verification agent (read-heavy with runtime checks)
512fn verification_permissions() -> PermissionPolicy {
513    PermissionPolicy::new()
514        .allow_all(&["read", "grep", "glob", "ls", "bash"])
515        .deny_all(&["write", "edit", "task"])
516}
517
518/// Permission policy for review agent (read-heavy with optional lightweight checks)
519fn review_permissions() -> PermissionPolicy {
520    PermissionPolicy::new()
521        .allow_all(&["read", "grep", "glob", "ls", "bash"])
522        .deny_all(&["write", "edit", "task"])
523}
524
525// ============================================================================
526// System Prompts for Built-in Agents
527// ============================================================================
528
529const 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// ============================================================================
542// Tests
543// ============================================================================
544
545#[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        // Should have built-in agents
569        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        // Register custom agent
600        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        // Unregister
606        assert!(registry.unregister("custom"));
607        assert_eq!(registry.len(), initial_count);
608        assert!(!registry.exists("custom"));
609
610        // Unregister non-existent
611        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        // Hidden agents should not be in visible list
622        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        // Check we have expected agents
631        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        // Check explore is read-only (has deny rules for write)
641        let explore = agents.iter().find(|a| a.name == "explore").unwrap();
642        assert!(!explore.permissions.deny.is_empty());
643
644        // Check general cannot spawn subagents
645        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    // ========================================================================
656    // Agent File Loading Tests
657    // ========================================================================
658
659    #[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        // Verify that deserialized rules actually match (tool_name populated)
693        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        // Users naturally write plain strings in allow/deny lists
701        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        // Verify rules are functional
717        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        // Create a YAML agent file
779        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        // Create a Markdown agent file
789        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        // Create an invalid file (should be skipped)
801        std::fs::write(temp_dir.path().join("invalid.yaml"), "not: valid: yaml: [").unwrap();
802
803        // Create a non-agent file (should be skipped)
804        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        // Create an agent file
825        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        // Should have built-in agents plus custom agent
838        assert!(registry.exists("explore"));
839        assert!(registry.exists("custom-agent"));
840        assert_eq!(registry.len(), 8); // 7 built-in + 1 custom
841    }
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}