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
419impl cordis::Service for AgentRegistry {
420    fn name(&self) -> &'static str { "agent_registry" }
421    fn init(&self, _ctx: &std::sync::Arc<cordis::Context>) -> cordis::ServiceInitFuture<'_> {
422        Box::pin(async { Ok(None) })
423    }
424    fn check(&self) -> bool { true }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use ares_llm::ProviderConfig;
431    use ares_tools::Tool;
432    use std::collections::HashMap;
433
434    struct MapToon(HashMap<String, AgentConfig>);
435    impl ToonAgents for MapToon {
436        fn get(&self, name: &str) -> Option<AgentConfig> {
437            self.0.get(name).cloned()
438        }
439        fn names(&self) -> Vec<String> {
440            self.0.keys().cloned().collect()
441        }
442    }
443
444    fn empty_tools() -> Arc<Tools> {
445        Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()))
446    }
447
448    fn create_test_agent_map() -> HashMap<String, AgentConfig> {
449        HashMap::new()
450    }
451
452    fn create_test_provider_registry() -> Arc<ProviderRegistry> {
453        let mut registry = ProviderRegistry::new();
454        registry.register_provider(
455            "ollama-local",
456            ProviderConfig::OpenAI {
457                api_key_env: "TEST_KEY".to_string(),
458                api_base: "https://test.example.com/v1".to_string(),
459                default_model: "ministral-3:3b".to_string(),
460            },
461        );
462        registry.register_model(
463            "default",
464            ares_llm::ModelConfig {
465                provider: "ollama-local".to_string(),
466                model: "ministral-3:3b".to_string(),
467                temperature: 0.7,
468                max_tokens: 512,
469            },
470        );
471        Arc::new(registry)
472    }
473
474    #[test]
475    fn intersect_agent_tools_defaults_to_deny() {
476        let tenant_tools = vec!["calendar".to_string(), "search".to_string()];
477        let allowed = intersect_agent_tools_with_tenant_allowlist(None, &tenant_tools);
478        assert!(allowed.is_empty());
479    }
480
481    #[test]
482    fn intersect_agent_tools_keeps_only_configured_and_tenant_allowed() {
483        let agent_tools = vec!["calendar".to_string(), "sql".to_string()];
484        let tenant_tools = vec!["calendar".to_string(), "search".to_string()];
485        let allowed =
486            intersect_agent_tools_with_tenant_allowlist(Some(&agent_tools), &tenant_tools);
487        assert_eq!(allowed, vec!["calendar".to_string()]);
488    }
489
490    #[test]
491    fn test_type_to_name() {
492        assert_eq!(AgentRegistry::type_to_name(&AgentType::Router), "router");
493        assert_eq!(AgentRegistry::type_to_name(&AgentType::Product), "product");
494        assert_eq!(AgentRegistry::type_to_name(&AgentType::HR), "hr");
495        assert_eq!(AgentRegistry::type_to_name(&AgentType::Invoice), "invoice");
496        assert_eq!(AgentRegistry::type_to_name(&AgentType::Sales), "sales");
497        assert_eq!(AgentRegistry::type_to_name(&AgentType::Finance), "finance");
498        assert_eq!(
499            AgentRegistry::type_to_name(&AgentType::Orchestrator),
500            "orchestrator"
501        );
502    }
503
504    #[test]
505    fn test_registry_register_and_get() {
506        let provider_registry = create_test_provider_registry();
507        let tools = empty_tools();
508        let mut registry = AgentRegistry::new(provider_registry, tools);
509
510        let config = AgentConfig {
511            model: "default".to_string(),
512            system_prompt: Some("Test prompt".to_string()),
513            tools: vec![],
514            max_tool_iterations: 5,
515            parallel_tools: false,
516            extra: HashMap::new(),
517            allowed_tools: None,
518            compaction_enabled: None,
519        };
520
521        registry.register("test-agent", config);
522
523        assert!(registry.has_agent("test-agent"));
524        assert!(!registry.has_agent("nonexistent"));
525        assert!(registry.get_config("test-agent").is_some());
526        assert!(registry.get_config("nonexistent").is_none());
527    }
528
529    #[test]
530    fn test_registry_agent_names() {
531        let provider_registry = create_test_provider_registry();
532        let tools = empty_tools();
533        let mut registry = AgentRegistry::new(provider_registry, tools);
534
535        registry.register(
536            "agent1",
537            AgentConfig {
538                model: "default".to_string(),
539                system_prompt: None,
540                tools: vec![],
541                max_tool_iterations: 10,
542                parallel_tools: false,
543                extra: HashMap::new(),
544                allowed_tools: None,
545                compaction_enabled: None,
546            },
547        );
548
549        registry.register(
550            "agent2",
551            AgentConfig {
552                model: "default".to_string(),
553                system_prompt: None,
554                tools: vec![],
555                max_tool_iterations: 10,
556                parallel_tools: false,
557                extra: HashMap::new(),
558                allowed_tools: None,
559                compaction_enabled: None,
560            },
561        );
562
563        let names = registry.agent_names();
564        assert_eq!(names.len(), 2);
565        assert!(names.contains(&"agent1".to_string()));
566        assert!(names.contains(&"agent2".to_string()));
567    }
568
569    #[test]
570    fn test_registry_get_agent_model() {
571        let provider_registry = create_test_provider_registry();
572        let tools = empty_tools();
573        let mut registry = AgentRegistry::new(provider_registry, tools);
574
575        registry.register(
576            "test",
577            AgentConfig {
578                model: "default".to_string(),
579                system_prompt: None,
580                tools: vec![],
581                max_tool_iterations: 10,
582                parallel_tools: false,
583                extra: HashMap::new(),
584                allowed_tools: None,
585                compaction_enabled: None,
586            },
587        );
588
589        assert_eq!(
590            registry.get_agent_model("test"),
591            Some("default".to_string())
592        );
593        assert_eq!(registry.get_agent_model("nonexistent"), None);
594    }
595
596    #[test]
597    fn test_registry_get_agent_tools() {
598        let provider_registry = create_test_provider_registry();
599        let tools = empty_tools();
600        let mut registry = AgentRegistry::new(provider_registry, tools);
601
602        registry.register(
603            "with_tools",
604            AgentConfig {
605                model: "default".to_string(),
606                system_prompt: None,
607                tools: vec!["calculator".to_string(), "web_search".to_string()],
608                max_tool_iterations: 10,
609                parallel_tools: false,
610                extra: HashMap::new(),
611                allowed_tools: None,
612                compaction_enabled: None,
613            },
614        );
615
616        registry.register(
617            "no_tools",
618            AgentConfig {
619                model: "default".to_string(),
620                system_prompt: None,
621                tools: vec![],
622                max_tool_iterations: 10,
623                parallel_tools: false,
624                extra: HashMap::new(),
625                allowed_tools: None,
626                compaction_enabled: None,
627            },
628        );
629
630        let tools = registry.get_agent_tools("with_tools");
631        assert_eq!(tools.len(), 2);
632        assert!(tools.contains(&"calculator".to_string()));
633
634        let no_tools = registry.get_agent_tools("no_tools");
635        assert!(no_tools.is_empty());
636    }
637
638    #[test]
639    fn test_builder_build_without_provider_registry() {
640        let result = AgentRegistryBuilder::new()
641            .with_tools(empty_tools())
642            .build();
643
644        assert!(result.is_err());
645    }
646
647    #[test]
648    fn test_builder_build_success() {
649        let provider_registry = create_test_provider_registry();
650
651        let result = AgentRegistryBuilder::new()
652            .with_provider_registry(provider_registry)
653            .with_agent(
654                "test",
655                AgentConfig {
656                    model: "default".to_string(),
657                    system_prompt: Some("Test".to_string()),
658                    tools: vec![],
659                    max_tool_iterations: 5,
660                    parallel_tools: false,
661                    extra: HashMap::new(),
662                    allowed_tools: None,
663                    compaction_enabled: None,
664                },
665            )
666            .build();
667
668        assert!(result.is_ok());
669        let _registry = result.unwrap();
670    }
671
672    // ============================================================
673    //  type_to_name for ALL variants (including Custom)
674    // ============================================================
675
676    #[test]
677    fn test_type_to_name_custom() {
678        assert_eq!(
679            AgentRegistry::type_to_name(&AgentType::Custom("my-custom-agent".to_string())),
680            "my-custom-agent"
681        );
682    }
683
684    // ============================================================
685    //  from_config() - build AgentRegistry from AresConfig
686    // ============================================================
687
688    #[test]
689    fn test_registry_from_config() {
690        let provider_registry = create_test_provider_registry();
691        let tools = empty_tools();
692
693        let overlay_config = {
694            let mut map = HashMap::new();
695            map.insert(
696                "toml-agent".to_string(),
697                AgentConfig {
698                    model: "default".to_string(),
699                    system_prompt: Some("TOML prompt".to_string()),
700                    tools: vec!["calculator".to_string()],
701                    max_tool_iterations: 5,
702                    parallel_tools: false,
703                    extra: HashMap::new(),
704                    allowed_tools: None,
705                    compaction_enabled: None,
706                },
707            );
708            map
709        };
710
711        let registry = AgentRegistry::from_config(overlay_config, provider_registry, tools);
712
713        assert!(registry.has_agent("toml-agent"));
714        assert!(registry.get_config("toml-agent").is_some());
715        assert_eq!(
716            registry.get_agent_model("toml-agent"),
717            Some("default".to_string())
718        );
719    }
720
721    // ============================================================
722    //  AgentRegistryBuilder::from_config(...).with_provider_registry(...).build()
723    // ============================================================
724
725    #[test]
726    fn test_builder_from_config_with_provider() {
727        let provider_registry = create_test_provider_registry();
728
729        let overlay_config = {
730            let mut map = HashMap::new();
731            map.insert(
732                "builder-agent".to_string(),
733                AgentConfig {
734                    model: "default".to_string(),
735                    system_prompt: Some("Builder prompt".to_string()),
736                    tools: vec![],
737                    max_tool_iterations: 10,
738                    parallel_tools: true,
739                    extra: HashMap::new(),
740                    allowed_tools: None,
741                    compaction_enabled: None,
742                },
743            );
744            map
745        };
746
747        let result = AgentRegistryBuilder::new()
748            .from_config(overlay_config)
749            .with_provider_registry(provider_registry)
750            .build();
751
752        assert!(result.is_ok());
753        let registry = result.unwrap();
754        assert!(registry.has_agent("builder-agent"));
755    }
756
757    // ============================================================
758    //  AgentRegistryBuilder::with_tools(...) and default tools
759    // ============================================================
760
761    #[test]
762    fn test_builder_with_tools_explicit() {
763        let provider_registry = create_test_provider_registry();
764        let custom_tools = empty_tools();
765
766        let result = AgentRegistryBuilder::new()
767            .with_provider_registry(provider_registry)
768            .with_tools(custom_tools.clone())
769            .with_agent(
770                "tool-agent",
771                AgentConfig {
772                    model: "default".to_string(),
773                    system_prompt: None,
774                    tools: vec!["calculator".to_string()],
775                    max_tool_iterations: 5,
776                    parallel_tools: false,
777                    extra: HashMap::new(),
778                    allowed_tools: None,
779                    compaction_enabled: None,
780                },
781            )
782            .build();
783
784        assert!(result.is_ok());
785        let registry = result.unwrap();
786        assert!(registry.has_agent("tool-agent"));
787    }
788
789    #[test]
790    fn test_builder_default_tools() {
791        let provider_registry = create_test_provider_registry();
792
793        // Build without calling with_tools - should use default empty Tools
794        let result = AgentRegistryBuilder::new()
795            .with_provider_registry(provider_registry)
796            .with_agent(
797                "no-tool-registry-agent",
798                AgentConfig {
799                    model: "default".to_string(),
800                    system_prompt: Some("No explicit tool registry".to_string()),
801                    tools: vec![],
802                    max_tool_iterations: 5,
803                    parallel_tools: false,
804                    extra: HashMap::new(),
805                    allowed_tools: None,
806                    compaction_enabled: None,
807                },
808            )
809            .build();
810
811        assert!(result.is_ok());
812        let registry = result.unwrap();
813        assert!(registry.has_agent("no-tool-registry-agent"));
814    }
815
816    // ============================================================
817    //  AgentRegistryBuilder::default() returns equivalent-to-new
818    // ============================================================
819
820    #[test]
821    fn test_builder_default_is_new() {
822        let from_new = AgentRegistryBuilder::new();
823        let from_default = AgentRegistryBuilder::default();
824
825        // Neither has provider_registry set, so both should error on build
826        assert!(from_new.build().is_err());
827        assert!(from_default.build().is_err());
828    }
829
830    // ============================================================
831    //  get_agent_system_prompt - TOML branch (Some) and None-when-missing
832    // ============================================================
833
834    #[test]
835    fn test_get_agent_system_prompt_toml_some() {
836        let provider_registry = create_test_provider_registry();
837        let tools = empty_tools();
838        let mut registry = AgentRegistry::new(provider_registry, tools);
839
840        registry.register(
841            "has-prompt",
842            AgentConfig {
843                model: "default".to_string(),
844                system_prompt: Some("My System Prompt".to_string()),
845                tools: vec![],
846                max_tool_iterations: 5,
847                parallel_tools: false,
848                extra: HashMap::new(),
849                allowed_tools: None,
850                compaction_enabled: None,
851            },
852        );
853
854        assert_eq!(
855            registry.get_agent_system_prompt("has-prompt"),
856            Some("My System Prompt".to_string())
857        );
858    }
859
860    #[test]
861    fn test_get_agent_system_prompt_toml_none() {
862        let provider_registry = create_test_provider_registry();
863        let tools = empty_tools();
864        let mut registry = AgentRegistry::new(provider_registry, tools);
865
866        registry.register(
867            "no-prompt",
868            AgentConfig {
869                model: "default".to_string(),
870                system_prompt: None,
871                tools: vec![],
872                max_tool_iterations: 5,
873                parallel_tools: false,
874                extra: HashMap::new(),
875                allowed_tools: None,
876                compaction_enabled: None,
877            },
878        );
879
880        assert_eq!(registry.get_agent_system_prompt("no-prompt"), None);
881        assert_eq!(registry.get_agent_system_prompt("missing"), None);
882    }
883
884    // ============================================================
885    //  create_agent(name) NOT-found path - Configuration error
886    // ============================================================
887
888    #[tokio::test]
889    async fn test_create_agent_not_found() {
890        let provider_registry = create_test_provider_registry();
891        let tools = empty_tools();
892        let registry = AgentRegistry::new(provider_registry, tools);
893
894        let result = registry.create_agent("nonexistent-agent").await;
895
896        assert!(result.is_err());
897        if let Err(AppError::Configuration(msg)) = result {
898            assert!(msg.contains("nonexistent-agent"));
899            assert!(msg.contains("not found"));
900        } else {
901            panic!("Expected Configuration error");
902        }
903    }
904
905    // ============================================================
906    //  create_agent_by_type - error path for unregistered type
907    // ============================================================
908
909    #[tokio::test]
910    async fn test_create_agent_by_type_not_registered() {
911        let provider_registry = create_test_provider_registry();
912        let tools = empty_tools();
913        let registry = AgentRegistry::new(provider_registry, tools);
914
915        // AgentType::Custom("custom-unregistered".to_string()) returns name "custom-unregistered"
916        let result = registry
917            .create_agent_by_type(AgentType::Custom("custom-unregistered".to_string()))
918            .await;
919
920        assert!(result.is_err());
921        if let Err(AppError::Configuration(msg)) = result {
922            assert!(msg.contains("custom-unregistered"));
923        } else {
924            panic!("Expected Configuration error");
925        }
926    }
927
928    // ============================================================
929    //  TOON-backed tests
930    // ============================================================
931
932    fn create_test_dynamic_config_manager() -> Arc<dyn ToonAgents> {
933        let mut agents = HashMap::new();
934        agents.insert(
935            "toon-only-agent".to_string(),
936            AgentConfig {
937                model: "default".to_string(),
938                system_prompt: Some("TOON system prompt".to_string()),
939                tools: vec!["calculator".to_string()],
940                allowed_tools: None,
941                max_tool_iterations: 10,
942                parallel_tools: false,
943                extra: HashMap::new(),
944                compaction_enabled: None,
945            },
946        );
947        Arc::new(MapToon(agents))
948    }
949
950    #[test]
951    fn test_with_dynamic_config() {
952        let provider_registry = create_test_provider_registry();
953        let tools = empty_tools();
954        let dcm = create_test_dynamic_config_manager();
955
956        let registry = AgentRegistry::with_dynamic_config(
957            create_test_agent_map(),
958            provider_registry,
959            tools,
960            dcm.clone(),
961        );
962
963        assert!(registry.get_toon_config("toon-only-agent").is_some());
964        assert!(registry.has_agent("toon-only-agent"));
965    }
966
967    #[test]
968    fn test_set_dynamic_config() {
969        let provider_registry = create_test_provider_registry();
970        let tools = empty_tools();
971        let dcm = create_test_dynamic_config_manager();
972
973        let mut registry = AgentRegistry::new(provider_registry, tools);
974        assert!(!registry.has_agent("toon-only-agent")); // Not set yet
975
976        registry.set_dynamic_config(dcm);
977        assert!(registry.has_agent("toon-only-agent"));
978        assert!(registry.get_toon_config("toon-only-agent").is_some());
979    }
980
981    #[test]
982    fn test_toon_merge_agent_names_no_duplicates() {
983        let provider_registry = create_test_provider_registry();
984        let tools = empty_tools();
985        let dcm = create_test_dynamic_config_manager();
986
987        let mut registry = AgentRegistry::new(provider_registry, tools);
988        registry.set_dynamic_config(dcm);
989
990        // Register a TOML agent with the same name as TOON - TOML should take precedence
991        registry.register(
992            "toon-only-agent",
993            AgentConfig {
994                model: "default".to_string(),
995                system_prompt: Some("TOML override".to_string()),
996                tools: vec![],
997                max_tool_iterations: 5,
998                parallel_tools: false,
999                extra: HashMap::new(),
1000                allowed_tools: None,
1001                compaction_enabled: None,
1002            },
1003        );
1004
1005        let names = registry.agent_names();
1006        let count = names.iter().filter(|n| *n == "toon-only-agent").count();
1007        assert_eq!(count, 1, "Should not have duplicate agent names");
1008        assert!(registry.has_agent("toon-only-agent"));
1009    }
1010
1011    #[test]
1012    fn test_get_agent_model_toon() {
1013        let provider_registry = create_test_provider_registry();
1014        let tools = empty_tools();
1015        let dcm = create_test_dynamic_config_manager();
1016
1017        let registry = AgentRegistry::with_dynamic_config(
1018            create_test_agent_map(),
1019            provider_registry,
1020            tools,
1021            dcm,
1022        );
1023
1024        // TOON-only agent
1025        assert_eq!(
1026            registry.get_agent_model("toon-only-agent"),
1027            Some("default".to_string())
1028        );
1029    }
1030
1031    #[test]
1032    fn test_get_agent_tools_toon() {
1033        let provider_registry = create_test_provider_registry();
1034        let tools = empty_tools();
1035        let dcm = create_test_dynamic_config_manager();
1036
1037        let registry = AgentRegistry::with_dynamic_config(
1038            create_test_agent_map(),
1039            provider_registry,
1040            tools,
1041            dcm,
1042        );
1043
1044        // TOON-only agent
1045        let tools = registry.get_agent_tools("toon-only-agent");
1046        assert_eq!(tools, vec!["calculator".to_string()]);
1047    }
1048
1049    #[test]
1050    fn test_get_agent_system_prompt_toon() {
1051        let provider_registry = create_test_provider_registry();
1052        let tools = empty_tools();
1053        let dcm = create_test_dynamic_config_manager();
1054
1055        let registry = AgentRegistry::with_dynamic_config(
1056            create_test_agent_map(),
1057            provider_registry,
1058            tools,
1059            dcm,
1060        );
1061
1062        // TOON-only agent system prompt
1063        assert_eq!(
1064            registry.get_agent_system_prompt("toon-only-agent"),
1065            Some("TOON system prompt".to_string())
1066        );
1067    }
1068}
1069
ares_agent/registry.rs

ares_agent/
registry.rs
