1use crate::config::CodeConfig;
58use crate::permissions::{PermissionChecker, PermissionPolicy};
59use serde::{Deserialize, Serialize};
60use std::collections::HashMap;
61use std::path::Path;
62use std::sync::RwLock;
63
64use crate::error::{read_or_recover, write_or_recover};
65
66#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
68#[serde(rename_all = "snake_case")]
69pub enum ConfirmationInheritance {
70 #[default]
73 AutoApprove,
74 DenyOnAsk,
76 InheritParent,
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct ModelConfig {
83 pub model: String,
85 pub provider: Option<String>,
87}
88
89impl ModelConfig {
90 pub fn new(model: impl Into<String>) -> Self {
92 Self {
93 model: model.into(),
94 provider: None,
95 }
96 }
97
98 pub fn with_provider(provider: impl Into<String>, model: impl Into<String>) -> Self {
100 Self {
101 model: model.into(),
102 provider: Some(provider.into()),
103 }
104 }
105
106 pub fn from_model_ref(model_ref: impl AsRef<str>) -> Self {
111 let model_ref = model_ref.as_ref();
112 if let Some((provider, model)) = model_ref.split_once('/') {
113 Self::with_provider(provider, model)
114 } else {
115 Self::new(model_ref)
116 }
117 }
118
119 pub fn model_ref(&self) -> String {
121 match &self.provider {
122 Some(provider) => format!("{}/{}", provider, self.model),
123 None => self.model.clone(),
124 }
125 }
126}
127
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
134#[serde(rename_all = "snake_case")]
135pub enum WorkerAgentKind {
136 #[serde(alias = "readonly", alias = "read-only", alias = "explore")]
138 ReadOnly,
139 #[serde(alias = "plan")]
141 Planner,
142 #[serde(alias = "implementation", alias = "general")]
144 Implementer,
145 #[serde(alias = "verification", alias = "verify")]
147 Verifier,
148 #[serde(alias = "review", alias = "code-review")]
150 Reviewer,
151 Custom,
153}
154
155#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct WorkerAgentSpec {
162 pub name: String,
164 pub description: String,
166 pub kind: WorkerAgentKind,
168 #[serde(default)]
170 pub hidden: bool,
171 #[serde(skip_serializing_if = "Option::is_none")]
173 pub permissions: Option<PermissionPolicy>,
174 #[serde(skip_serializing_if = "Option::is_none")]
176 pub model: Option<ModelConfig>,
177 #[serde(skip_serializing_if = "Option::is_none")]
179 pub prompt: Option<String>,
180 #[serde(skip_serializing_if = "Option::is_none")]
182 pub max_steps: Option<usize>,
183 #[serde(default, skip_serializing_if = "Option::is_none")]
185 pub confirmation_inheritance: Option<ConfirmationInheritance>,
186}
187
188impl WorkerAgentKind {
189 pub fn as_str(self) -> &'static str {
191 match self {
192 Self::ReadOnly => "read_only",
193 Self::Planner => "planner",
194 Self::Implementer => "implementer",
195 Self::Verifier => "verifier",
196 Self::Reviewer => "reviewer",
197 Self::Custom => "custom",
198 }
199 }
200
201 fn default_permissions(self) -> PermissionPolicy {
202 match self {
203 Self::ReadOnly => explore_permissions(),
204 Self::Planner => plan_permissions(),
205 Self::Implementer => general_permissions(),
206 Self::Verifier => verification_permissions(),
207 Self::Reviewer => review_permissions(),
208 Self::Custom => PermissionPolicy::strict(),
209 }
210 }
211
212 fn default_prompt(self) -> Option<&'static str> {
213 match self {
214 Self::ReadOnly => Some(EXPLORE_PROMPT),
215 Self::Planner => Some(PLAN_PROMPT),
216 Self::Verifier => Some(VERIFICATION_PROMPT),
217 Self::Reviewer => Some(REVIEW_PROMPT),
218 Self::Implementer | Self::Custom => None,
219 }
220 }
221
222 fn default_max_steps(self) -> usize {
223 match self {
224 Self::ReadOnly => 20,
225 Self::Planner => 30,
226 Self::Implementer => 50,
227 Self::Verifier => 30,
228 Self::Reviewer => 25,
229 Self::Custom => 30,
230 }
231 }
232}
233
234impl std::fmt::Display for WorkerAgentKind {
235 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236 f.write_str(self.as_str())
237 }
238}
239
240impl std::str::FromStr for WorkerAgentKind {
241 type Err = anyhow::Error;
242
243 fn from_str(value: &str) -> anyhow::Result<Self> {
244 match value.trim().to_ascii_lowercase().as_str() {
245 "read_only" | "readonly" | "read-only" | "explore" | "scanner" => Ok(Self::ReadOnly),
246 "planner" | "plan" => Ok(Self::Planner),
247 "implementer" | "implementation" | "general" | "executor" => Ok(Self::Implementer),
248 "verifier" | "verification" | "verify" | "tester" => Ok(Self::Verifier),
249 "reviewer" | "review" | "code-review" | "code_reviewer" => Ok(Self::Reviewer),
250 "custom" => Ok(Self::Custom),
251 other => Err(anyhow::anyhow!("unknown worker agent kind '{}'", other)),
252 }
253 }
254}
255
256pub type CattleAgentKind = WorkerAgentKind;
258pub type CattleAgentSpec = WorkerAgentSpec;
260
261impl WorkerAgentSpec {
262 pub fn new(
264 kind: WorkerAgentKind,
265 name: impl Into<String>,
266 description: impl Into<String>,
267 ) -> Self {
268 Self {
269 name: name.into(),
270 description: description.into(),
271 kind,
272 hidden: false,
273 permissions: None,
274 model: None,
275 prompt: None,
276 max_steps: None,
277 confirmation_inheritance: None,
278 }
279 }
280
281 pub fn read_only(name: impl Into<String>, description: impl Into<String>) -> Self {
283 Self::new(WorkerAgentKind::ReadOnly, name, description)
284 }
285
286 pub fn planner(name: impl Into<String>, description: impl Into<String>) -> Self {
288 Self::new(WorkerAgentKind::Planner, name, description)
289 }
290
291 pub fn implementer(name: impl Into<String>, description: impl Into<String>) -> Self {
293 Self::new(WorkerAgentKind::Implementer, name, description)
294 }
295
296 pub fn verifier(name: impl Into<String>, description: impl Into<String>) -> Self {
298 Self::new(WorkerAgentKind::Verifier, name, description)
299 }
300
301 pub fn reviewer(name: impl Into<String>, description: impl Into<String>) -> Self {
303 Self::new(WorkerAgentKind::Reviewer, name, description)
304 }
305
306 pub fn custom(name: impl Into<String>, description: impl Into<String>) -> Self {
308 Self::new(WorkerAgentKind::Custom, name, description)
309 }
310
311 pub fn hidden(mut self, hidden: bool) -> Self {
313 self.hidden = hidden;
314 self
315 }
316
317 pub fn with_permissions(mut self, permissions: PermissionPolicy) -> Self {
319 self.permissions = Some(permissions);
320 self
321 }
322
323 pub fn with_model(mut self, model: ModelConfig) -> Self {
325 self.model = Some(model);
326 self
327 }
328
329 pub fn with_model_ref(mut self, model_ref: impl AsRef<str>) -> Self {
331 self.model = Some(ModelConfig::from_model_ref(model_ref));
332 self
333 }
334
335 pub fn with_provider_model(
337 mut self,
338 provider: impl Into<String>,
339 model: impl Into<String>,
340 ) -> Self {
341 self.model = Some(ModelConfig::with_provider(provider, model));
342 self
343 }
344
345 pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
347 self.prompt = Some(prompt.into());
348 self
349 }
350
351 pub fn with_max_steps(mut self, max_steps: usize) -> Self {
353 self.max_steps = Some(max_steps);
354 self
355 }
356
357 pub fn with_confirmation(mut self, inheritance: ConfirmationInheritance) -> Self {
359 self.confirmation_inheritance = Some(inheritance);
360 self
361 }
362
363 pub fn into_agent_definition(self) -> AgentDefinition {
365 let mut agent = AgentDefinition::new(&self.name, &self.description)
366 .with_permissions(
367 self.permissions
368 .unwrap_or_else(|| self.kind.default_permissions()),
369 )
370 .with_max_steps(
371 self.max_steps
372 .unwrap_or_else(|| self.kind.default_max_steps()),
373 );
374
375 if self.hidden {
376 agent = agent.hidden();
377 }
378 if let Some(model) = self.model {
379 agent = agent.with_model(model);
380 }
381 if let Some(prompt) = self
382 .prompt
383 .or_else(|| self.kind.default_prompt().map(str::to_string))
384 {
385 agent = agent.with_prompt(&prompt);
386 }
387 if let Some(ci) = self.confirmation_inheritance {
388 agent = agent.with_confirmation(ci);
389 }
390 agent
391 }
392}
393
394impl From<WorkerAgentSpec> for AgentDefinition {
395 fn from(spec: WorkerAgentSpec) -> Self {
396 spec.into_agent_definition()
397 }
398}
399
400#[derive(Debug, Clone, Serialize, Deserialize)]
404pub struct AgentDefinition {
405 pub name: String,
407 pub description: String,
409 #[serde(default)]
411 pub native: bool,
412 #[serde(default)]
414 pub hidden: bool,
415 #[serde(default)]
417 pub permissions: PermissionPolicy,
418 #[serde(skip_serializing_if = "Option::is_none")]
420 pub model: Option<ModelConfig>,
421 #[serde(skip_serializing_if = "Option::is_none")]
423 pub prompt: Option<String>,
424 #[serde(skip_serializing_if = "Option::is_none")]
426 pub max_steps: Option<usize>,
427 #[serde(default, skip_serializing_if = "Option::is_none")]
430 pub confirmation_inheritance: Option<ConfirmationInheritance>,
431}
432
433impl AgentDefinition {
434 pub fn new(name: &str, description: &str) -> Self {
436 Self {
437 name: name.to_string(),
438 description: description.to_string(),
439 native: false,
440 hidden: false,
441 permissions: PermissionPolicy::default(),
442 model: None,
443 prompt: None,
444 max_steps: None,
445 confirmation_inheritance: None,
446 }
447 }
448
449 pub fn worker(spec: WorkerAgentSpec) -> Self {
451 spec.into_agent_definition()
452 }
453
454 pub fn native(mut self) -> Self {
456 self.native = true;
457 self
458 }
459
460 pub fn hidden(mut self) -> Self {
462 self.hidden = true;
463 self
464 }
465
466 pub fn with_permissions(mut self, permissions: PermissionPolicy) -> Self {
468 self.permissions = permissions;
469 self
470 }
471
472 pub fn with_model(mut self, model: ModelConfig) -> Self {
474 self.model = Some(model);
475 self
476 }
477
478 pub fn with_prompt(mut self, prompt: &str) -> Self {
480 self.prompt = Some(prompt.to_string());
481 self
482 }
483
484 pub fn with_max_steps(mut self, max_steps: usize) -> Self {
486 self.max_steps = Some(max_steps);
487 self
488 }
489
490 pub fn has_defined_permissions(&self) -> bool {
492 !self.permissions.allow.is_empty() || !self.permissions.deny.is_empty()
493 }
494
495 pub(crate) fn apply_to(&self, config: &mut crate::agent::AgentConfig) {
501 use std::sync::Arc;
502
503 if config.permission_checker.is_none() && self.has_defined_permissions() {
504 config.permission_checker =
505 Some(Arc::new(self.permissions.clone()) as Arc<dyn PermissionChecker>);
506 config.permission_policy = Some(self.permissions.clone());
507 }
508
509 if let Some(ref prompt) = self.prompt {
510 if config.prompt_slots.extra.is_none() {
511 config.prompt_slots.extra = Some(prompt.clone());
512 }
513 }
514
515 if let Some(max_steps) = self.max_steps {
516 if config.max_tool_rounds == crate::agent::MAX_TOOL_ROUNDS {
517 config.max_tool_rounds = max_steps;
518 }
519 }
520
521 if config.confirmation_manager.is_none() {
523 let inheritance = self.confirmation_inheritance.clone().unwrap_or_else(|| {
524 if self.has_defined_permissions() {
525 ConfirmationInheritance::AutoApprove
526 } else {
527 ConfirmationInheritance::DenyOnAsk
528 }
529 });
530 match inheritance {
531 ConfirmationInheritance::AutoApprove => {
532 config.confirmation_manager =
533 Some(Arc::new(crate::hitl::AutoApproveConfirmation));
534 }
535 ConfirmationInheritance::DenyOnAsk => { }
536 ConfirmationInheritance::InheritParent => { }
537 }
538 }
539 }
540
541 pub fn with_confirmation(mut self, inheritance: ConfirmationInheritance) -> Self {
543 self.confirmation_inheritance = Some(inheritance);
544 self
545 }
546}
547
548pub struct AgentRegistry {
553 agents: RwLock<HashMap<String, AgentDefinition>>,
554}
555
556impl Default for AgentRegistry {
557 fn default() -> Self {
558 Self::new()
559 }
560}
561
562impl AgentRegistry {
563 pub fn new() -> Self {
565 let registry = Self {
566 agents: RwLock::new(HashMap::new()),
567 };
568
569 for agent in builtin_agents() {
571 registry.register(agent);
572 }
573
574 registry
575 }
576
577 pub fn with_config(config: &CodeConfig) -> Self {
581 let registry = Self::new();
582
583 for dir in &config.agent_dirs {
585 let agents = load_agents_from_dir(dir);
586 for agent in agents {
587 tracing::info!("Loaded agent '{}' from {}", agent.name, dir.display());
588 registry.register(agent);
589 }
590 }
591
592 registry
593 }
594
595 pub fn register(&self, agent: AgentDefinition) {
597 let mut agents = write_or_recover(&self.agents);
598 tracing::debug!("Registering agent: {}", agent.name);
599 agents.insert(agent.name.clone(), agent);
600 }
601
602 pub fn register_worker(&self, spec: WorkerAgentSpec) -> AgentDefinition {
607 let agent = spec.into_agent_definition();
608 self.register(agent.clone());
609 agent
610 }
611
612 pub fn register_workers<I>(&self, specs: I) -> Vec<AgentDefinition>
614 where
615 I: IntoIterator<Item = WorkerAgentSpec>,
616 {
617 specs
618 .into_iter()
619 .map(|spec| self.register_worker(spec))
620 .collect()
621 }
622
623 pub fn unregister(&self, name: &str) -> bool {
627 let mut agents = write_or_recover(&self.agents);
628 agents.remove(name).is_some()
629 }
630
631 pub fn get(&self, name: &str) -> Option<AgentDefinition> {
633 let agents = read_or_recover(&self.agents);
634 agents.get(name).cloned()
635 }
636
637 pub fn list(&self) -> Vec<AgentDefinition> {
639 let agents = read_or_recover(&self.agents);
640 agents.values().cloned().collect()
641 }
642
643 pub fn list_visible(&self) -> Vec<AgentDefinition> {
645 let agents = read_or_recover(&self.agents);
646 agents.values().filter(|a| !a.hidden).cloned().collect()
647 }
648
649 pub fn exists(&self, name: &str) -> bool {
651 let agents = read_or_recover(&self.agents);
652 agents.contains_key(name)
653 }
654
655 pub fn len(&self) -> usize {
657 let agents = read_or_recover(&self.agents);
658 agents.len()
659 }
660
661 pub fn is_empty(&self) -> bool {
663 self.len() == 0
664 }
665}
666
667pub fn parse_agent_yaml(content: &str) -> anyhow::Result<AgentDefinition> {
676 let value: serde_yaml::Value = serde_yaml::from_str(content)
677 .map_err(|e| anyhow::anyhow!("Failed to parse agent YAML: {}", e))?;
678
679 parse_agent_yaml_value(value, "agent YAML")
680}
681
682fn parse_agent_yaml_value(
683 value: serde_yaml::Value,
684 context: &str,
685) -> anyhow::Result<AgentDefinition> {
686 if yaml_value_has_key(&value, "kind") {
687 let spec: WorkerAgentSpec = serde_yaml::from_value(value)
688 .map_err(|e| anyhow::anyhow!("Failed to parse worker {}: {}", context, e))?;
689 validate_agent_name(&spec.name)?;
690 return Ok(spec.into_agent_definition());
691 }
692
693 let agent: AgentDefinition = serde_yaml::from_value(value)
694 .map_err(|e| anyhow::anyhow!("Failed to parse {}: {}", context, e))?;
695 validate_agent_name(&agent.name)?;
696 Ok(agent)
697}
698
699fn parse_worker_yaml_value(
700 value: serde_yaml::Value,
701 context: &str,
702) -> anyhow::Result<WorkerAgentSpec> {
703 let spec: WorkerAgentSpec = serde_yaml::from_value(value)
704 .map_err(|e| anyhow::anyhow!("Failed to parse worker {}: {}", context, e))?;
705 validate_agent_name(&spec.name)?;
706 Ok(spec)
707}
708
709fn yaml_value_has_key(value: &serde_yaml::Value, key: &str) -> bool {
710 value
711 .as_mapping()
712 .map(|mapping| mapping.contains_key(serde_yaml::Value::String(key.to_string())))
713 .unwrap_or(false)
714}
715
716fn validate_agent_name(name: &str) -> anyhow::Result<()> {
717 if name.trim().is_empty() {
718 return Err(anyhow::anyhow!("Agent name is required"));
719 }
720 Ok(())
721}
722
723pub fn parse_agent_md(content: &str) -> anyhow::Result<AgentDefinition> {
727 let parts: Vec<&str> = content.splitn(3, "---").collect();
729
730 if parts.len() < 3 {
731 return Err(anyhow::anyhow!(
732 "Invalid markdown format: missing YAML frontmatter"
733 ));
734 }
735
736 let frontmatter = parts[1].trim();
737 let body = parts[2].trim();
738
739 let value: serde_yaml::Value = serde_yaml::from_str(frontmatter)
741 .map_err(|e| anyhow::anyhow!("Failed to parse agent frontmatter: {}", e))?;
742
743 if yaml_value_has_key(&value, "kind") {
744 let mut spec = parse_worker_yaml_value(value, "frontmatter")?;
745 if spec.prompt.is_none() && !body.is_empty() {
746 spec.prompt = Some(body.to_string());
747 }
748 return Ok(spec.into_agent_definition());
749 }
750
751 let mut agent = parse_agent_yaml_value(value, "agent frontmatter")?;
752
753 if agent.prompt.is_none() && !body.is_empty() {
755 agent.prompt = Some(body.to_string());
756 }
757
758 Ok(agent)
759}
760
761pub fn load_agents_from_dir(dir: &Path) -> Vec<AgentDefinition> {
766 let mut agents = Vec::new();
767
768 let Ok(entries) = std::fs::read_dir(dir) else {
769 tracing::warn!("Failed to read agent directory: {}", dir.display());
770 return agents;
771 };
772
773 for entry in entries.flatten() {
774 let path = entry.path();
775
776 if !path.is_file() {
778 continue;
779 }
780
781 let Some(ext) = path.extension().and_then(|e| e.to_str()) else {
782 continue;
783 };
784
785 let Ok(content) = std::fs::read_to_string(&path) else {
787 tracing::warn!("Failed to read agent file: {}", path.display());
788 continue;
789 };
790
791 let result = match ext {
793 "yaml" | "yml" => parse_agent_yaml(&content),
794 "md" => parse_agent_md(&content),
795 _ => continue,
796 };
797
798 match result {
799 Ok(agent) => {
800 tracing::debug!("Loaded agent '{}' from {}", agent.name, path.display());
801 agents.push(agent);
802 }
803 Err(e) => {
804 tracing::warn!("Failed to parse agent file {}: {}", path.display(), e);
805 }
806 }
807 }
808
809 agents
810}
811
812pub fn builtin_agents() -> Vec<AgentDefinition> {
814 vec![
815 AgentDefinition::new(
817 "explore",
818 "Fast codebase exploration agent. Use for searching files, reading code, \
819 and understanding codebase structure. Read-only operations only.",
820 )
821 .native()
822 .with_permissions(explore_permissions())
823 .with_max_steps(20)
824 .with_prompt(EXPLORE_PROMPT),
825 AgentDefinition::new(
827 "general",
828 "General-purpose agent for multi-step task execution. Can read, write, \
829 and execute commands.",
830 )
831 .native()
832 .with_permissions(general_permissions())
833 .with_max_steps(50),
834 AgentDefinition::new(
836 "plan",
837 "Planning agent for designing implementation approaches. Read-only access \
838 to explore codebase and create plans.",
839 )
840 .native()
841 .with_permissions(plan_permissions())
842 .with_max_steps(30)
843 .with_prompt(PLAN_PROMPT),
844 AgentDefinition::new(
846 "verification",
847 "Verification agent for adversarial validation. Prefer real checks, \
848 reproductions, and regression testing over code reading alone.",
849 )
850 .native()
851 .with_permissions(verification_permissions())
852 .with_max_steps(30)
853 .with_prompt(VERIFICATION_PROMPT),
854 AgentDefinition::new(
856 "review",
857 "Code review agent focused on correctness, regressions, security, \
858 maintainability, and clear findings.",
859 )
860 .native()
861 .with_permissions(review_permissions())
862 .with_max_steps(25)
863 .with_prompt(REVIEW_PROMPT),
864 ]
865}
866
867fn explore_permissions() -> PermissionPolicy {
873 PermissionPolicy::new()
874 .allow_all(&["read", "grep", "glob", "ls"])
875 .deny_all(&["write", "edit", "task"])
876 .allow("Bash(ls:*)")
877 .allow("Bash(cat:*)")
878 .allow("Bash(head:*)")
879 .allow("Bash(tail:*)")
880 .allow("Bash(find:*)")
881 .allow("Bash(wc:*)")
882 .deny("Bash(rm:*)")
883 .deny("Bash(mv:*)")
884 .deny("Bash(cp:*)")
885}
886
887fn general_permissions() -> PermissionPolicy {
889 PermissionPolicy::new()
890 .allow_all(&["read", "write", "edit", "grep", "glob", "ls", "bash"])
891 .deny("task")
892}
893
894fn plan_permissions() -> PermissionPolicy {
896 PermissionPolicy::new()
897 .allow_all(&["read", "grep", "glob", "ls"])
898 .deny_all(&["write", "edit", "bash", "task"])
899}
900
901fn verification_permissions() -> PermissionPolicy {
903 PermissionPolicy::new()
904 .allow_all(&["read", "grep", "glob", "ls", "bash"])
905 .deny_all(&["write", "edit", "task"])
906}
907
908fn review_permissions() -> PermissionPolicy {
910 PermissionPolicy::new()
911 .allow_all(&["read", "grep", "glob", "ls", "bash"])
912 .deny_all(&["write", "edit", "task"])
913}
914
915const EXPLORE_PROMPT: &str = crate::prompts::AGENT_EXPLORE;
920
921const PLAN_PROMPT: &str = crate::prompts::AGENT_PLAN;
922
923const VERIFICATION_PROMPT: &str = crate::prompts::AGENT_VERIFICATION;
924
925const REVIEW_PROMPT: &str = crate::prompts::AGENT_CODE_REVIEW;
926
927#[cfg(test)]
932mod tests {
933 use super::*;
934
935 #[test]
936 fn test_agent_definition_builder() {
937 let agent = AgentDefinition::new("test", "Test agent")
938 .native()
939 .hidden()
940 .with_max_steps(10);
941
942 assert_eq!(agent.name, "test");
943 assert_eq!(agent.description, "Test agent");
944 assert!(agent.native);
945 assert!(agent.hidden);
946 assert_eq!(agent.max_steps, Some(10));
947 }
948
949 #[test]
950 fn test_agent_registry_new() {
951 let registry = AgentRegistry::new();
952
953 assert!(registry.exists("explore"));
955 assert!(registry.exists("general"));
956 assert!(registry.exists("plan"));
957 assert!(registry.exists("verification"));
958 assert!(registry.exists("review"));
959 assert_eq!(registry.len(), 5);
960 }
961
962 #[test]
963 fn test_agent_registry_get() {
964 let registry = AgentRegistry::new();
965
966 let explore = registry.get("explore").unwrap();
967 assert_eq!(explore.name, "explore");
968 assert!(explore.native);
969 assert!(!explore.hidden);
970
971 assert!(registry.get("nonexistent").is_none());
972 }
973
974 #[test]
975 fn test_agent_registry_register_unregister() {
976 let registry = AgentRegistry::new();
977 let initial_count = registry.len();
978
979 let custom = AgentDefinition::new("custom", "Custom agent");
981 registry.register(custom);
982 assert_eq!(registry.len(), initial_count + 1);
983 assert!(registry.exists("custom"));
984
985 assert!(registry.unregister("custom"));
987 assert_eq!(registry.len(), initial_count);
988 assert!(!registry.exists("custom"));
989
990 assert!(!registry.unregister("nonexistent"));
992 }
993
994 #[test]
995 fn test_agent_registry_list_visible() {
996 let registry = AgentRegistry::new();
997
998 let visible = registry.list_visible();
999 let all = registry.list();
1000
1001 assert_eq!(visible.len(), all.len());
1002 assert!(visible.iter().all(|a| !a.hidden));
1003 }
1004
1005 #[test]
1006 fn test_builtin_agents() {
1007 let agents = builtin_agents();
1008
1009 let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
1011 assert!(names.contains(&"explore"));
1012 assert!(names.contains(&"general"));
1013 assert!(names.contains(&"plan"));
1014 assert!(names.contains(&"verification"));
1015 assert!(names.contains(&"review"));
1016
1017 let explore = agents.iter().find(|a| a.name == "explore").unwrap();
1019 assert!(!explore.permissions.deny.is_empty());
1020 }
1021
1022 #[test]
1027 fn test_parse_agent_yaml() {
1028 let yaml = r#"
1029name: test-agent
1030description: A test agent
1031hidden: false
1032max_steps: 20
1033"#;
1034 let agent = parse_agent_yaml(yaml).unwrap();
1035 assert_eq!(agent.name, "test-agent");
1036 assert_eq!(agent.description, "A test agent");
1037 assert!(!agent.hidden);
1038 assert_eq!(agent.max_steps, Some(20));
1039 }
1040
1041 #[test]
1042 fn test_parse_agent_yaml_with_permissions() {
1043 let yaml = r#"
1044name: restricted-agent
1045description: Agent with permissions
1046permissions:
1047 allow:
1048 - rule: read
1049 - rule: grep
1050 deny:
1051 - rule: write
1052"#;
1053 let agent = parse_agent_yaml(yaml).unwrap();
1054 assert_eq!(agent.name, "restricted-agent");
1055 assert_eq!(agent.permissions.allow.len(), 2);
1056 assert_eq!(agent.permissions.deny.len(), 1);
1057 assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
1059 assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
1060 assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
1061 }
1062
1063 #[test]
1064 fn test_parse_agent_yaml_with_plain_string_permissions() {
1065 let yaml = r#"
1067name: plain-agent
1068description: Agent with plain string permissions
1069permissions:
1070 allow:
1071 - read
1072 - grep
1073 - "Bash(cargo:*)"
1074 deny:
1075 - write
1076"#;
1077 let agent = parse_agent_yaml(yaml).unwrap();
1078 assert_eq!(agent.name, "plain-agent");
1079 assert_eq!(agent.permissions.allow.len(), 3);
1080 assert_eq!(agent.permissions.deny.len(), 1);
1081 assert!(agent.permissions.allow[0].matches("read", &serde_json::json!({})));
1083 assert!(agent.permissions.allow[1].matches("grep", &serde_json::json!({})));
1084 assert!(agent.permissions.allow[2]
1085 .matches("Bash", &serde_json::json!({"command": "cargo build"})));
1086 assert!(agent.permissions.deny[0].matches("write", &serde_json::json!({})));
1087 }
1088
1089 #[test]
1090 fn test_parse_worker_agent_yaml_uses_cattle_defaults() {
1091 let yaml = r#"
1092name: frontend-fixer
1093description: Disposable frontend implementer
1094kind: implementer
1095max_steps: 7
1096"#;
1097 let agent = parse_agent_yaml(yaml).unwrap();
1098
1099 assert_eq!(agent.name, "frontend-fixer");
1100 assert_eq!(agent.max_steps, Some(7));
1101 assert!(agent
1102 .permissions
1103 .allow
1104 .iter()
1105 .any(|r| r.matches("write", &serde_json::json!({}))));
1106 assert!(agent
1107 .permissions
1108 .deny
1109 .iter()
1110 .any(|r| r.matches("task", &serde_json::json!({}))));
1111 }
1112
1113 #[test]
1114 fn test_parse_agent_yaml_missing_name() {
1115 let yaml = r#"
1116description: Agent without name
1117"#;
1118 let result = parse_agent_yaml(yaml);
1119 assert!(result.is_err());
1120 }
1121
1122 #[test]
1123 fn test_parse_agent_md() {
1124 let md = r#"---
1125name: md-agent
1126description: Agent from markdown
1127max_steps: 15
1128---
1129# System Prompt
1130
1131You are a helpful agent.
1132Do your best work.
1133"#;
1134 let agent = parse_agent_md(md).unwrap();
1135 assert_eq!(agent.name, "md-agent");
1136 assert_eq!(agent.description, "Agent from markdown");
1137 assert_eq!(agent.max_steps, Some(15));
1138 assert!(agent.prompt.is_some());
1139 assert!(agent.prompt.unwrap().contains("helpful agent"));
1140 }
1141
1142 #[test]
1143 fn test_parse_agent_md_with_prompt_in_frontmatter() {
1144 let md = r#"---
1145name: prompt-agent
1146description: Agent with prompt in frontmatter
1147prompt: "Frontmatter prompt"
1148---
1149Body content that should be ignored
1150"#;
1151 let agent = parse_agent_md(md).unwrap();
1152 assert_eq!(agent.prompt.unwrap(), "Frontmatter prompt");
1153 }
1154
1155 #[test]
1156 fn test_parse_worker_agent_md_uses_body_prompt() {
1157 let md = r#"---
1158name: review-cow
1159description: Disposable review worker
1160kind: reviewer
1161---
1162Review only the staged diff and return prioritized findings.
1163"#;
1164 let agent = parse_agent_md(md).unwrap();
1165
1166 assert_eq!(agent.name, "review-cow");
1167 assert_eq!(
1168 agent.prompt.as_deref(),
1169 Some("Review only the staged diff and return prioritized findings.")
1170 );
1171 assert!(agent
1172 .permissions
1173 .deny
1174 .iter()
1175 .any(|r| r.matches("write", &serde_json::json!({}))));
1176 }
1177
1178 #[test]
1179 fn test_parse_agent_md_missing_frontmatter() {
1180 let md = "Just markdown without frontmatter";
1181 let result = parse_agent_md(md);
1182 assert!(result.is_err());
1183 }
1184
1185 #[test]
1186 fn test_load_agents_from_dir() {
1187 let temp_dir = tempfile::tempdir().unwrap();
1188
1189 std::fs::write(
1191 temp_dir.path().join("agent1.yaml"),
1192 r#"
1193name: yaml-agent
1194description: Agent from YAML file
1195"#,
1196 )
1197 .unwrap();
1198
1199 std::fs::write(
1201 temp_dir.path().join("agent2.md"),
1202 r#"---
1203name: md-agent
1204description: Agent from Markdown file
1205---
1206System prompt here
1207"#,
1208 )
1209 .unwrap();
1210
1211 std::fs::write(temp_dir.path().join("invalid.yaml"), "not: valid: yaml: [").unwrap();
1213
1214 std::fs::write(temp_dir.path().join("readme.txt"), "Just a text file").unwrap();
1216
1217 let agents = load_agents_from_dir(temp_dir.path());
1218 assert_eq!(agents.len(), 2);
1219
1220 let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
1221 assert!(names.contains(&"yaml-agent"));
1222 assert!(names.contains(&"md-agent"));
1223 }
1224
1225 #[test]
1226 fn test_load_agents_from_nonexistent_dir() {
1227 let agents = load_agents_from_dir(std::path::Path::new("/nonexistent/dir"));
1228 assert!(agents.is_empty());
1229 }
1230
1231 #[test]
1232 fn test_registry_with_config() {
1233 let temp_dir = tempfile::tempdir().unwrap();
1234
1235 std::fs::write(
1237 temp_dir.path().join("custom.yaml"),
1238 r#"
1239name: custom-agent
1240description: Custom agent from config
1241"#,
1242 )
1243 .unwrap();
1244
1245 let config = CodeConfig::new().add_agent_dir(temp_dir.path());
1246 let registry = AgentRegistry::with_config(&config);
1247
1248 assert!(registry.exists("explore"));
1250 assert!(registry.exists("custom-agent"));
1251 assert_eq!(registry.len(), 6); }
1253
1254 #[test]
1255 fn test_agent_definition_with_model() {
1256 let model = ModelConfig {
1257 model: "claude-3-5-sonnet".to_string(),
1258 provider: Some("anthropic".to_string()),
1259 };
1260 let agent = AgentDefinition::new("test", "Test").with_model(model);
1261 assert!(agent.model.is_some());
1262 assert_eq!(agent.model.unwrap().provider, Some("anthropic".to_string()));
1263 }
1264
1265 #[test]
1266 fn test_model_config_from_model_ref() {
1267 let model = ModelConfig::from_model_ref("openai/gpt-4o");
1268 assert_eq!(model.provider.as_deref(), Some("openai"));
1269 assert_eq!(model.model, "gpt-4o");
1270 assert_eq!(model.model_ref(), "openai/gpt-4o");
1271
1272 let inherited = ModelConfig::from_model_ref("claude-sonnet");
1273 assert_eq!(inherited.provider, None);
1274 assert_eq!(inherited.model_ref(), "claude-sonnet");
1275 }
1276
1277 #[test]
1278 fn test_worker_agent_kind_from_str_accepts_aliases() {
1279 assert_eq!(
1280 "explore".parse::<WorkerAgentKind>().unwrap(),
1281 WorkerAgentKind::ReadOnly
1282 );
1283 assert_eq!(
1284 "general".parse::<WorkerAgentKind>().unwrap(),
1285 WorkerAgentKind::Implementer
1286 );
1287 assert!("unknown".parse::<WorkerAgentKind>().is_err());
1288 }
1289
1290 #[test]
1291 fn worker_spec_implementer_creates_cattle_agent_definition() {
1292 let agent = WorkerAgentSpec::implementer("frontend-fixer", "Fix frontend issues")
1293 .with_prompt("Focus on small, verified patches.")
1294 .with_provider_model("anthropic", "claude-sonnet")
1295 .with_max_steps(12)
1296 .into_agent_definition();
1297
1298 assert_eq!(agent.name, "frontend-fixer");
1299 assert_eq!(agent.max_steps, Some(12));
1300 assert_eq!(
1301 agent.prompt.as_deref(),
1302 Some("Focus on small, verified patches.")
1303 );
1304 assert_eq!(agent.model.unwrap().provider.as_deref(), Some("anthropic"));
1305 assert!(agent
1306 .permissions
1307 .allow
1308 .iter()
1309 .any(|r| r.matches("write", &serde_json::json!({}))));
1310 assert!(agent
1311 .permissions
1312 .deny
1313 .iter()
1314 .any(|r| r.matches("task", &serde_json::json!({}))));
1315 }
1316
1317 #[test]
1318 fn worker_spec_read_only_uses_safe_defaults() {
1319 let agent = WorkerAgentSpec::read_only("scanner", "Scan repository")
1320 .hidden(true)
1321 .into_agent_definition();
1322
1323 assert!(agent.hidden);
1324 assert_eq!(agent.max_steps, Some(20));
1325 assert!(agent.prompt.is_some());
1326 assert!(agent
1327 .permissions
1328 .allow
1329 .iter()
1330 .any(|r| r.matches("read", &serde_json::json!({}))));
1331 assert!(agent
1332 .permissions
1333 .deny
1334 .iter()
1335 .any(|r| r.matches("write", &serde_json::json!({}))));
1336 }
1337
1338 #[test]
1339 fn registry_register_worker_returns_and_stores_definition() {
1340 let registry = AgentRegistry::new();
1341 let agent =
1342 registry.register_worker(WorkerAgentSpec::custom("strict-worker", "Strict worker"));
1343
1344 assert_eq!(agent.name, "strict-worker");
1345 assert!(registry.exists("strict-worker"));
1346 assert_eq!(
1347 agent
1348 .permissions
1349 .check("bash", &serde_json::json!({"command":"echo hi"})),
1350 crate::permissions::PermissionDecision::Ask
1351 );
1352 }
1353
1354 #[test]
1355 fn registry_register_workers_batches_cattle_specs() {
1356 let registry = AgentRegistry::new();
1357 let agents = registry.register_workers([
1358 WorkerAgentSpec::planner("planner-cow", "Plan work"),
1359 WorkerAgentSpec::verifier("verify-cow", "Verify work"),
1360 ]);
1361
1362 assert_eq!(agents.len(), 2);
1363 assert!(registry.exists("planner-cow"));
1364 assert!(registry.exists("verify-cow"));
1365 }
1366
1367 #[test]
1368 fn test_agent_registry_default() {
1369 let registry = AgentRegistry::default();
1370 assert!(!registry.is_empty());
1371 assert_eq!(registry.len(), 5);
1372 }
1373
1374 #[test]
1375 fn test_agent_registry_is_empty() {
1376 let registry = AgentRegistry {
1377 agents: RwLock::new(HashMap::new()),
1378 };
1379 assert!(registry.is_empty());
1380 assert_eq!(registry.len(), 0);
1381 }
1382
1383 #[test]
1384 fn test_apply_to_sets_permissions() {
1385 use crate::agent::AgentConfig;
1386 use crate::permissions::PermissionDecision;
1387
1388 let def = AgentDefinition::new("writer", "Write files")
1389 .with_permissions(PermissionPolicy::new().allow("write(*)"));
1390
1391 let mut config = AgentConfig::default();
1392 assert!(config.permission_checker.is_none());
1393
1394 def.apply_to(&mut config);
1395
1396 assert!(config.permission_checker.is_some());
1397 assert!(config.permission_policy.is_some());
1398 let checker = config.permission_checker.unwrap();
1399 assert_eq!(
1400 checker.check(
1401 "write",
1402 &serde_json::json!({"file_path": "x.txt", "content": "hi"})
1403 ),
1404 PermissionDecision::Allow
1405 );
1406 }
1407
1408 #[test]
1409 fn test_apply_to_sets_prompt() {
1410 use crate::agent::AgentConfig;
1411
1412 let def = AgentDefinition::new("helper", "Help").with_prompt("Be helpful.");
1413 let mut config = AgentConfig::default();
1414
1415 def.apply_to(&mut config);
1416
1417 assert_eq!(config.prompt_slots.extra.as_deref(), Some("Be helpful."));
1418 }
1419
1420 #[test]
1421 fn test_apply_to_sets_max_steps() {
1422 use crate::agent::AgentConfig;
1423
1424 let def = AgentDefinition::new("fast", "Fast agent").with_max_steps(7);
1425 let mut config = AgentConfig::default();
1426
1427 def.apply_to(&mut config);
1428
1429 assert_eq!(config.max_tool_rounds, 7);
1430 }
1431
1432 #[test]
1433 fn test_apply_to_respects_host_overrides() {
1434 use crate::agent::AgentConfig;
1435
1436 let def = AgentDefinition::new("agent", "Agent")
1437 .with_permissions(PermissionPolicy::new().allow("bash(*)"))
1438 .with_prompt("Agent prompt.")
1439 .with_max_steps(10);
1440
1441 let mut config = AgentConfig::default();
1442 config.prompt_slots.extra = Some("Host prompt.".to_string());
1443 config.max_tool_rounds = 25;
1444 config.permission_checker = Some(std::sync::Arc::new(PermissionPolicy::new().allow("*")));
1445
1446 def.apply_to(&mut config);
1447
1448 assert_eq!(config.prompt_slots.extra.as_deref(), Some("Host prompt."));
1450 assert_eq!(config.max_tool_rounds, 25);
1451 }
1452
1453 #[test]
1454 fn test_apply_to_skips_empty_permissions() {
1455 use crate::agent::AgentConfig;
1456
1457 let def = AgentDefinition::new("empty", "No permissions");
1458 let mut config = AgentConfig::default();
1459
1460 def.apply_to(&mut config);
1461
1462 assert!(config.permission_checker.is_none());
1463 assert!(config.permission_policy.is_none());
1464 }
1465}