use crate::config::CodeConfig;
use crate::permissions::PermissionPolicy;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::sync::RwLock;
use crate::error::{read_or_recover, write_or_recover};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum AgentMode {
#[default]
Primary,
Subagent,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
pub model: String,
pub provider: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentDefinition {
pub name: String,
pub description: String,
#[serde(default)]
pub mode: AgentMode,
#[serde(default)]
pub native: bool,
#[serde(default)]
pub hidden: bool,
#[serde(default)]
pub permissions: PermissionPolicy,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<ModelConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_steps: Option<usize>,
#[serde(default)]
pub can_spawn_subagents: bool,
}
impl AgentDefinition {
pub fn new(name: &str, description: &str) -> Self {
Self {
name: name.to_string(),
description: description.to_string(),
mode: AgentMode::Subagent,
native: false,
hidden: false,
permissions: PermissionPolicy::default(),
model: None,
prompt: None,
max_steps: None,
can_spawn_subagents: false,
}
}
pub fn with_mode(mut self, mode: AgentMode) -> Self {
self.mode = mode;
self
}
pub fn native(mut self) -> Self {
self.native = true;
self
}
pub fn hidden(mut self) -> Self {
self.hidden = true;
self
}
pub fn with_permissions(mut self, permissions: PermissionPolicy) -> Self {
self.permissions = permissions;
self
}
pub fn with_model(mut self, model: ModelConfig) -> Self {
self.model = Some(model);
self
}
pub fn with_prompt(mut self, prompt: &str) -> Self {
self.prompt = Some(prompt.to_string());
self
}
pub fn with_max_steps(mut self, max_steps: usize) -> Self {
self.max_steps = Some(max_steps);
self
}
pub fn allow_subagents(mut self) -> Self {
self.can_spawn_subagents = true;
self
}
}
pub struct AgentRegistry {
agents: RwLock<HashMap<String, AgentDefinition>>,
}
impl Default for AgentRegistry {
fn default() -> Self {
Self::new()
}
}
impl AgentRegistry {
pub fn new() -> Self {
let registry = Self {
agents: RwLock::new(HashMap::new()),
};
for agent in builtin_agents() {
registry.register(agent);
}
registry
}
pub fn with_config(config: &CodeConfig) -> Self {
let registry = Self::new();
for dir in &config.agent_dirs {
let agents = load_agents_from_dir(dir);
for agent in agents {
tracing::info!("Loaded agent '{}' from {}", agent.name, dir.display());
registry.register(agent);
}
}
registry
}
pub fn register(&self, agent: AgentDefinition) {
let mut agents = write_or_recover(&self.agents);
tracing::debug!("Registering agent: {}", agent.name);
agents.insert(agent.name.clone(), agent);
}
pub fn unregister(&self, name: &str) -> bool {
let mut agents = write_or_recover(&self.agents);
agents.remove(name).is_some()
}
pub fn get(&self, name: &str) -> Option<AgentDefinition> {
let agents = read_or_recover(&self.agents);
agents.get(name).cloned()
}
pub fn list(&self) -> Vec<AgentDefinition> {
let agents = read_or_recover(&self.agents);
agents.values().cloned().collect()
}
pub fn list_visible(&self) -> Vec<AgentDefinition> {
let agents = read_or_recover(&self.agents);
agents.values().filter(|a| !a.hidden).cloned().collect()
}
pub fn exists(&self, name: &str) -> bool {
let agents = read_or_recover(&self.agents);
agents.contains_key(name)
}
pub fn len(&self) -> usize {
let agents = read_or_recover(&self.agents);
agents.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
pub fn parse_agent_yaml(content: &str) -> anyhow::Result<AgentDefinition> {
let agent: AgentDefinition = serde_yaml::from_str(content)
.map_err(|e| anyhow::anyhow!("Failed to parse agent YAML: {}", e))?;
if agent.name.is_empty() {
return Err(anyhow::anyhow!("Agent name is required"));
}
Ok(agent)
}
pub fn parse_agent_md(content: &str) -> anyhow::Result<AgentDefinition> {
let parts: Vec<&str> = content.splitn(3, "---").collect();
if parts.len() < 3 {
return Err(anyhow::anyhow!(
"Invalid markdown format: missing YAML frontmatter"
));
}
let frontmatter = parts[1].trim();
let body = parts[2].trim();
let mut agent: AgentDefinition = serde_yaml::from_str(frontmatter)
.map_err(|e| anyhow::anyhow!("Failed to parse agent frontmatter: {}", e))?;
if agent.name.is_empty() {
return Err(anyhow::anyhow!("Agent name is required"));
}
if agent.prompt.is_none() && !body.is_empty() {
agent.prompt = Some(body.to_string());
}
Ok(agent)
}
pub fn load_agents_from_dir(dir: &Path) -> Vec<AgentDefinition> {
let mut agents = Vec::new();
let Ok(entries) = std::fs::read_dir(dir) else {
tracing::warn!("Failed to read agent directory: {}", dir.display());
return agents;
};
for entry in entries.flatten() {
let path = entry.path();
if !path.is_file() {
continue;
}
let Some(ext) = path.extension().and_then(|e| e.to_str()) else {
continue;
};
let Ok(content) = std::fs::read_to_string(&path) else {
tracing::warn!("Failed to read agent file: {}", path.display());
continue;
};
let result = match ext {
"yaml" | "yml" => parse_agent_yaml(&content),
"md" => parse_agent_md(&content),
_ => continue,
};
match result {
Ok(agent) => {
tracing::debug!("Loaded agent '{}' from {}", agent.name, path.display());
agents.push(agent);
}
Err(e) => {
tracing::warn!("Failed to parse agent file {}: {}", path.display(), e);
}
}
}
agents
}
pub fn builtin_agents() -> Vec<AgentDefinition> {
vec![
AgentDefinition::new(
"explore",
"Fast codebase exploration agent. Use for searching files, reading code, \
and understanding codebase structure. Read-only operations only.",
)
.native()
.with_permissions(explore_permissions())
.with_max_steps(20)
.with_prompt(EXPLORE_PROMPT),
AgentDefinition::new(
"general",
"General-purpose agent for multi-step task execution. Can read, write, \
and execute commands. Cannot spawn subagents.",
)
.native()
.with_permissions(general_permissions())
.with_max_steps(50),
AgentDefinition::new(
"plan",
"Planning agent for designing implementation approaches. Read-only access \
to explore codebase and create plans.",
)
.native()
.with_mode(AgentMode::Primary)
.with_permissions(plan_permissions())
.with_max_steps(30)
.with_prompt(PLAN_PROMPT),
AgentDefinition::new(
"verification",
"Verification agent for adversarial validation. Prefer real checks, \
reproductions, and regression testing over code reading alone.",
)
.native()
.with_mode(AgentMode::Primary)
.with_permissions(verification_permissions())
.with_max_steps(30)
.with_prompt(VERIFICATION_PROMPT),
AgentDefinition::new(
"review",
"Code review agent focused on correctness, regressions, security, \
maintainability, and clear findings.",
)
.native()
.with_mode(AgentMode::Primary)
.with_permissions(review_permissions())
.with_max_steps(25)
.with_prompt(REVIEW_PROMPT),
AgentDefinition::new(
"title",
"Generate a concise title for the session based on conversation content.",
)
.native()
.hidden()
.with_mode(AgentMode::Primary)
.with_permissions(PermissionPolicy::new())
.with_max_steps(1)
.with_prompt(TITLE_PROMPT),
AgentDefinition::new(
"summary",
"Summarize the session conversation for context compaction.",
)
.native()
.hidden()
.with_mode(AgentMode::Primary)
.with_permissions(summary_permissions())
.with_max_steps(5)
.with_prompt(SUMMARY_PROMPT),
]
}
fn explore_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow_all(&["read", "grep", "glob", "ls"])
.deny_all(&["write", "edit", "task"])
.allow("Bash(ls:*)")
.allow("Bash(cat:*)")
.allow("Bash(head:*)")
.allow("Bash(tail:*)")
.allow("Bash(find:*)")
.allow("Bash(wc:*)")
.deny("Bash(rm:*)")
.deny("Bash(mv:*)")
.deny("Bash(cp:*)")
}
fn general_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow_all(&["read", "write", "edit", "grep", "glob", "ls", "bash"])
.deny("task")
}
fn plan_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow_all(&["read", "grep", "glob", "ls"])
.deny_all(&["write", "edit", "bash", "task"])
}
fn summary_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow("read")
.deny_all(&["write", "edit", "bash", "grep", "glob", "ls", "task"])
}
fn verification_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow_all(&["read", "grep", "glob", "ls", "bash"])
.deny_all(&["write", "edit", "task"])
}
fn review_permissions() -> PermissionPolicy {
PermissionPolicy::new()
.allow_all(&["read", "grep", "glob", "ls", "bash"])
.deny_all(&["write", "edit", "task"])
}
const EXPLORE_PROMPT: &str = crate::prompts::SUBAGENT_EXPLORE;
const PLAN_PROMPT: &str = crate::prompts::SUBAGENT_PLAN;
const VERIFICATION_PROMPT: &str = crate::prompts::AGENT_VERIFICATION;
const REVIEW_PROMPT: &str = crate::prompts::SUBAGENT_CODE_REVIEW;
const TITLE_PROMPT: &str = crate::prompts::SUBAGENT_TITLE;
const SUMMARY_PROMPT: &str = crate::prompts::SUBAGENT_SUMMARY;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_agent_definition_builder() {
let agent = AgentDefinition::new("test", "Test agent")
.native()
.hidden()
.with_max_steps(10);
assert_eq!(agent.name, "test");
assert_eq!(agent.description, "Test agent");
assert!(agent.native);
assert!(agent.hidden);
assert_eq!(agent.max_steps, Some(10));
assert!(!agent.can_spawn_subagents);
}
#[test]
fn test_agent_registry_new() {
let registry = AgentRegistry::new();
assert!(registry.exists("explore"));
assert!(registry.exists("general"));
assert!(registry.exists("plan"));
assert!(registry.exists("verification"));
assert!(registry.exists("review"));
assert!(registry.exists("title"));
assert!(registry.exists("summary"));
assert_eq!(registry.len(), 7);
}
#[test]
fn test_agent_registry_get() {
let registry = AgentRegistry::new();
let explore = registry.get("explore").unwrap();
assert_eq!(explore.name, "explore");
assert!(explore.native);
assert!(!explore.hidden);
let title = registry.get("title").unwrap();
assert!(title.hidden);
assert!(registry.get("nonexistent").is_none());
}
#[test]
fn test_agent_registry_register_unregister() {
let registry = AgentRegistry::new();
let initial_count = registry.len();
let custom = AgentDefinition::new("custom", "Custom agent");
registry.register(custom);
assert_eq!(registry.len(), initial_count + 1);
assert!(registry.exists("custom"));
assert!(registry.unregister("custom"));
assert_eq!(registry.len(), initial_count);
assert!(!registry.exists("custom"));
assert!(!registry.unregister("nonexistent"));
}
#[test]
fn test_agent_registry_list_visible() {
let registry = AgentRegistry::new();
let visible = registry.list_visible();
let all = registry.list();
assert!(visible.len() < all.len());
assert!(visible.iter().all(|a| !a.hidden));
}
#[test]
fn test_builtin_agents() {
let agents = builtin_agents();
let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
assert!(names.contains(&"explore"));
assert!(names.contains(&"general"));
assert!(names.contains(&"plan"));
assert!(names.contains(&"verification"));
assert!(names.contains(&"review"));
assert!(names.contains(&"title"));
assert!(names.contains(&"summary"));
let explore = agents.iter().find(|a| a.name == "explore").unwrap();
assert!(!explore.permissions.deny.is_empty());
let general = agents.iter().find(|a| a.name == "general").unwrap();
assert!(!general.can_spawn_subagents);
}
#[test]
fn test_agent_mode_default() {
let mode = AgentMode::default();
assert_eq!(mode, AgentMode::Primary);
}
#[test]
fn test_parse_agent_yaml() {
let yaml = r#"
name: test-agent
description: A test agent
mode: subagent
hidden: false
max_steps: 20
"#;
let agent = parse_agent_yaml(yaml).unwrap();
assert_eq!(agent.name, "test-agent");
assert_eq!(agent.description, "A test agent");
assert_eq!(agent.mode, AgentMode::Subagent);
assert!(!agent.hidden);
assert_eq!(agent.max_steps, Some(20));
}
#[test]
fn test_parse_agent_yaml_with_permissions() {
let yaml = r#"
name: restricted-agent
description: Agent with permissions
permissions:
allow:
- rule: read
- rule: grep
deny:
- rule: write
"#;
let agent = parse_agent_yaml(yaml).unwrap();
assert_eq!(agent.name, "restricted-agent");
assert_eq!(agent.permissions.allow.len(), 2);
assert_eq!(agent.permissions.deny.len(), 1);
assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
}
#[test]
fn test_parse_agent_yaml_with_plain_string_permissions() {
let yaml = r#"
name: plain-agent
description: Agent with plain string permissions
permissions:
allow:
- read
- grep
- "Bash(cargo:*)"
deny:
- write
"#;
let agent = parse_agent_yaml(yaml).unwrap();
assert_eq!(agent.name, "plain-agent");
assert_eq!(agent.permissions.allow.len(), 3);
assert_eq!(agent.permissions.deny.len(), 1);
assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
assert!(agent.permissions.allow[2]
.matches("Bash", &serde_json::json!({"command": "cargo build"})));
assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
}
#[test]
fn test_parse_agent_yaml_missing_name() {
let yaml = r#"
description: Agent without name
"#;
let result = parse_agent_yaml(yaml);
assert!(result.is_err());
}
#[test]
fn test_parse_agent_md() {
let md = r#"---
name: md-agent
description: Agent from markdown
mode: subagent
max_steps: 15
---
# System Prompt
You are a helpful agent.
Do your best work.
"#;
let agent = parse_agent_md(md).unwrap();
assert_eq!(agent.name, "md-agent");
assert_eq!(agent.description, "Agent from markdown");
assert_eq!(agent.max_steps, Some(15));
assert!(agent.prompt.is_some());
assert!(agent.prompt.unwrap().contains("helpful agent"));
}
#[test]
fn test_parse_agent_md_with_prompt_in_frontmatter() {
let md = r#"---
name: prompt-agent
description: Agent with prompt in frontmatter
prompt: "Frontmatter prompt"
---
Body content that should be ignored
"#;
let agent = parse_agent_md(md).unwrap();
assert_eq!(agent.prompt.unwrap(), "Frontmatter prompt");
}
#[test]
fn test_parse_agent_md_missing_frontmatter() {
let md = "Just markdown without frontmatter";
let result = parse_agent_md(md);
assert!(result.is_err());
}
#[test]
fn test_load_agents_from_dir() {
let temp_dir = tempfile::tempdir().unwrap();
std::fs::write(
temp_dir.path().join("agent1.yaml"),
r#"
name: yaml-agent
description: Agent from YAML file
"#,
)
.unwrap();
std::fs::write(
temp_dir.path().join("agent2.md"),
r#"---
name: md-agent
description: Agent from Markdown file
---
System prompt here
"#,
)
.unwrap();
std::fs::write(temp_dir.path().join("invalid.yaml"), "not: valid: yaml: [").unwrap();
std::fs::write(temp_dir.path().join("readme.txt"), "Just a text file").unwrap();
let agents = load_agents_from_dir(temp_dir.path());
assert_eq!(agents.len(), 2);
let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
assert!(names.contains(&"yaml-agent"));
assert!(names.contains(&"md-agent"));
}
#[test]
fn test_load_agents_from_nonexistent_dir() {
let agents = load_agents_from_dir(std::path::Path::new("/nonexistent/dir"));
assert!(agents.is_empty());
}
#[test]
fn test_registry_with_config() {
let temp_dir = tempfile::tempdir().unwrap();
std::fs::write(
temp_dir.path().join("custom.yaml"),
r#"
name: custom-agent
description: Custom agent from config
"#,
)
.unwrap();
let config = CodeConfig::new().add_agent_dir(temp_dir.path());
let registry = AgentRegistry::with_config(&config);
assert!(registry.exists("explore"));
assert!(registry.exists("custom-agent"));
assert_eq!(registry.len(), 8); }
#[test]
fn test_agent_definition_with_model() {
let model = ModelConfig {
model: "claude-3-5-sonnet".to_string(),
provider: Some("anthropic".to_string()),
};
let agent = AgentDefinition::new("test", "Test").with_model(model);
assert!(agent.model.is_some());
assert_eq!(agent.model.unwrap().provider, Some("anthropic".to_string()));
}
#[test]
fn test_agent_definition_allow_subagents() {
let agent = AgentDefinition::new("test", "Test").allow_subagents();
assert!(agent.can_spawn_subagents);
}
#[test]
fn test_agent_registry_default() {
let registry = AgentRegistry::default();
assert!(!registry.is_empty());
assert_eq!(registry.len(), 7);
}
#[test]
fn test_agent_registry_is_empty() {
let registry = AgentRegistry {
agents: RwLock::new(HashMap::new()),
};
assert!(registry.is_empty());
assert_eq!(registry.len(), 0);
}
}