ares_agent/
registry.rs

1//! Agent Registry for managing configurable agents
2//!
3//! This module provides a registry for creating and managing agents
4//! based on both TOML and TOON configuration.
5//!
6//! ## Configuration Precedence
7//!
8//! When looking up an agent by name:
9//! 1. TOML config (`ares.toml` [agents.*]) is checked first
10//! 2. TOON config (`config/agents/*.toon`) is checked second
11//!
12//! This allows TOML to override TOON configs for specific deployments.
13
14use crate::configurable::ConfigurableAgent;
15use crate::{AgentConfig, ToonAgents};
16use ares_llm::ProviderRegistry;
17use ares_tools::Tools;
18use ares_types::types::{AgentType, AppError, Result};
19use std::collections::HashMap;
20use std::sync::Arc;
21
22/// Registry for managing agent configurations and creating agent instances
23///
24/// Supports both TOML-based static config and TOON-based dynamic config.
25/// TOML configs take precedence over TOON configs when both exist.
26pub struct AgentRegistry {
27    /// Agent configurations from TOML keyed by name
28    configs: HashMap<String, AgentConfig>,
29    /// Provider registry for creating LLM clients
30    provider_registry: Arc<ProviderRegistry>,
31    /// Tools capability shared across agents
32    tools: Arc<Tools>,
33    /// Optional TOON-based dynamic agent lookup (Overlay keeps this live)
34    dynamic_config: Option<Arc<dyn ToonAgents>>,
35}
36
37fn intersect_agent_tools_with_tenant_allowlist(
38    agent_tools: Option<&[String]>,
39    tenant_allowed_tools: &[String],
40) -> Vec<String> {
41    agent_tools
42        .unwrap_or(&[])
43        .iter()
44        .filter(|tool| tenant_allowed_tools.contains(*tool))
45        .cloned()
46        .collect()
47}
48
49impl AgentRegistry {
50    /// Create a new agent registry
51    pub fn new(provider_registry: Arc<ProviderRegistry>, tools: Arc<Tools>) -> Self {
52        Self {
53            configs: HashMap::new(),
54            provider_registry,
55            tools,
56            dynamic_config: None,
57        }
58    }
59
60    /// Create an agent registry from TOML agent configs
61    pub fn from_config(
62        agents: HashMap<String, AgentConfig>,
63        provider_registry: Arc<ProviderRegistry>,
64        tools: Arc<Tools>,
65    ) -> Self {
66        Self {
67            configs: agents,
68            provider_registry,
69            tools,
70            dynamic_config: None,
71        }
72    }
73
74    /// Create an agent registry with both TOML and TOON config support
75    pub fn with_dynamic_config(
76        agents: HashMap<String, AgentConfig>,
77        provider_registry: Arc<ProviderRegistry>,
78        tools: Arc<Tools>,
79        dynamic_config: Arc<dyn ToonAgents>,
80    ) -> Self {
81        Self {
82            configs: agents,
83            provider_registry,
84            tools,
85            dynamic_config: Some(dynamic_config),
86        }
87    }
88
89    /// Set the dynamic TOON agent lookup
90    pub fn set_dynamic_config(&mut self, dynamic_config: Arc<dyn ToonAgents>) {
91        self.dynamic_config = Some(dynamic_config);
92    }
93
94    /// Register an agent configuration
95    pub fn register(&mut self, name: &str, config: AgentConfig) {
96        self.configs.insert(name.to_string(), config);
97    }
98
99    /// Get an agent configuration by name (TOML only)
100    ///
101    /// Note: For lookups that include TOON, use `get_config_any` instead.
102    pub fn get_config(&self, name: &str) -> Option<&AgentConfig> {
103        self.configs.get(name)
104    }
105
106    /// Get an agent configuration by name from TOML or TOON.
107    pub fn get_config_any(&self, name: &str) -> Option<AgentConfig> {
108        self.configs.get(name).cloned().or_else(|| self.get_toon_config(name))
109    }
110
111    /// Get TOON agent config by name (already converted to AgentConfig)
112    pub fn get_toon_config(&self, name: &str) -> Option<AgentConfig> {
113        self.dynamic_config.as_ref().and_then(|dc| dc.get(name))
114    }
115
116    /// Check if an agent exists in TOML config
117    fn has_toml_agent(&self, name: &str) -> bool {
118        self.configs.contains_key(name)
119    }
120
121    /// Check if an agent exists in TOON config
122    fn has_toon_agent(&self, name: &str) -> bool {
123        self.dynamic_config
124            .as_ref()
125            .map(|dc| dc.get(name).is_some())
126            .unwrap_or(false)
127    }
128
129    /// Get all agent names (from both TOML and TOON)
130    pub fn agent_names(&self) -> Vec<String> {
131        let mut names: Vec<String> = self.configs.keys().cloned().collect();
132
133        // Add TOON agent names that aren't already in TOML
134        if let Some(dc) = &self.dynamic_config {
135            for name in dc.names() {
136                if !names.contains(&name) {
137                    names.push(name);
138                }
139            }
140        }
141
142        names
143    }
144
145    /// Check if an agent exists (in either TOML or TOON config)
146    pub fn has_agent(&self, name: &str) -> bool {
147        self.has_toml_agent(name) || self.has_toon_agent(name)
148    }
149
150    /// Create an agent instance by name
151    ///
152    /// This creates a new ConfigurableAgent with the appropriate LLM client
153    /// and tool registry based on the agent's configuration.
154    ///
155    /// Lookup order:
156    /// 1. TOML config (`ares.toml` [agents.*])
157    /// 2. TOON config (`config/agents/*.toon`)
158    pub async fn create_agent(&self, name: &str) -> Result<ConfigurableAgent> {
159        // First check TOML config
160        if let Some(config) = self.get_config(name) {
161            return self.create_agent_from_config(name, config).await;
162        }
163
164        // Then check TOON config
165        if let Some(config) = self.get_toon_config(name) {
166            return self.create_agent_from_config(name, &config).await;
167        }
168
169        Err(AppError::Configuration(format!(
170            "Agent '{}' not found in TOML or TOON configuration",
171            name
172        )))
173    }
174
175    /// Create an agent instance from an explicit configuration
176    pub async fn create_agent_from_config(
177        &self,
178        name: &str,
179        config: &AgentConfig,
180    ) -> Result<ConfigurableAgent> {
181        // Create the LLM client for this agent's model
182        let llm = self
183            .provider_registry
184            .create_client_for_model(&config.model)
185            .await?;
186        let provider_name = self
187            .provider_registry
188            .get_model(&config.model)
189            .map(|model| model.provider.clone())
190            .unwrap_or_else(|| config.model.clone());
191
192        let mut agent = ConfigurableAgent::new_with_provider(
193            name,
194            config,
195            llm,
196            None,
197            provider_name,
198        );
199        agent.set_tools(Arc::clone(&self.tools));
200        Ok(agent)
201    }
202
203    /// Create an agent instance from an explicit configuration with tier
204    /// resolution and fallback providers wired in.
205    #[cfg(feature = "postgres")]
206    pub async fn create_agent_from_config_with_fallbacks(
207        &self,
208        name: &str,
209        config: &AgentConfig,
210        tenant_id: &str,
211        pool: &sqlx::PgPool,
212        fleet_secrets: &ares_store::FleetSecrets,
213    ) -> Result<ConfigurableAgent> {
214        let chain = self
215            .provider_registry
216            .resolve_with_fallback(&config.model, tenant_id, pool, fleet_secrets)
217            .await?;
218
219        let allowlist_store = ares_store::tenant_allowlist::TenantAllowlistStore::new(pool);
220        for resolved in &chain {
221            if !allowlist_store
222                .is_model_allowed(tenant_id, &resolved.model_name)
223                .await
224                .map_err(|e| AppError::Auth(format!("Failed to check model allowlist: {}", e)))?
225            {
226                return Err(AppError::Auth(format!(
227                    "Model '{}' is not allowed for this tenant",
228                    resolved.model_name
229                )));
230            }
231        }
232
233        let mut iter = chain.into_iter();
234        let primary = iter.next().ok_or_else(|| {
235            AppError::Configuration(format!(
236                "No provider resolved for model/tier '{}'",
237                config.model
238            ))
239        })?;
240
241        let primary_provider_name = primary.provider_name.clone();
242
243        let llm = self
244            .provider_registry
245            .create_client_for_resolved_provider(&primary)
246            .await
247            .map_err(|e| {
248                AppError::Configuration(format!(
249                    "Failed to construct primary provider '{}' for model '{}': {}",
250                    primary.provider_name, primary.model_name, e
251                ))
252            })?;
253
254        let mut fallback_llms = Vec::new();
255        for fallback in iter {
256            let fallback_provider_name = fallback.provider_name.clone();
257            let client = self
258                .provider_registry
259                .create_client_for_resolved_provider(&fallback)
260                .await
261                .map_err(|e| {
262                    AppError::Configuration(format!(
263                        "Failed to construct fallback provider '{}' for model '{}': {}",
264                        fallback.provider_name, fallback.model_name, e
265                    ))
266                })?;
267            fallback_llms.push((fallback_provider_name, client));
268        }
269
270        let mut agent = ConfigurableAgent::new_with_provider(
271            name,
272            config,
273            llm,
274            None,
275            primary_provider_name,
276        );
277        agent.set_tools(Arc::clone(&self.tools));
278        agent.set_fallback_llms_with_providers(fallback_llms);
279
280        // --- tenant allowlist enforcement ---
281        // Tool enforcement: intersect config allowed_tools with tenant allowlist.
282        // Empty tenant allowlist rows mean default-deny.
283        let db_tools = allowlist_store
284            .list_tools(tenant_id)
285            .await
286            .map_err(|e| AppError::Auth(format!("Failed to check tool allowlist: {}", e)))?;
287        let db_tool_names: Vec<String> = db_tools.iter().map(|t| t.tool_name.clone()).collect();
288        let new_allowed =
289            intersect_agent_tools_with_tenant_allowlist(agent.allowed_tools(), &db_tool_names);
290        agent.set_allowed_tools(Some(new_allowed));
291
292        Ok(agent)
293    }
294
295    /// Create an agent instance for a specific AgentType
296    pub async fn create_agent_by_type(&self, agent_type: AgentType) -> Result<ConfigurableAgent> {
297        let name = Self::type_to_name(&agent_type);
298        self.create_agent(name).await
299    }
300
301    /// Convert AgentType to agent name
302    pub fn type_to_name(agent_type: &AgentType) -> &str {
303        agent_type.as_str()
304    }
305
306    /// Get the model name for an agent (checks both TOML and TOON)
307    pub fn get_agent_model(&self, name: &str) -> Option<String> {
308        // Check TOML first
309        if let Some(config) = self.configs.get(name) {
310            return Some(config.model.clone());
311        }
312        // Check TOON
313        self.get_toon_config(name).map(|c| c.model)
314    }
315
316    /// Get the tools for an agent (checks both TOML and TOON).
317    /// Returns the explicit allowed_tools list, falling back to the legacy
318    /// tools field.  An empty result means no configured tools; runtime
319    /// execution remains deny-by-default.
320    pub fn get_agent_tools(&self, name: &str) -> Vec<String> {
321        // Check TOML first
322        if let Some(config) = self.configs.get(name) {
323            return config
324                .allowed_tools
325                .clone()
326                .unwrap_or_else(|| config.tools.clone());
327        }
328        // Check TOON
329        self.get_toon_config(name)
330            .map(|c| c.allowed_tools.unwrap_or(c.tools))
331            .unwrap_or_default()
332    }
333
334    /// Get the system prompt for an agent (checks both TOML and TOON)
335    pub fn get_agent_system_prompt(&self, name: &str) -> Option<String> {
336        // Check TOML first
337        if let Some(config) = self.configs.get(name) {
338            return config.system_prompt.clone();
339        }
340        // Check TOON
341        self.get_toon_config(name).and_then(|c| c.system_prompt)
342    }
343}
344
345/// Builder for creating AgentRegistry with fluent API
346pub struct AgentRegistryBuilder {
347    configs: HashMap<String, AgentConfig>,
348    provider_registry: Option<Arc<ProviderRegistry>>,
349    tools: Option<Arc<Tools>>,
350    dynamic_config: Option<Arc<dyn ToonAgents>>,
351}
352
353impl AgentRegistryBuilder {
354    /// Create a new builder
355    pub fn new() -> Self {
356        Self {
357            configs: HashMap::new(),
358            provider_registry: None,
359            tools: None,
360            dynamic_config: None,
361        }
362    }
363
364    /// Set the provider registry
365    pub fn with_provider_registry(mut self, registry: Arc<ProviderRegistry>) -> Self {
366        self.provider_registry = Some(registry);
367        self
368    }
369
370    /// Set the unified tools capability
371    pub fn with_tools(mut self, tools: Arc<Tools>) -> Self {
372        self.tools = Some(tools);
373        self
374    }
375
376    /// Set the dynamic config manager for TOON support
377    pub fn with_dynamic_config(mut self, dynamic_config: Arc<dyn ToonAgents>) -> Self {
378        self.dynamic_config = Some(dynamic_config);
379        self
380    }
381
382    /// Add an agent configuration
383    pub fn with_agent(mut self, name: &str, config: AgentConfig) -> Self {
384        self.configs.insert(name.to_string(), config);
385        self
386    }
387
388    /// Load agent configurations from TOML agent map
389    pub fn from_config(mut self, agents: HashMap<String, AgentConfig>) -> Self {
390        self.configs = agents;
391        self
392    }
393
394    /// Build the AgentRegistry
395    pub fn build(self) -> Result<AgentRegistry> {
396        let provider_registry = self.provider_registry.ok_or_else(|| {
397            AppError::Configuration("ProviderRegistry is required for AgentRegistry".into())
398        })?;
399
400        let tools = self.tools.unwrap_or_else(|| {
401            Arc::new(Tools::from_static(std::iter::empty::<Arc<dyn ares_tools::Tool>>()))
402        });
403
404        Ok(AgentRegistry {
405            configs: self.configs,
406            provider_registry,
407            tools,
408            dynamic_config: self.dynamic_config,
409        })
410    }
411}
412
413impl Default for AgentRegistryBuilder {
414    fn default() -> Self {
415        Self::new()
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use ares_llm::ProviderConfig;
423    use ares_tools::Tool;
424    use std::collections::HashMap;
425
426    struct MapToon(HashMap<String, AgentConfig>);
427    impl ToonAgents for MapToon {
428        fn get(&self, name: &str) -> Option<AgentConfig> {
429            self.0.get(name).cloned()
430        }
431        fn names(&self) -> Vec<String> {
432            self.0.keys().cloned().collect()
433        }
434    }
435
436    fn empty_tools() -> Arc<Tools> {
437        Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()))
438    }
439
440    fn create_test_agent_map() -> HashMap<String, AgentConfig> {
441        HashMap::new()
442    }
443
444    fn create_test_provider_registry() -> Arc<ProviderRegistry> {
445        let mut registry = ProviderRegistry::new();
446        registry.register_provider(
447            "ollama-local",
448            ProviderConfig::OpenAI {
449                api_key_env: "TEST_KEY".to_string(),
450                api_base: "https://test.example.com/v1".to_string(),
451                default_model: "ministral-3:3b".to_string(),
452            },
453        );
454        registry.register_model(
455            "default",
456            ares_llm::ModelConfig {
457                provider: "ollama-local".to_string(),
458                model: "ministral-3:3b".to_string(),
459                temperature: 0.7,
460                max_tokens: 512,
461            },
462        );
463        Arc::new(registry)
464    }
465
466    #[test]
467    fn intersect_agent_tools_defaults_to_deny() {
468        let tenant_tools = vec!["calendar".to_string(), "search".to_string()];
469        let allowed = intersect_agent_tools_with_tenant_allowlist(None, &tenant_tools);
470        assert!(allowed.is_empty());
471    }
472
473    #[test]
474    fn intersect_agent_tools_keeps_only_configured_and_tenant_allowed() {
475        let agent_tools = vec!["calendar".to_string(), "sql".to_string()];
476        let tenant_tools = vec!["calendar".to_string(), "search".to_string()];
477        let allowed =
478            intersect_agent_tools_with_tenant_allowlist(Some(&agent_tools), &tenant_tools);
479        assert_eq!(allowed, vec!["calendar".to_string()]);
480    }
481
482    #[test]
483    fn test_type_to_name() {
484        assert_eq!(AgentRegistry::type_to_name(&AgentType::Router), "router");
485        assert_eq!(AgentRegistry::type_to_name(&AgentType::Product), "product");
486        assert_eq!(AgentRegistry::type_to_name(&AgentType::HR), "hr");
487        assert_eq!(AgentRegistry::type_to_name(&AgentType::Invoice), "invoice");
488        assert_eq!(AgentRegistry::type_to_name(&AgentType::Sales), "sales");
489        assert_eq!(AgentRegistry::type_to_name(&AgentType::Finance), "finance");
490        assert_eq!(
491            AgentRegistry::type_to_name(&AgentType::Orchestrator),
492            "orchestrator"
493        );
494    }
495
496    #[test]
497    fn test_registry_register_and_get() {
498        let provider_registry = create_test_provider_registry();
499        let tools = empty_tools();
500        let mut registry = AgentRegistry::new(provider_registry, tools);
501
502        let config = AgentConfig {
503            model: "default".to_string(),
504            system_prompt: Some("Test prompt".to_string()),
505            tools: vec![],
506            max_tool_iterations: 5,
507            parallel_tools: false,
508            extra: HashMap::new(),
509            allowed_tools: None,
510            compaction_enabled: None,
511        };
512
513        registry.register("test-agent", config);
514
515        assert!(registry.has_agent("test-agent"));
516        assert!(!registry.has_agent("nonexistent"));
517        assert!(registry.get_config("test-agent").is_some());
518        assert!(registry.get_config("nonexistent").is_none());
519    }
520
521    #[test]
522    fn test_registry_agent_names() {
523        let provider_registry = create_test_provider_registry();
524        let tools = empty_tools();
525        let mut registry = AgentRegistry::new(provider_registry, tools);
526
527        registry.register(
528            "agent1",
529            AgentConfig {
530                model: "default".to_string(),
531                system_prompt: None,
532                tools: vec![],
533                max_tool_iterations: 10,
534                parallel_tools: false,
535                extra: HashMap::new(),
536                allowed_tools: None,
537                compaction_enabled: None,
538            },
539        );
540
541        registry.register(
542            "agent2",
543            AgentConfig {
544                model: "default".to_string(),
545                system_prompt: None,
546                tools: vec![],
547                max_tool_iterations: 10,
548                parallel_tools: false,
549                extra: HashMap::new(),
550                allowed_tools: None,
551                compaction_enabled: None,
552            },
553        );
554
555        let names = registry.agent_names();
556        assert_eq!(names.len(), 2);
557        assert!(names.contains(&"agent1".to_string()));
558        assert!(names.contains(&"agent2".to_string()));
559    }
560
561    #[test]
562    fn test_registry_get_agent_model() {
563        let provider_registry = create_test_provider_registry();
564        let tools = empty_tools();
565        let mut registry = AgentRegistry::new(provider_registry, tools);
566
567        registry.register(
568            "test",
569            AgentConfig {
570                model: "default".to_string(),
571                system_prompt: None,
572                tools: vec![],
573                max_tool_iterations: 10,
574                parallel_tools: false,
575                extra: HashMap::new(),
576                allowed_tools: None,
577                compaction_enabled: None,
578            },
579        );
580
581        assert_eq!(
582            registry.get_agent_model("test"),
583            Some("default".to_string())
584        );
585        assert_eq!(registry.get_agent_model("nonexistent"), None);
586    }
587
588    #[test]
589    fn test_registry_get_agent_tools() {
590        let provider_registry = create_test_provider_registry();
591        let tools = empty_tools();
592        let mut registry = AgentRegistry::new(provider_registry, tools);
593
594        registry.register(
595            "with_tools",
596            AgentConfig {
597                model: "default".to_string(),
598                system_prompt: None,
599                tools: vec!["calculator".to_string(), "web_search".to_string()],
600                max_tool_iterations: 10,
601                parallel_tools: false,
602                extra: HashMap::new(),
603                allowed_tools: None,
604                compaction_enabled: None,
605            },
606        );
607
608        registry.register(
609            "no_tools",
610            AgentConfig {
611                model: "default".to_string(),
612                system_prompt: None,
613                tools: vec![],
614                max_tool_iterations: 10,
615                parallel_tools: false,
616                extra: HashMap::new(),
617                allowed_tools: None,
618                compaction_enabled: None,
619            },
620        );
621
622        let tools = registry.get_agent_tools("with_tools");
623        assert_eq!(tools.len(), 2);
624        assert!(tools.contains(&"calculator".to_string()));
625
626        let no_tools = registry.get_agent_tools("no_tools");
627        assert!(no_tools.is_empty());
628    }
629
630    #[test]
631    fn test_builder_build_without_provider_registry() {
632        let result = AgentRegistryBuilder::new()
633            .with_tools(empty_tools())
634            .build();
635
636        assert!(result.is_err());
637    }
638
639    #[test]
640    fn test_builder_build_success() {
641        let provider_registry = create_test_provider_registry();
642
643        let result = AgentRegistryBuilder::new()
644            .with_provider_registry(provider_registry)
645            .with_agent(
646                "test",
647                AgentConfig {
648                    model: "default".to_string(),
649                    system_prompt: Some("Test".to_string()),
650                    tools: vec![],
651                    max_tool_iterations: 5,
652                    parallel_tools: false,
653                    extra: HashMap::new(),
654                    allowed_tools: None,
655                    compaction_enabled: None,
656                },
657            )
658            .build();
659
660        assert!(result.is_ok());
661        let _registry = result.unwrap();
662    }
663
664    // ============================================================
665    //  type_to_name for ALL variants (including Custom)
666    // ============================================================
667
668    #[test]
669    fn test_type_to_name_custom() {
670        assert_eq!(
671            AgentRegistry::type_to_name(&AgentType::Custom("my-custom-agent".to_string())),
672            "my-custom-agent"
673        );
674    }
675
676    // ============================================================
677    //  from_config() - build AgentRegistry from AresConfig
678    // ============================================================
679
680    #[test]
681    fn test_registry_from_config() {
682        let provider_registry = create_test_provider_registry();
683        let tools = empty_tools();
684
685        let overlay_config = {
686            let mut map = HashMap::new();
687            map.insert(
688                "toml-agent".to_string(),
689                AgentConfig {
690                    model: "default".to_string(),
691                    system_prompt: Some("TOML prompt".to_string()),
692                    tools: vec!["calculator".to_string()],
693                    max_tool_iterations: 5,
694                    parallel_tools: false,
695                    extra: HashMap::new(),
696                    allowed_tools: None,
697                    compaction_enabled: None,
698                },
699            );
700            map
701        };
702
703        let registry = AgentRegistry::from_config(overlay_config, provider_registry, tools);
704
705        assert!(registry.has_agent("toml-agent"));
706        assert!(registry.get_config("toml-agent").is_some());
707        assert_eq!(
708            registry.get_agent_model("toml-agent"),
709            Some("default".to_string())
710        );
711    }
712
713    // ============================================================
714    //  AgentRegistryBuilder::from_config(...).with_provider_registry(...).build()
715    // ============================================================
716
717    #[test]
718    fn test_builder_from_config_with_provider() {
719        let provider_registry = create_test_provider_registry();
720
721        let overlay_config = {
722            let mut map = HashMap::new();
723            map.insert(
724                "builder-agent".to_string(),
725                AgentConfig {
726                    model: "default".to_string(),
727                    system_prompt: Some("Builder prompt".to_string()),
728                    tools: vec![],
729                    max_tool_iterations: 10,
730                    parallel_tools: true,
731                    extra: HashMap::new(),
732                    allowed_tools: None,
733                    compaction_enabled: None,
734                },
735            );
736            map
737        };
738
739        let result = AgentRegistryBuilder::new()
740            .from_config(overlay_config)
741            .with_provider_registry(provider_registry)
742            .build();
743
744        assert!(result.is_ok());
745        let registry = result.unwrap();
746        assert!(registry.has_agent("builder-agent"));
747    }
748
749    // ============================================================
750    //  AgentRegistryBuilder::with_tools(...) and default tools
751    // ============================================================
752
753    #[test]
754    fn test_builder_with_tools_explicit() {
755        let provider_registry = create_test_provider_registry();
756        let custom_tools = empty_tools();
757
758        let result = AgentRegistryBuilder::new()
759            .with_provider_registry(provider_registry)
760            .with_tools(custom_tools.clone())
761            .with_agent(
762                "tool-agent",
763                AgentConfig {
764                    model: "default".to_string(),
765                    system_prompt: None,
766                    tools: vec!["calculator".to_string()],
767                    max_tool_iterations: 5,
768                    parallel_tools: false,
769                    extra: HashMap::new(),
770                    allowed_tools: None,
771                    compaction_enabled: None,
772                },
773            )
774            .build();
775
776        assert!(result.is_ok());
777        let registry = result.unwrap();
778        assert!(registry.has_agent("tool-agent"));
779    }
780
781    #[test]
782    fn test_builder_default_tools() {
783        let provider_registry = create_test_provider_registry();
784
785        // Build without calling with_tools - should use default empty Tools
786        let result = AgentRegistryBuilder::new()
787            .with_provider_registry(provider_registry)
788            .with_agent(
789                "no-tool-registry-agent",
790                AgentConfig {
791                    model: "default".to_string(),
792                    system_prompt: Some("No explicit tool registry".to_string()),
793                    tools: vec![],
794                    max_tool_iterations: 5,
795                    parallel_tools: false,
796                    extra: HashMap::new(),
797                    allowed_tools: None,
798                    compaction_enabled: None,
799                },
800            )
801            .build();
802
803        assert!(result.is_ok());
804        let registry = result.unwrap();
805        assert!(registry.has_agent("no-tool-registry-agent"));
806    }
807
808    // ============================================================
809    //  AgentRegistryBuilder::default() returns equivalent-to-new
810    // ============================================================
811
812    #[test]
813    fn test_builder_default_is_new() {
814        let from_new = AgentRegistryBuilder::new();
815        let from_default = AgentRegistryBuilder::default();
816
817        // Neither has provider_registry set, so both should error on build
818        assert!(from_new.build().is_err());
819        assert!(from_default.build().is_err());
820    }
821
822    // ============================================================
823    //  get_agent_system_prompt - TOML branch (Some) and None-when-missing
824    // ============================================================
825
826    #[test]
827    fn test_get_agent_system_prompt_toml_some() {
828        let provider_registry = create_test_provider_registry();
829        let tools = empty_tools();
830        let mut registry = AgentRegistry::new(provider_registry, tools);
831
832        registry.register(
833            "has-prompt",
834            AgentConfig {
835                model: "default".to_string(),
836                system_prompt: Some("My System Prompt".to_string()),
837                tools: vec![],
838                max_tool_iterations: 5,
839                parallel_tools: false,
840                extra: HashMap::new(),
841                allowed_tools: None,
842                compaction_enabled: None,
843            },
844        );
845
846        assert_eq!(
847            registry.get_agent_system_prompt("has-prompt"),
848            Some("My System Prompt".to_string())
849        );
850    }
851
852    #[test]
853    fn test_get_agent_system_prompt_toml_none() {
854        let provider_registry = create_test_provider_registry();
855        let tools = empty_tools();
856        let mut registry = AgentRegistry::new(provider_registry, tools);
857
858        registry.register(
859            "no-prompt",
860            AgentConfig {
861                model: "default".to_string(),
862                system_prompt: None,
863                tools: vec![],
864                max_tool_iterations: 5,
865                parallel_tools: false,
866                extra: HashMap::new(),
867                allowed_tools: None,
868                compaction_enabled: None,
869            },
870        );
871
872        assert_eq!(registry.get_agent_system_prompt("no-prompt"), None);
873        assert_eq!(registry.get_agent_system_prompt("missing"), None);
874    }
875
876    // ============================================================
877    //  create_agent(name) NOT-found path - Configuration error
878    // ============================================================
879
880    #[tokio::test]
881    async fn test_create_agent_not_found() {
882        let provider_registry = create_test_provider_registry();
883        let tools = empty_tools();
884        let registry = AgentRegistry::new(provider_registry, tools);
885
886        let result = registry.create_agent("nonexistent-agent").await;
887
888        assert!(result.is_err());
889        if let Err(AppError::Configuration(msg)) = result {
890            assert!(msg.contains("nonexistent-agent"));
891            assert!(msg.contains("not found"));
892        } else {
893            panic!("Expected Configuration error");
894        }
895    }
896
897    // ============================================================
898    //  create_agent_by_type - error path for unregistered type
899    // ============================================================
900
901    #[tokio::test]
902    async fn test_create_agent_by_type_not_registered() {
903        let provider_registry = create_test_provider_registry();
904        let tools = empty_tools();
905        let registry = AgentRegistry::new(provider_registry, tools);
906
907        // AgentType::Custom("custom-unregistered".to_string()) returns name "custom-unregistered"
908        let result = registry
909            .create_agent_by_type(AgentType::Custom("custom-unregistered".to_string()))
910            .await;
911
912        assert!(result.is_err());
913        if let Err(AppError::Configuration(msg)) = result {
914            assert!(msg.contains("custom-unregistered"));
915        } else {
916            panic!("Expected Configuration error");
917        }
918    }
919
920    // ============================================================
921    //  TOON-backed tests
922    // ============================================================
923
924    fn create_test_dynamic_config_manager() -> Arc<dyn ToonAgents> {
925        let mut agents = HashMap::new();
926        agents.insert(
927            "toon-only-agent".to_string(),
928            AgentConfig {
929                model: "default".to_string(),
930                system_prompt: Some("TOON system prompt".to_string()),
931                tools: vec!["calculator".to_string()],
932                allowed_tools: None,
933                max_tool_iterations: 10,
934                parallel_tools: false,
935                extra: HashMap::new(),
936                compaction_enabled: None,
937            },
938        );
939        Arc::new(MapToon(agents))
940    }
941
942    #[test]
943    fn test_with_dynamic_config() {
944        let provider_registry = create_test_provider_registry();
945        let tools = empty_tools();
946        let dcm = create_test_dynamic_config_manager();
947
948        let registry = AgentRegistry::with_dynamic_config(
949            create_test_agent_map(),
950            provider_registry,
951            tools,
952            dcm.clone(),
953        );
954
955        assert!(registry.get_toon_config("toon-only-agent").is_some());
956        assert!(registry.has_agent("toon-only-agent"));
957    }
958
959    #[test]
960    fn test_set_dynamic_config() {
961        let provider_registry = create_test_provider_registry();
962        let tools = empty_tools();
963        let dcm = create_test_dynamic_config_manager();
964
965        let mut registry = AgentRegistry::new(provider_registry, tools);
966        assert!(!registry.has_agent("toon-only-agent")); // Not set yet
967
968        registry.set_dynamic_config(dcm);
969        assert!(registry.has_agent("toon-only-agent"));
970        assert!(registry.get_toon_config("toon-only-agent").is_some());
971    }
972
973    #[test]
974    fn test_toon_merge_agent_names_no_duplicates() {
975        let provider_registry = create_test_provider_registry();
976        let tools = empty_tools();
977        let dcm = create_test_dynamic_config_manager();
978
979        let mut registry = AgentRegistry::new(provider_registry, tools);
980        registry.set_dynamic_config(dcm);
981
982        // Register a TOML agent with the same name as TOON - TOML should take precedence
983        registry.register(
984            "toon-only-agent",
985            AgentConfig {
986                model: "default".to_string(),
987                system_prompt: Some("TOML override".to_string()),
988                tools: vec![],
989                max_tool_iterations: 5,
990                parallel_tools: false,
991                extra: HashMap::new(),
992                allowed_tools: None,
993                compaction_enabled: None,
994            },
995        );
996
997        let names = registry.agent_names();
998        let count = names.iter().filter(|n| *n == "toon-only-agent").count();
999        assert_eq!(count, 1, "Should not have duplicate agent names");
1000        assert!(registry.has_agent("toon-only-agent"));
1001    }
1002
1003    #[test]
1004    fn test_get_agent_model_toon() {
1005        let provider_registry = create_test_provider_registry();
1006        let tools = empty_tools();
1007        let dcm = create_test_dynamic_config_manager();
1008
1009        let registry = AgentRegistry::with_dynamic_config(
1010            create_test_agent_map(),
1011            provider_registry,
1012            tools,
1013            dcm,
1014        );
1015
1016        // TOON-only agent
1017        assert_eq!(
1018            registry.get_agent_model("toon-only-agent"),
1019            Some("default".to_string())
1020        );
1021    }
1022
1023    #[test]
1024    fn test_get_agent_tools_toon() {
1025        let provider_registry = create_test_provider_registry();
1026        let tools = empty_tools();
1027        let dcm = create_test_dynamic_config_manager();
1028
1029        let registry = AgentRegistry::with_dynamic_config(
1030            create_test_agent_map(),
1031            provider_registry,
1032            tools,
1033            dcm,
1034        );
1035
1036        // TOON-only agent
1037        let tools = registry.get_agent_tools("toon-only-agent");
1038        assert_eq!(tools, vec!["calculator".to_string()]);
1039    }
1040
1041    #[test]
1042    fn test_get_agent_system_prompt_toon() {
1043        let provider_registry = create_test_provider_registry();
1044        let tools = empty_tools();
1045        let dcm = create_test_dynamic_config_manager();
1046
1047        let registry = AgentRegistry::with_dynamic_config(
1048            create_test_agent_map(),
1049            provider_registry,
1050            tools,
1051            dcm,
1052        );
1053
1054        // TOON-only agent system prompt
1055        assert_eq!(
1056            registry.get_agent_system_prompt("toon-only-agent"),
1057            Some("TOON system prompt".to_string())
1058        );
1059    }
1060}
1061
1062impl cordis::Service for AgentRegistry {
1063    fn name(&self) -> &'static str { "agent_registry" }
1064    fn init(&self, _ctx: &std::sync::Arc<cordis::Context>) -> cordis::ServiceInitFuture<'_> {
1065        Box::pin(async { Ok(None) })
1066    }
1067    fn check(&self) -> bool { true }
1068}
1069
ares_agent/registry.rs

ares_agent/
registry.rs
