1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::str::FromStr;
4use std::sync::Arc;
5
6use ai_agents_context::ContextManager;
7use ai_agents_core::{AgentError, AgentStorage, LLMFeature, LLMProvider, Result, Tool};
8use ai_agents_hitl::{ApprovalHandler, HITLEngine, RejectAllHandler};
9use ai_agents_hooks::{AgentHooks, CompositeHooks};
10use ai_agents_llm::providers::{ProviderType, UnifiedLLMProvider};
11use ai_agents_llm::{LLMRegistry, LLMRole};
12use ai_agents_memory::{
13 CompactingMemory, InMemoryStore, LLMSummarizer, Memory, NoopSummarizer, Summarizer,
14};
15use ai_agents_observability::{
16 ObservabilityConfig, ObservabilityHooks, ObservabilityManager, ObservedLLMProvider,
17 ObservedTool,
18};
19use ai_agents_process::ProcessProcessor;
20use ai_agents_reasoning::{ReasoningConfig, ReflectionConfig};
21use ai_agents_recovery::{MessageFilter, RecoveryManager};
22use ai_agents_relationships::{
23 RelationshipEvaluator, RelationshipEvaluatorTrait, RelationshipManager,
24};
25use ai_agents_skills::{SkillDefinition, SkillLoader};
26use ai_agents_state::{LLMTransitionEvaluator, StateMachine, TransitionEvaluator};
27use ai_agents_template::{TemplateInheritance, TemplateLoader, TemplateRenderer};
28use ai_agents_tools::mcp::view::MCPViewTool;
29use ai_agents_tools::mcp::wrapper::MCPWrapperTool;
30use ai_agents_tools::{ToolRegistry, ToolSecurityEngine, create_builtin_registry};
31
32use super::AgentInfo;
33use super::StreamingConfig;
34use super::runtime::{RuntimeAgent, ToolResourceLocks};
35use crate::spec::{AgentSpec, StorageConfig};
36
37fn feature_overrides_from_config(config: &crate::spec::LLMConfig) -> HashMap<LLMFeature, bool> {
38 let mut overrides = HashMap::new();
39 if let Some(enabled) = config.function_calling {
40 overrides.insert(LLMFeature::FunctionCalling, enabled);
41 }
42 if let Some(enabled) = config.vision {
43 overrides.insert(LLMFeature::Vision, enabled);
44 }
45 if let Some(enabled) = config.json_mode {
46 overrides.insert(LLMFeature::JsonMode, enabled);
47 }
48 overrides
49}
50
51fn model_by_alias_from_spec(spec: Option<&AgentSpec>) -> HashMap<String, String> {
53 spec.map(|spec| {
54 let mut models: HashMap<String, String> = spec
55 .llms
56 .iter()
57 .map(|(alias, config)| (alias.clone(), config.model.clone()))
58 .collect();
59 if let Some(config) = spec.llm.as_config() {
60 models.insert("default".to_string(), config.model.clone());
61 }
62 models
63 })
64 .unwrap_or_default()
65}
66
67fn wrap_registry_with_observability(
69 registry: LLMRegistry,
70 manager: Arc<ObservabilityManager>,
71 model_by_alias: &HashMap<String, String>,
72) -> LLMRegistry {
73 registry.map_providers(|alias, provider| {
74 let provider_name = provider.provider_name().to_string();
75 let model = model_by_alias
76 .get(alias)
77 .cloned()
78 .unwrap_or_else(|| alias.to_string());
79 Arc::new(ObservedLLMProvider::new(
80 provider,
81 Arc::clone(&manager),
82 Some(alias.to_string()),
83 provider_name,
84 model,
85 )) as Arc<dyn LLMProvider>
86 })
87}
88
89pub struct AgentBuilder {
90 spec: Option<AgentSpec>,
91 llm: Option<Arc<dyn LLMProvider>>,
92 llm_registry: Option<LLMRegistry>,
93 memory: Option<Arc<dyn Memory>>,
94 tools: Option<ToolRegistry>,
95 skills: Vec<SkillDefinition>,
96 skill_loader: Option<SkillLoader>,
97 yaml_dir: Option<PathBuf>,
98 system_prompt: Option<String>,
99 tools_prompt: Option<String>,
100 auto_tools_prompt: bool,
101 max_iterations: Option<u32>,
102 max_context_tokens: Option<u32>,
103 recovery_manager: Option<RecoveryManager>,
104 tool_security: Option<ToolSecurityEngine>,
105 process_processor: Option<ProcessProcessor>,
106 message_filters: HashMap<String, Arc<dyn MessageFilter>>,
107 context_manager: Option<Arc<ContextManager>>,
108 state_machine: Option<Arc<StateMachine>>,
109 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
110 hooks: Option<Arc<dyn AgentHooks>>,
111 hitl_engine: Option<HITLEngine>,
112 approval_handler: Option<Arc<dyn ApprovalHandler>>,
113 storage_config: Option<StorageConfig>,
114 storage: Option<Arc<dyn AgentStorage>>,
115 reasoning: Option<ReasoningConfig>,
116 reflection: Option<ReflectionConfig>,
117 streaming: Option<StreamingConfig>,
118 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
119 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
120 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
121 persona_templates: Option<Arc<ai_agents_persona::PersonaTemplateRegistry>>,
122 observability_manager: Option<Arc<ObservabilityManager>>,
123 resource_locks: Option<ToolResourceLocks>,
124 llm_registry_observed: bool,
125 skills_prepared: bool,
126 routing_frozen: Option<LLMRegistry>,
127 routing_dirty: bool,
128}
129
130impl AgentBuilder {
131 pub fn new() -> Self {
132 Self {
133 reasoning: None,
134 reflection: None,
135 spec: None,
136 llm: None,
137 llm_registry: None,
138 memory: None,
139 tools: None,
140 skills: Vec::new(),
141 skill_loader: None,
142 yaml_dir: None,
143 system_prompt: None,
144 tools_prompt: None,
145 auto_tools_prompt: true,
146 max_iterations: None,
147 max_context_tokens: None,
148 recovery_manager: None,
149 tool_security: None,
150 process_processor: None,
151 message_filters: HashMap::new(),
152 context_manager: None,
153 state_machine: None,
154 transition_evaluator: None,
155 hooks: None,
156 hitl_engine: None,
157 approval_handler: None,
158 storage_config: None,
159 storage: None,
160 streaming: None,
161 spawner: None,
162 spawner_registry: None,
163 persona_manager: None,
164 persona_templates: None,
165 observability_manager: None,
166 resource_locks: None,
167 llm_registry_observed: false,
168 skills_prepared: false,
169 routing_frozen: None,
170 routing_dirty: false,
171 }
172 }
173
174 pub fn from_spec(spec: AgentSpec) -> Self {
175 let system_prompt = spec.system_prompt.clone();
176 let max_iterations = Some(spec.max_iterations);
177 let max_context_tokens = Some(spec.max_context_tokens);
178 let reasoning = Some(spec.reasoning.clone());
179 let reflection = Some(spec.reflection.clone());
180
181 Self {
182 spec: Some(spec),
183 llm: None,
184 llm_registry: None,
185 memory: None,
186 tools: None,
187 skills: Vec::new(),
188 skill_loader: None,
189 yaml_dir: None,
190 system_prompt: Some(system_prompt),
191 tools_prompt: None,
192 auto_tools_prompt: true,
193 max_iterations,
194 max_context_tokens,
195 recovery_manager: None,
196 tool_security: None,
197 process_processor: None,
198 message_filters: HashMap::new(),
199 context_manager: None,
200 state_machine: None,
201 transition_evaluator: None,
202 hooks: None,
203 hitl_engine: None,
204 approval_handler: None,
205 storage_config: None,
206 storage: None,
207 reasoning,
208 reflection,
209 streaming: None,
210 spawner: None,
211 spawner_registry: None,
212 persona_manager: None,
213 persona_templates: None,
214 observability_manager: None,
215 resource_locks: None,
216 llm_registry_observed: false,
217 skills_prepared: false,
218 routing_frozen: None,
219 routing_dirty: false,
220 }
221 }
222
223 pub fn from_spec_with_base_dir(spec: AgentSpec, base_dir: impl Into<PathBuf>) -> Self {
225 let mut builder = Self::from_spec(spec);
226 builder.yaml_dir = Some(base_dir.into());
227 builder
228 }
229
230 pub fn from_yaml(yaml_content: &str) -> Result<Self> {
231 let spec = AgentSpec::from_yaml_strict(yaml_content)?;
232 spec.validate()?;
233 Ok(Self::from_spec(spec))
234 }
235
236 pub(crate) fn shared_resource_locks(&mut self) -> ToolResourceLocks {
237 self.resource_locks
238 .get_or_insert_with(|| Arc::new(parking_lot::RwLock::new(HashMap::new())))
239 .clone()
240 }
241
242 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
243 self.resource_locks = Some(locks);
244 self
245 }
246
247 pub fn from_yaml_file(path: impl AsRef<Path>) -> Result<Self> {
248 let path = path.as_ref();
249 let content = std::fs::read_to_string(path).map_err(AgentError::IoError)?;
250 let spec = AgentSpec::from_yaml_strict(&content)?;
251 spec.validate()?;
252 Ok(match path.parent() {
253 Some(parent) => Self::from_spec_with_base_dir(spec, parent),
254 None => Self::from_spec(spec),
255 })
256 }
257
258 pub fn from_template(template_name: &str) -> Result<Self> {
259 let loader = TemplateLoader::new();
260 Self::from_template_with_loader(template_name, &loader)
261 }
262
263 pub fn from_template_with_loader(template_name: &str, loader: &TemplateLoader) -> Result<Self> {
264 let renderer = TemplateRenderer::new();
265 let variables = loader.variables();
266
267 let load_and_render = |name: &str| -> Result<String> {
268 let content = loader.load_template(name)?;
269 renderer.render(&content, variables)
270 };
271
272 let rendered_root = load_and_render(template_name)?;
273 let processed = TemplateInheritance::process(&rendered_root, load_and_render)?;
274 let spec = AgentSpec::from_yaml_strict(&processed)?;
275 spec.validate()?;
276 Ok(Self::from_spec(spec))
277 }
278
279 pub fn auto_configure_llms(mut self) -> Result<Self> {
281 if self.routing_frozen.is_some() {
282 self.routing_dirty = true;
283 }
284 let spec = self
285 .spec
286 .as_ref()
287 .ok_or_else(|| AgentError::Config("Cannot auto-configure LLMs without spec".into()))?;
288
289 if !spec.llms.is_empty() {
290 let mut registry = LLMRegistry::new();
291
292 for (alias, config) in &spec.llms {
293 let provider_type = ProviderType::from_str(&config.provider)
294 .map_err(|e| AgentError::Config(e.to_string()))?;
295
296 let core_config = ai_agents_core::LLMConfig {
297 temperature: Some(config.temperature),
298 max_tokens: Some(config.max_tokens),
299 top_p: config.top_p,
300 top_k: None,
301 frequency_penalty: None,
302 presence_penalty: None,
303 stop_sequences: None,
304 timeout_seconds: config.timeout_seconds,
305 reasoning: config.reasoning,
306 reasoning_effort: config.reasoning_effort.clone(),
307 reasoning_budget_tokens: config.reasoning_budget_tokens,
308 extra: config.extra.clone(),
309 };
310 let base_url = config.base_url.clone().or_else(|| {
312 config
313 .extra
314 .get("base_url")
315 .and_then(|v| v.as_str())
316 .map(|s| s.to_string())
317 });
318
319 let api_key = config
321 .api_key_env
322 .as_ref()
323 .and_then(|env_var| std::env::var(env_var).ok());
324
325 let mut provider = UnifiedLLMProvider::from_spec_config(
326 provider_type,
327 &config.model,
328 api_key,
329 base_url,
330 core_config,
331 )
332 .map_err(|e| AgentError::LLM(e.to_string()))?
333 .with_feature_overrides(feature_overrides_from_config(config));
334 if let Some(choice) = config.tool_choice.clone() {
335 provider = provider.with_tool_choice(choice);
336 }
337
338 registry.register(alias, Arc::new(provider));
339 }
340
341 let default_alias = spec.llm.get_default_alias();
342 let router_alias = spec.llm.get_router_alias();
343
344 registry.set_default(&default_alias);
345 if let Some(router) = router_alias {
346 registry.set_router(&router);
347 }
348
349 self.llm_registry = Some(registry);
350 self.llm_registry_observed = false;
351 } else if let Some(config) = spec.llm.as_config() {
352 let provider_type = ProviderType::from_str(&config.provider)
353 .map_err(|e| AgentError::Config(e.to_string()))?;
354
355 let core_config = ai_agents_core::LLMConfig {
356 temperature: Some(config.temperature),
357 max_tokens: Some(config.max_tokens),
358 top_p: config.top_p,
359 top_k: None,
360 frequency_penalty: None,
361 presence_penalty: None,
362 stop_sequences: None,
363 timeout_seconds: config.timeout_seconds,
364 reasoning: config.reasoning,
365 reasoning_effort: config.reasoning_effort.clone(),
366 reasoning_budget_tokens: config.reasoning_budget_tokens,
367 extra: config.extra.clone(),
368 };
369 let base_url = config.base_url.clone().or_else(|| {
371 config
372 .extra
373 .get("base_url")
374 .and_then(|v| v.as_str())
375 .map(|s| s.to_string())
376 });
377
378 let api_key = config
380 .api_key_env
381 .as_ref()
382 .and_then(|env_var| std::env::var(env_var).ok());
383
384 let mut provider = UnifiedLLMProvider::from_spec_config(
385 provider_type,
386 &config.model,
387 api_key,
388 base_url,
389 core_config,
390 )
391 .map_err(|e| AgentError::LLM(e.to_string()))?
392 .with_feature_overrides(feature_overrides_from_config(config));
393 if let Some(choice) = config.tool_choice.clone() {
394 provider = provider.with_tool_choice(choice);
395 }
396
397 self.llm = Some(Arc::new(provider));
398 self.llm_registry_observed = false;
399 }
400
401 Ok(self)
402 }
403
404 pub fn auto_configure_features(mut self) -> Result<Self> {
421 if let Some(ref spec) = self.spec {
422 self.recovery_manager = Some(RecoveryManager::try_new(spec.error_recovery.clone())?);
423 self.tool_security = Some(ToolSecurityEngine::try_new(spec.tool_security.clone())?);
424
425 if spec.has_process() {
426 let mut processor = ProcessProcessor::new(spec.process.clone());
427 if let Some(ref registry) = self.llm_registry {
428 processor = processor.with_llm_registry(Arc::new(registry.clone()));
429 }
430 self.process_processor = Some(processor);
431 }
432
433 if self.tools.is_none() {
435 self.tools = Some(create_builtin_registry());
436 }
437 }
438 Ok(self)
439 }
440
441 pub async fn auto_configure_mcp(mut self) -> Result<Self> {
448 if let Some(ref spec) = self.spec {
449 let mcp_configs: Vec<_> = spec
451 .tools
452 .as_ref()
453 .map(|tools| {
454 tools
455 .iter()
456 .filter_map(|entry| entry.to_mcp_config())
457 .collect()
458 })
459 .unwrap_or_default();
460
461 if !mcp_configs.is_empty() {
462 let registry = self.tools.get_or_insert_with(create_builtin_registry);
463
464 for config in mcp_configs {
465 let tool_name = config.name.clone();
466 let timeout_ms = config.startup_timeout_ms;
467 let views_config = config.views.clone();
468
469 let wrapper = MCPWrapperTool::new(config);
470
471 match tokio::time::timeout(
473 std::time::Duration::from_millis(timeout_ms),
474 wrapper.initialized(),
475 )
476 .await
477 {
478 Ok(Ok(initialized_tool)) => {
479 tracing::info!(
480 tool = %tool_name,
481 functions = initialized_tool.function_count(),
482 "MCP wrapper tool registered"
483 );
484
485 let parent = Arc::new(initialized_tool);
486
487 registry.register(parent.clone()).map_err(|e| {
489 AgentError::Config(format!(
490 "Failed to register MCP tool '{}': {}",
491 tool_name, e
492 ))
493 })?;
494
495 for (view_name, view_config) in &views_config {
497 let view_tool = MCPViewTool::new(
498 view_name.clone(),
499 parent.clone(),
500 view_config.functions.clone(),
501 view_config.description.clone(),
502 )
503 .map_err(|e| {
504 AgentError::Config(format!(
505 "Failed to create MCP view '{}': {}",
506 view_name, e
507 ))
508 })?;
509
510 tracing::info!(
511 view = %view_name,
512 parent = %tool_name,
513 functions = view_config.functions.len(),
514 "MCP view tool registered"
515 );
516
517 registry.register(Arc::new(view_tool)).map_err(|e| {
518 AgentError::Config(format!(
519 "Failed to register MCP view '{}': {}",
520 view_name, e
521 ))
522 })?;
523 }
524 }
525 Ok(Err(e)) => {
526 return Err(AgentError::Config(format!(
527 "MCP tool '{}' initialization failed: {}",
528 tool_name, e
529 )));
530 }
531 Err(_) => {
532 return Err(AgentError::Config(format!(
533 "MCP tool '{}' timed out after {}ms",
534 tool_name, timeout_ms
535 )));
536 }
537 }
538 }
539 }
540 }
541 Ok(self)
542 }
543
544 pub fn llm(mut self, llm: Arc<dyn LLMProvider>) -> Self {
546 if self.routing_frozen.is_some() {
547 self.routing_dirty = true;
548 }
549 self.llm = Some(llm);
550 self
551 }
552
553 pub fn llm_alias(mut self, alias: impl Into<String>, provider: Arc<dyn LLMProvider>) -> Self {
555 if self.routing_frozen.is_some() {
556 self.routing_dirty = true;
557 }
558 if self.llm_registry.is_none() {
559 self.llm_registry = Some(LLMRegistry::new());
560 }
561 if let Some(ref mut registry) = self.llm_registry {
562 registry.register(alias, provider);
563 }
564 self
565 }
566
567 pub fn llm_registry(mut self, registry: LLMRegistry) -> Self {
569 if self.routing_frozen.is_some() {
570 self.routing_dirty = true;
571 }
572 self.llm_registry = Some(registry);
573 self.llm_registry_observed = false;
574 self
575 }
576
577 pub(crate) fn authoritative_llm_registry(
579 mut self,
580 registry: LLMRegistry,
581 observed: bool,
582 ) -> Self {
583 self.llm_registry = Some(registry);
584 self.llm_registry_observed = observed;
585 self
586 }
587
588 pub fn memory(mut self, memory: Arc<dyn Memory>) -> Self {
589 self.memory = Some(memory);
590 self
591 }
592
593 pub fn tools(mut self, tools: ToolRegistry) -> Self {
597 self.tools = Some(tools);
598 self
599 }
600
601 pub fn tool(mut self, tool: Arc<dyn Tool>) -> Self {
606 let registry = self.tools.get_or_insert_with(ToolRegistry::new);
607 let _ = registry.register(tool);
608 self
609 }
610
611 pub fn extend_tools(mut self, additional: ToolRegistry) -> Self {
616 let registry = self.tools.get_or_insert_with(ToolRegistry::new);
617 for id in additional.list_ids() {
618 if registry.get(&id).is_none()
619 && let Some(tool) = additional.get(&id)
620 {
621 let _ = registry.register(tool);
622 }
623 }
624 self
625 }
626
627 pub fn skill(mut self, skill: SkillDefinition) -> Self {
629 if self.routing_frozen.is_some() {
630 self.routing_dirty = true;
631 }
632 self.skills.push(skill);
633 self
634 }
635
636 pub fn skills(mut self, skills: Vec<SkillDefinition>) -> Self {
638 if self.routing_frozen.is_some() {
639 self.routing_dirty = true;
640 }
641 self.skills.extend(skills);
642 self
643 }
644
645 pub fn skill_loader(mut self, loader: SkillLoader) -> Self {
647 if self.routing_frozen.is_some() {
648 self.routing_dirty = true;
649 }
650 self.skill_loader = Some(loader);
651 self
652 }
653
654 pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
655 self.system_prompt = Some(prompt.into());
656 self
657 }
658
659 pub fn tools_prompt(mut self, prompt: impl Into<String>) -> Self {
660 self.tools_prompt = Some(prompt.into());
661 self.auto_tools_prompt = false;
662 self
663 }
664
665 pub fn auto_tools_prompt(mut self, auto: bool) -> Self {
666 self.auto_tools_prompt = auto;
667 self
668 }
669
670 pub fn max_iterations(mut self, max: u32) -> Self {
671 self.max_iterations = Some(max);
672 self
673 }
674
675 pub fn max_context_tokens(mut self, tokens: u32) -> Self {
676 self.max_context_tokens = Some(tokens);
677 self
678 }
679
680 pub fn recovery_manager(mut self, manager: RecoveryManager) -> Self {
682 if self.routing_frozen.is_some() {
683 self.routing_dirty = true;
684 }
685 self.recovery_manager = Some(manager);
686 self
687 }
688
689 pub fn tool_security(mut self, engine: ToolSecurityEngine) -> Self {
690 self.tool_security = Some(engine);
691 self
692 }
693
694 pub fn process_processor(mut self, processor: ProcessProcessor) -> Self {
696 if self.routing_frozen.is_some() {
697 self.routing_dirty = true;
698 }
699 self.process_processor = Some(processor);
700 self
701 }
702
703 pub fn message_filter(
704 mut self,
705 name: impl Into<String>,
706 filter: Arc<dyn MessageFilter>,
707 ) -> Self {
708 self.message_filters.insert(name.into(), filter);
709 self
710 }
711
712 pub fn context_manager(mut self, manager: Arc<ContextManager>) -> Self {
713 self.context_manager = Some(manager);
714 self
715 }
716
717 pub fn state_machine(mut self, machine: Arc<StateMachine>) -> Self {
719 if self.routing_frozen.is_some() {
720 self.routing_dirty = true;
721 }
722 self.state_machine = Some(machine);
723 self
724 }
725
726 pub fn transition_evaluator(mut self, evaluator: Arc<dyn TransitionEvaluator>) -> Self {
728 if self.routing_frozen.is_some() {
729 self.routing_dirty = true;
730 }
731 self.transition_evaluator = Some(evaluator);
732 self
733 }
734
735 pub fn hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
736 self.hooks = Some(hooks);
737 self
738 }
739
740 pub fn approval_handler(mut self, handler: Arc<dyn ApprovalHandler>) -> Self {
741 self.approval_handler = Some(handler);
742 self
743 }
744
745 pub fn hitl_engine(mut self, engine: HITLEngine) -> Self {
747 if self.routing_frozen.is_some() {
748 self.routing_dirty = true;
749 }
750 self.hitl_engine = Some(engine);
751 self
752 }
753
754 pub fn storage_config(mut self, config: StorageConfig) -> Self {
755 self.storage_config = Some(config);
756 self
757 }
758
759 pub fn storage(mut self, storage: Arc<dyn AgentStorage>) -> Self {
760 self.storage = Some(storage);
761 self
762 }
763
764 pub fn reasoning(mut self, config: ReasoningConfig) -> Self {
766 if self.routing_frozen.is_some() {
767 self.routing_dirty = true;
768 }
769 self.reasoning = Some(config);
770 self
771 }
772
773 pub fn reflection(mut self, config: ReflectionConfig) -> Self {
775 if self.routing_frozen.is_some() {
776 self.routing_dirty = true;
777 }
778 self.reflection = Some(config);
779 self
780 }
781
782 pub fn persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
784 self.persona_manager = Some(manager);
785 self
786 }
787
788 pub fn persona_templates(
790 mut self,
791 registry: Arc<ai_agents_persona::PersonaTemplateRegistry>,
792 ) -> Self {
793 self.persona_templates = Some(registry);
794 self
795 }
796
797 pub fn observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
799 if self.routing_frozen.is_some() {
800 self.routing_dirty = true;
801 }
802 self.observability_manager = Some(manager);
803 self
804 }
805
806 fn ensure_observability_manager(&mut self) -> Result<Option<Arc<ObservabilityManager>>> {
808 if let Some(manager) = self.observability_manager.as_ref() {
809 return Ok(Some(Arc::clone(manager)));
810 }
811 let Some(ref spec) = self.spec else {
812 return Ok(None);
813 };
814 if !spec.observability.enabled {
815 return Ok(None);
816 }
817 let config = self.observability_config_with_pricing(&spec.observability)?;
818 config
819 .validate()
820 .map_err(|e| AgentError::Config(e.to_string()))?;
821 let manager = ObservabilityManager::new(config);
822 self.observability_manager = Some(Arc::clone(&manager));
823 Ok(Some(manager))
824 }
825
826 fn observability_config_with_pricing(
828 &self,
829 config: &ObservabilityConfig,
830 ) -> Result<ObservabilityConfig> {
831 config
832 .clone()
833 .with_pricing_file_loaded(self.yaml_dir.as_deref())
834 .map_err(|e| AgentError::Config(e.to_string()))
835 }
836
837 pub(crate) fn prepare_routing(&mut self, freeze: bool) -> Result<()> {
839 if self.routing_frozen.as_ref().is_some_and(|frozen| {
840 self.routing_dirty
841 || !self
842 .llm_registry
843 .as_ref()
844 .is_some_and(|registry| registry.same_bindings(frozen))
845 }) {
846 return Err(AgentError::Config(
847 "Hierarchical LLM bindings changed after consumer construction".into(),
848 ));
849 }
850 let hierarchy_requested = self.spec.as_ref().map_or_else(
851 || {
852 self.llm_registry
853 .as_ref()
854 .is_some_and(|registry| registry.router_roles().is_some())
855 },
856 |spec| spec.llm.router_roles().is_some(),
857 );
858 if !hierarchy_requested {
859 if let (Some(spec), Some(registry)) = (&self.spec, &mut self.llm_registry)
860 && registry.router_roles().is_some()
861 {
862 spec.install_routing(registry);
863 }
864 return Ok(());
865 }
866 if let Some(spec) = &self.spec {
867 spec.validate()?;
868 }
869 if let Some(llm) = &self.llm {
870 let registry = self.llm_registry.get_or_insert_with(LLMRegistry::new);
871 if !registry.has("default") {
872 registry.register("default", llm.clone());
873 }
874 }
875 if let Some(spec) = &self.spec {
876 let registry = self.llm_registry.get_or_insert_with(LLMRegistry::new);
877 spec.install_routing(registry);
878 }
879 let hierarchy = self
880 .llm_registry
881 .as_ref()
882 .is_some_and(|registry| registry.router_roles().is_some());
883 if !hierarchy {
884 return Ok(());
885 }
886 if !self.skills_prepared {
887 if let Some(spec) = &self.spec {
888 let mut loader = self.skill_loader.take().unwrap_or_default();
889 if let Some(dir) = &self.yaml_dir {
890 loader.set_base_dir(dir);
891 }
892 self.skills.extend(loader.load_refs(&spec.skills)?);
893 }
894 self.skills_prepared = true;
895 }
896 let registry = self
897 .llm_registry
898 .as_ref()
899 .expect("hierarchy registry exists");
900 registry
901 .validate_router_roles()
902 .map_err(|e| AgentError::Config(e.to_string()))?;
903 let mut prepared = self.spec.clone().unwrap_or_default();
904 if let Some(recovery) = &self.recovery_manager {
905 prepared.error_recovery = recovery.config().clone();
906 }
907 prepared.llm = crate::spec::LLMConfigOrSelector::Selector(
908 crate::spec::LLMSelector::new(registry.default_alias())
909 .with_router_roles(registry.router_roles().expect("hierarchy exists").clone()),
910 );
911 if let Some(processor) = &self.process_processor {
912 prepared.process = processor.config().clone();
913 }
914 if let Some(engine) = &self.hitl_engine {
915 prepared.hitl = Some(engine.config().clone());
916 }
917 prepared.skills = self
918 .skills
919 .iter()
920 .cloned()
921 .map(ai_agents_skills::SkillRef::Inline)
922 .collect();
923 if let Some(state) = &self.state_machine {
924 prepared.states = Some(state.config().clone());
925 }
926 if let Some(reasoning) = &self.reasoning {
927 prepared.reasoning = reasoning.clone();
928 }
929 if let Some(reflection) = &self.reflection {
930 prepared.reflection = reflection.clone();
931 }
932 for alias in prepared.referenced_llm_aliases() {
933 registry.get(&alias).map_err(|_| {
934 AgentError::Config(format!(
935 "Configured local LLM alias '{alias}' is not registered"
936 ))
937 })?;
938 }
939 if self.transition_evaluator.is_some() {
940 if registry
941 .router_roles()
942 .and_then(|roles| roles.state.as_ref())
943 .and_then(|state| state.transition.as_ref())
944 .is_some()
945 {
946 return Err(AgentError::Config(
947 "Custom evaluator conflicts with llm.router.state.transition".into(),
948 ));
949 }
950 if prepared.has_semantic_parallel_transition() {
951 return Err(AgentError::Config(
952 "Custom evaluator cannot implement hierarchical semantic parallel transitions"
953 .into(),
954 ));
955 }
956 }
957 if freeze {
958 self.routing_frozen = Some(registry.clone());
959 }
960 Ok(())
961 }
962
963 pub(crate) fn prepared_registry_observed(&self) -> bool {
966 self.llm_registry_observed
967 }
968
969 pub(crate) fn prepared_registry(&self) -> Option<&LLMRegistry> {
970 self.llm_registry.as_ref()
971 }
972
973 pub(crate) fn wrap_llm_registry_for_observability(&mut self) -> Result<()> {
975 if self.llm_registry_observed {
976 return Ok(());
977 }
978 let Some(manager) = self.ensure_observability_manager()? else {
979 return Ok(());
980 };
981 let Some(registry) = self.llm_registry.take() else {
982 return Ok(());
983 };
984 let model_by_alias = model_by_alias_from_spec(self.spec.as_ref());
985 let wrapped = wrap_registry_with_observability(registry, manager, &model_by_alias);
986 let wrapped_arc = Arc::new(wrapped.clone());
987 if let Some(processor) = self.process_processor.take() {
988 self.process_processor = Some(processor.with_llm_registry(wrapped_arc));
989 }
990 if self.routing_frozen.is_some() {
991 self.routing_frozen = Some(wrapped.clone());
992 }
993 self.llm_registry = Some(wrapped);
994 self.llm_registry_observed = true;
995 Ok(())
996 }
997
998 pub fn streaming(mut self, enabled: bool) -> Self {
999 let mut config = self.streaming.unwrap_or_default();
1000 config.enabled = enabled;
1001 self.streaming = Some(config);
1002 self
1003 }
1004
1005 pub async fn auto_configure_spawner(mut self) -> Result<Self> {
1008 self.prepare_routing(false)?;
1009 let spawner_config = match self.spec.as_ref().and_then(|s| s.spawner.as_ref()) {
1010 Some(c) => c.clone(),
1011 None => return Ok(self),
1012 };
1013
1014 if spawner_config.shared_llms && self.llm_registry.is_none() {
1015 let provider = self.llm.as_ref().cloned().ok_or_else(|| {
1016 AgentError::Config(
1017 "spawner.shared_llms requires the parent LLM provider to be configured"
1018 .to_string(),
1019 )
1020 })?;
1021 let mut registry = LLMRegistry::new();
1022 registry.register("default", provider);
1023 registry.set_default("default");
1024 self.llm_registry = Some(registry);
1025 self.llm_registry_observed = false;
1026 }
1027
1028 self.wrap_llm_registry_for_observability()?;
1029 self.prepare_routing(true)?;
1030 let observability_manager = self.observability_manager.clone();
1031
1032 use crate::spawner::{
1033 AgentRegistry, AgentSpawner,
1034 config::{configure_spawner_tools, resolve_templates},
1035 };
1036
1037 let mut spawner = AgentSpawner::new();
1038
1039 if let Some(ref manager) = observability_manager {
1040 spawner = spawner.with_observability(Arc::clone(manager));
1041 }
1042 spawner = spawner.with_resource_locks(self.shared_resource_locks());
1043
1044 if spawner_config.shared_llms {
1045 let reg = self.llm_registry.as_ref().ok_or_else(|| {
1046 AgentError::Config(
1047 "spawner.shared_llms requires the parent LLM registry to be configured"
1048 .to_string(),
1049 )
1050 })?;
1051 spawner = if self.llm_registry_observed {
1052 spawner.with_shared_observed_llms(reg.clone())
1053 } else {
1054 spawner.with_shared_llms(reg.clone())
1055 };
1056 }
1057
1058 if !spawner_config.shared_context.is_empty() {
1059 spawner = spawner.with_shared_context_map(spawner_config.shared_context.clone());
1060 }
1061
1062 if let Some(max) = spawner_config.max_agents {
1063 spawner = spawner.with_max_agents(max);
1064 }
1065
1066 if let Some(ref prefix) = spawner_config.name_prefix {
1067 spawner = spawner.with_name_prefix(prefix.clone())?;
1068 }
1069
1070 if !spawner_config.templates.is_empty() {
1072 let resolved = resolve_templates(&spawner_config.templates, self.yaml_dir.as_deref())?;
1073 spawner = spawner.with_templates(resolved);
1074 }
1075
1076 if let Some(ref allowed) = spawner_config.allowed_tools {
1077 spawner = spawner.with_allowed_tools(allowed.clone());
1078 }
1079
1080 if let Some(ref sc) = spawner_config.shared_storage {
1082 let converted = crate::spec::storage::to_storage_config(sc);
1083 if let Some(st) = ai_agents_storage::create_storage(&converted).await? {
1084 spawner = spawner.with_shared_storage(Arc::clone(&st));
1085
1086 let parent_has_storage = self.storage.is_some()
1088 || self.storage_config.is_some()
1089 || self.spec.as_ref().is_some_and(|s| s.has_storage());
1090 if !parent_has_storage {
1091 self.storage = Some(st);
1092 }
1093 }
1094 }
1095
1096 let spawner = Arc::new(spawner);
1097 let registry = Arc::new(AgentRegistry::new());
1098
1099 self.spawner = Some(Arc::clone(&spawner));
1100 self.spawner_registry = Some(Arc::clone(®istry));
1101 let llm_for_tools = Arc::new(self.llm_registry.clone().unwrap_or_default());
1102 let agent_name = self
1103 .spec
1104 .as_ref()
1105 .map(|s| s.name.clone())
1106 .unwrap_or_default();
1107
1108 let tools = configure_spawner_tools(
1109 Arc::clone(&spawner),
1110 Arc::clone(®istry),
1111 Arc::clone(&llm_for_tools),
1112 &agent_name,
1113 );
1114
1115 let tool_registry = self.tools.get_or_insert_with(create_builtin_registry);
1116 for tool in tools {
1117 let _ = tool_registry.register(tool);
1118 }
1119
1120 tracing::info!("Spawner tools registered");
1121
1122 if spawner_config.orchestration_tools.is_enabled() {
1124 let orch_tools = crate::orchestration::tools::configure_orchestration_tools(
1125 &spawner_config.orchestration_tools,
1126 Arc::clone(®istry),
1127 Arc::clone(&llm_for_tools),
1128 );
1129 let tool_registry = self.tools.get_or_insert_with(create_builtin_registry);
1130 for tool in orch_tools {
1131 let _ = tool_registry.register(tool);
1132 }
1133 tracing::info!("Orchestration tools registered");
1134 }
1135
1136 for entry in &spawner_config.auto_spawn {
1137 let yaml_path = if let Some(ref dir) = self.yaml_dir {
1138 dir.join(&entry.agent)
1139 } else {
1140 std::path::PathBuf::from(&entry.agent)
1141 };
1142
1143 tracing::info!(id = %entry.id, path = %yaml_path.display(), "Auto-spawning agent");
1144 let spawned = spawner
1145 .spawn_from_yaml_file_with_id(entry.id.clone(), &yaml_path)
1146 .await
1147 .map_err(|error| {
1148 AgentError::Config(format!(
1149 "Failed to auto-spawn agent '{}' from '{}': {}",
1150 entry.id,
1151 yaml_path.display(),
1152 error
1153 ))
1154 })?;
1155 registry.register(spawned).await.map_err(|error| {
1156 AgentError::Config(format!(
1157 "Failed to register auto-spawned agent '{}': {}",
1158 entry.id, error
1159 ))
1160 })?;
1161 tracing::info!(id = %entry.id, "Auto-spawned agent registered");
1162 }
1163
1164 if let Some(ref spec) = self.spec
1166 && let Some(ref state_config) = spec.states
1167 {
1168 let refs = collect_orchestration_refs(&state_config.states);
1169 let mut missing: Vec<String> = Vec::new();
1170
1171 for (agent_id, state_name, pattern) in &refs {
1172 if !registry.contains(agent_id) {
1173 missing.push(format!(
1174 " - '{}' (referenced by state '{}' via {})",
1175 agent_id, state_name, pattern
1176 ));
1177 }
1178 }
1179
1180 if !missing.is_empty() {
1181 missing.sort();
1182 missing.dedup();
1183 return Err(AgentError::Config(format!(
1184 "Auto-spawn validation failed. These agents are referenced by \
1185 orchestration states but were not successfully spawned:\n\n{}\n\n\
1186 Check that agent YAML files exist and contain valid specs.",
1187 missing.join("\n")
1188 )));
1189 }
1190 }
1191
1192 Ok(self)
1193 }
1194
1195 pub fn build(mut self) -> Result<RuntimeAgent> {
1197 self.prepare_routing(false)?;
1198 let resource_locks = self.shared_resource_locks();
1199 let actor_memory_config = self
1201 .spec
1202 .as_ref()
1203 .and_then(|s| s.memory.actor_memory.clone());
1204 let facts_config = self.spec.as_ref().and_then(|s| s.memory.facts.clone());
1205 let relationships_config = self
1206 .spec
1207 .as_ref()
1208 .and_then(|s| s.memory.relationships.clone());
1209
1210 let observability_manager = self.ensure_observability_manager()?;
1211
1212 let base_prompt = self
1213 .system_prompt
1214 .ok_or_else(|| AgentError::Config("System prompt is required".into()))?;
1215
1216 let mut tools = self.tools.unwrap_or_default();
1217
1218 if self
1219 .llm_registry
1220 .as_ref()
1221 .is_some_and(|registry| registry.router_roles().is_some())
1222 {
1223 tools = tools.map_tools_agent_local_web(|tool| tool);
1224 }
1225 let system_prompt = base_prompt;
1228
1229 let max_iterations = self.max_iterations.unwrap_or(10);
1230
1231 let info = if let Some(ref spec) = self.spec {
1232 AgentInfo::new(&spec.name, &spec.name, &spec.version)
1233 .with_description(spec.description.clone().unwrap_or_default())
1234 } else {
1235 AgentInfo::new("agent", "Agent", "1.0.0")
1236 };
1237
1238 if !self.skills_prepared
1239 && let Some(ref spec) = self.spec
1240 && !spec.skills.is_empty()
1241 {
1242 let mut loader = self.skill_loader.take().unwrap_or_default();
1243 if let Some(ref dir) = self.yaml_dir {
1244 loader.set_base_dir(dir);
1245 }
1246 let loaded_skills = loader.load_refs(&spec.skills)?;
1247 self.skills.extend(loaded_skills);
1248 }
1249
1250 let mut llm_registry = self.llm_registry.unwrap_or_default();
1251
1252 if let Some(llm) = self.llm
1253 && !llm_registry.has("default")
1254 {
1255 let provider = if self.llm_registry_observed {
1256 if let Some(ref manager) = observability_manager {
1257 Arc::new(ObservedLLMProvider::new(
1258 llm.clone(),
1259 Arc::clone(manager),
1260 Some("default".to_string()),
1261 llm.provider_name().to_string(),
1262 model_by_alias_from_spec(self.spec.as_ref())
1263 .get("default")
1264 .cloned()
1265 .unwrap_or_else(|| "default".to_string()),
1266 )) as Arc<dyn LLMProvider>
1267 } else {
1268 llm.clone()
1269 }
1270 } else {
1271 llm.clone()
1272 };
1273 llm_registry.register("default", provider);
1274 }
1275
1276 if let Some(ref spec) = self.spec {
1277 spec.install_routing(&mut llm_registry);
1278 }
1279
1280 if llm_registry.is_empty() {
1281 return Err(AgentError::Config(
1282 "At least one LLM provider is required".into(),
1283 ));
1284 }
1285
1286 if let Some(ref manager) = observability_manager
1287 && !self.llm_registry_observed
1288 {
1289 let model_by_alias = model_by_alias_from_spec(self.spec.as_ref());
1290 llm_registry = wrap_registry_with_observability(
1291 llm_registry,
1292 Arc::clone(manager),
1293 &model_by_alias,
1294 );
1295 self.llm_registry_observed = true;
1296 }
1297
1298 let memory_local = self
1299 .spec
1300 .as_ref()
1301 .and_then(|spec| spec.memory.summarizer_llm.as_deref());
1302 let (role_summary, role_merge) = if self.memory.is_none()
1303 && self
1304 .spec
1305 .as_ref()
1306 .is_some_and(|spec| spec.memory.is_compacting())
1307 {
1308 (
1309 llm_registry
1310 .resolve_role_override(LLMRole::MemorySummarize, memory_local)
1311 .map_err(|e| AgentError::Config(e.to_string()))?
1312 .map(|resolved| resolved.provider),
1313 llm_registry
1314 .resolve_role_override(LLMRole::MemoryMerge, memory_local)
1315 .map_err(|e| AgentError::Config(e.to_string()))?
1316 .map(|resolved| resolved.provider),
1317 )
1318 } else {
1319 (None, None)
1320 };
1321 let memory = self.memory.unwrap_or_else(|| {
1323 if let Some(ref spec) = self.spec {
1324 if spec.memory.is_compacting() {
1325 let summarizer_llm = role_summary.or_else(|| {
1326 spec.memory
1327 .summarizer_llm
1328 .as_ref()
1329 .and_then(|alias| llm_registry.get(alias).ok())
1330 .or_else(|| llm_registry.router().ok())
1331 .or_else(|| llm_registry.default().ok())
1332 });
1333
1334 let summarizer: Arc<dyn Summarizer> = match summarizer_llm {
1335 Some(llm) => {
1336 let mut summarizer = LLMSummarizer::new(llm);
1337 if let Some(merge) = role_merge {
1338 summarizer = summarizer.with_merge_llm(merge);
1339 }
1340 Arc::new(summarizer)
1341 }
1342 None => Arc::new(NoopSummarizer),
1343 };
1344 let config = spec.memory.to_compacting_config();
1345 return Arc::new(CompactingMemory::new(summarizer, config));
1346 }
1347 Arc::new(InMemoryStore::new(spec.memory.max_messages))
1348 } else {
1349 Arc::new(InMemoryStore::new(100))
1350 }
1351 });
1352
1353 let persona_manager: Option<Arc<ai_agents_persona::PersonaManager>> =
1355 if let Some(pm) = self.persona_manager.take() {
1356 Some(pm)
1357 } else if let Some(ref spec) = self.spec {
1358 if spec.has_persona() {
1359 let persona_config = spec.persona.clone().unwrap();
1360 let renderer = ai_agents_context::TemplateRenderer::new();
1361 let registry = self.persona_templates.clone();
1362 let manager = ai_agents_persona::PersonaManager::from_config(
1363 persona_config,
1364 registry,
1365 renderer,
1366 )
1367 .map_err(|e| {
1368 AgentError::Config(format!("Failed to create PersonaManager: {}", e))
1369 })?;
1370 Some(Arc::new(manager))
1371 } else {
1372 None
1373 }
1374 } else {
1375 None
1376 };
1377
1378 if let Some(ref pm) = persona_manager
1380 && pm.should_register_evolve_tool()
1381 {
1382 let evolve_tool = ai_agents_persona::PersonaEvolveTool::new(pm.clone());
1383 let _ = tools.register(Arc::new(evolve_tool));
1384 }
1385
1386 if let Some(ref manager) = observability_manager {
1387 tools = tools.map_tools(|tool| {
1388 Arc::new(ObservedTool::new(tool, Arc::clone(manager))) as Arc<dyn Tool>
1389 });
1390 }
1391
1392 let relationship_manager: Option<Arc<RelationshipManager>> =
1393 if let Some(ref config) = relationships_config {
1394 if config.enabled {
1395 let evaluator: Option<Arc<dyn RelationshipEvaluatorTrait>> =
1396 if config.auto_update.enabled {
1397 let llm = match llm_registry
1398 .resolve_role_override(
1399 LLMRole::MemoryRelationships,
1400 config.auto_update.llm.as_deref(),
1401 )
1402 .map_err(|e| AgentError::Config(e.to_string()))?
1403 {
1404 Some(resolved) => Some(resolved.provider),
1405 None => config
1406 .auto_update
1407 .llm
1408 .as_ref()
1409 .and_then(|alias| llm_registry.get(alias).ok())
1410 .or_else(|| llm_registry.router().ok())
1411 .or_else(|| llm_registry.default().ok()),
1412 };
1413 llm.map(|llm| {
1414 Arc::new(RelationshipEvaluator::new(llm))
1415 as Arc<dyn RelationshipEvaluatorTrait>
1416 })
1417 } else {
1418 None
1419 };
1420 Some(Arc::new(RelationshipManager::from_config_with_evaluator(
1421 config.clone(),
1422 evaluator,
1423 )?))
1424 } else {
1425 None
1426 }
1427 } else {
1428 None
1429 };
1430
1431 let tools_arc = Arc::new(tools);
1432 let llm_registry_arc = Arc::new(llm_registry);
1433 let web_provider = match llm_registry_arc
1434 .resolve_role_override(LLMRole::WebExtract, None)
1435 .map_err(|e| AgentError::Config(e.to_string()))?
1436 {
1437 Some(resolved) => Some(resolved.provider),
1438 None => llm_registry_arc
1439 .router()
1440 .ok()
1441 .or_else(|| llm_registry_arc.default().ok()),
1442 };
1443 tools_arc.set_web_fetch_extractor(web_provider);
1444
1445 let declared_tool_ids: Option<Vec<String>> = Some(if self.spec.is_none() {
1448 tools_arc.list_ids()
1449 } else {
1450 let mut ids: Vec<String> = self
1451 .spec
1452 .as_ref()
1453 .and_then(|s| s.tools.as_ref())
1454 .map(|tools| {
1455 let mut ids: Vec<String> = tools
1456 .iter()
1457 .filter_map(|t| {
1458 tools_arc
1459 .canonical_id(t.name())
1460 .or_else(|| Some(t.name().to_string()))
1461 })
1462 .collect();
1463
1464 for entry in tools {
1465 if let Some(mcp_config) = entry.to_mcp_config() {
1466 for view_name in mcp_config.views.keys() {
1467 ids.push(
1468 tools_arc
1469 .canonical_id(view_name)
1470 .unwrap_or_else(|| view_name.clone()),
1471 );
1472 }
1473 }
1474 }
1475
1476 ids
1477 })
1478 .unwrap_or_default();
1479
1480 if let Some(ref spec) = self.spec
1481 && let Some(ref spawner) = spec.spawner
1482 {
1483 ids.extend(spawner.management_tools.granted_management_tool_ids());
1485 ids.extend(spawner.orchestration_tools.granted_orchestration_tool_ids());
1486 }
1487
1488 if persona_manager
1489 .as_ref()
1490 .is_some_and(|pm| pm.should_register_evolve_tool())
1491 {
1492 ids.push("persona_evolve".to_string());
1494 }
1495
1496 ids.sort();
1497 ids.dedup();
1498 ids
1499 });
1500
1501 if let Some(ref ids) = declared_tool_ids {
1506 let mcp_names: Vec<String> = self
1507 .spec
1508 .as_ref()
1509 .and_then(|s| s.tools.as_ref())
1510 .map(|tools| {
1511 let mut names = Vec::new();
1512 for entry in tools {
1513 if entry.is_mcp() {
1514 names.push(entry.name().to_string());
1515 if let Some(cfg) = entry.to_mcp_config() {
1516 names.extend(cfg.views.keys().cloned());
1517 }
1518 }
1519 }
1520 names
1521 })
1522 .unwrap_or_default();
1523
1524 let missing: Vec<&str> = ids
1525 .iter()
1526 .filter(|id| !mcp_names.contains(id))
1527 .filter(|id| tools_arc.get(id).is_none())
1528 .map(|s| s.as_str())
1529 .collect();
1530
1531 if !missing.is_empty() {
1532 return Err(AgentError::Config(format!(
1533 "Tools granted by YAML but not registered: [{}]. \
1534 Register them via .tool(Arc::new(...)) or the matching auto_configure_* method before .build(), \
1535 or remove the grant from YAML.",
1536 missing.join(", ")
1537 )));
1538 }
1539 }
1540
1541 let mut agent = RuntimeAgent::try_new(
1542 info,
1543 llm_registry_arc.clone(),
1544 memory,
1545 tools_arc,
1546 self.skills,
1547 system_prompt,
1548 max_iterations,
1549 )?
1550 .with_shared_resource_locks(resource_locks)
1551 .with_declared_tool_ids(declared_tool_ids);
1552
1553 if let Some(tokens) = self.max_context_tokens {
1554 agent = agent.with_max_context_tokens(tokens);
1555 }
1556
1557 if let Some(manager) = self.recovery_manager {
1558 agent = agent.with_recovery_manager(manager);
1559 } else if let Some(ref spec) = self.spec {
1560 agent =
1561 agent.with_recovery_manager(RecoveryManager::try_new(spec.error_recovery.clone())?);
1562 }
1563
1564 if let Some(engine) = self.tool_security {
1565 agent = agent.with_tool_security(engine);
1566 } else if let Some(ref spec) = self.spec {
1567 agent =
1568 agent.with_tool_security(ToolSecurityEngine::try_new(spec.tool_security.clone())?);
1569 }
1570
1571 if let Some(processor) = self.process_processor {
1572 agent =
1573 agent.with_process_processor(processor.with_llm_registry(llm_registry_arc.clone()));
1574 } else if let Some(ref spec) = self.spec
1575 && spec.has_process()
1576 {
1577 let processor = ProcessProcessor::new(spec.process.clone())
1578 .with_llm_registry(llm_registry_arc.clone());
1579 agent = agent.with_process_processor(processor);
1580 }
1581
1582 for (name, filter) in self.message_filters {
1583 agent.register_message_filter(name, filter);
1584 }
1585
1586 let role_transition = if self.transition_evaluator.is_none()
1587 && (self.state_machine.is_some()
1588 || self.spec.as_ref().is_some_and(|spec| spec.states.is_some()))
1589 {
1590 llm_registry_arc
1591 .resolve_role_override(LLMRole::StateTransition, None)
1592 .map_err(|e| AgentError::Config(e.to_string()))?
1593 .map(|resolved| resolved.provider)
1594 } else {
1595 None
1596 };
1597 if let Some(state_machine) = self.state_machine {
1599 let evaluator = self.transition_evaluator.unwrap_or_else(|| {
1600 let eval_llm = role_transition.clone().unwrap_or_else(|| {
1601 llm_registry_arc
1602 .get("evaluator")
1603 .or_else(|_| llm_registry_arc.router())
1604 .or_else(|_| llm_registry_arc.default())
1605 .expect("At least one LLM required for transition evaluator")
1606 });
1607 Arc::new(LLMTransitionEvaluator::new(eval_llm))
1608 });
1609 agent = agent.with_state_machine(state_machine, evaluator);
1610 } else if let Some(ref spec) = self.spec
1611 && let Some(ref state_config) = spec.states
1612 {
1613 let state_machine = StateMachine::new(state_config.clone())?;
1614 let evaluator = self.transition_evaluator.unwrap_or_else(|| {
1615 let eval_llm = role_transition.clone().unwrap_or_else(|| {
1616 llm_registry_arc
1617 .get("evaluator")
1618 .or_else(|_| llm_registry_arc.router())
1619 .or_else(|_| llm_registry_arc.default())
1620 .expect("At least one LLM required for transition evaluator")
1621 });
1622 Arc::new(LLMTransitionEvaluator::new(eval_llm))
1623 });
1624 agent = agent.with_state_machine(Arc::new(state_machine), evaluator);
1625 }
1626
1627 if let Some(context_manager) = self.context_manager {
1629 agent = agent.with_context_manager(context_manager);
1630 } else if let Some(ref spec) = self.spec
1631 && !spec.context.is_empty()
1632 {
1633 let context_manager = ContextManager::new(
1634 spec.context.clone(),
1635 spec.name.clone(),
1636 spec.version.clone(),
1637 );
1638 agent = agent.with_context_manager(Arc::new(context_manager));
1639 }
1640
1641 if let Some(ref spec) = self.spec {
1643 agent = agent.with_parallel_tools(spec.parallel_tools.clone());
1644 let streaming_config = self
1645 .streaming
1646 .clone()
1647 .unwrap_or_else(|| spec.streaming.clone());
1648 agent = agent.with_streaming(streaming_config);
1649
1650 agent = agent.with_runtime_config(spec.runtime.clone());
1651
1652 if let Some(ref budget) = spec.memory.token_budget {
1654 agent = agent.with_memory_token_budget(budget.clone());
1655 }
1656
1657 if self.storage_config.is_none() && spec.has_storage() {
1659 agent = agent.with_storage_config(spec.storage.clone());
1660 }
1661 }
1662
1663 if let Some(storage_config) = self.storage_config {
1665 agent = agent.with_storage_config(storage_config);
1666 }
1667 if let Some(storage) = self.storage {
1668 agent = agent.with_storage(storage);
1669 }
1670
1671 if let (Some(spawner), Some(registry)) = (self.spawner, self.spawner_registry) {
1673 agent = agent.with_spawner_handles(spawner, registry);
1674 }
1675
1676 if let Some(manager) = observability_manager {
1678 agent = agent.with_observability(Arc::clone(&manager));
1679 let observability_hooks: Arc<dyn AgentHooks> =
1680 Arc::new(ObservabilityHooks::new(manager));
1681 let hooks: Arc<dyn AgentHooks> = if let Some(user_hooks) = self.hooks {
1682 Arc::new(
1683 CompositeHooks::new()
1684 .add(user_hooks)
1685 .add(observability_hooks),
1686 )
1687 } else {
1688 observability_hooks
1689 };
1690 agent = agent.with_hooks(hooks);
1691 } else if let Some(hooks) = self.hooks {
1692 agent = agent.with_hooks(hooks);
1693 }
1694
1695 if let Some(hitl_engine) = self.hitl_engine {
1697 let handler = self
1698 .approval_handler
1699 .unwrap_or_else(|| Arc::new(RejectAllHandler::new()));
1700 agent = agent.with_hitl(hitl_engine, handler);
1701 } else if let Some(ref spec) = self.spec
1702 && let Some(ref hitl_config) = spec.hitl
1703 {
1704 let hitl_engine = HITLEngine::new(hitl_config.clone());
1705 let handler = self
1706 .approval_handler
1707 .unwrap_or_else(|| Arc::new(RejectAllHandler::new()));
1708 agent = agent.with_hitl(hitl_engine, handler);
1709 }
1710
1711 if let Some(reasoning) = self.reasoning {
1712 agent = agent.with_reasoning(reasoning);
1713 } else if let Some(ref spec) = self.spec {
1714 agent = agent.with_reasoning(spec.reasoning.clone());
1715 }
1716
1717 if let Some(reflection) = self.reflection {
1718 agent = agent.with_reflection(reflection);
1719 } else if let Some(ref spec) = self.spec {
1720 agent = agent.with_reflection(spec.reflection.clone());
1721 }
1722
1723 if let Some(ref spec) = self.spec
1725 && spec.disambiguation.is_enabled()
1726 {
1727 agent = agent.with_disambiguation(spec.disambiguation.clone());
1728 }
1729
1730 if let Some(pm) = persona_manager {
1732 agent = agent.with_persona(pm);
1733 }
1734
1735 if let Some(relationship_manager) = relationship_manager {
1736 agent = agent.with_relationships(relationship_manager);
1737 }
1738
1739 if actor_memory_config.is_some() || facts_config.is_some() {
1745 agent = agent.with_facts_config(actor_memory_config, facts_config);
1746 }
1747
1748 Ok(agent)
1749 }
1750}
1751
1752fn collect_orchestration_refs(
1754 states: &std::collections::HashMap<String, ai_agents_state::StateDefinition>,
1755) -> Vec<(String, String, &'static str)> {
1756 let mut refs = Vec::new();
1757
1758 for (state_name, def) in states {
1759 if let Some(ref delegate_id) = def.delegate {
1760 refs.push((delegate_id.clone(), state_name.clone(), "delegate"));
1761 }
1762 if let Some(ref concurrent) = def.concurrent {
1763 for agent_ref in &concurrent.agents {
1764 refs.push((agent_ref.id().to_string(), state_name.clone(), "concurrent"));
1765 }
1766 }
1767 if let Some(ref gc) = def.group_chat {
1768 for participant in &gc.participants {
1769 refs.push((participant.id.clone(), state_name.clone(), "group_chat"));
1770 }
1771 }
1772 if let Some(ref pipeline) = def.pipeline {
1773 for stage in &pipeline.stages {
1774 refs.push((stage.id().to_string(), state_name.clone(), "pipeline"));
1775 }
1776 }
1777 if let Some(ref handoff) = def.handoff {
1778 refs.push((handoff.initial_agent.clone(), state_name.clone(), "handoff"));
1779 for agent_id in &handoff.available_agents {
1780 refs.push((agent_id.clone(), state_name.clone(), "handoff"));
1781 }
1782 }
1783
1784 if let Some(ref sub_states) = def.states {
1786 refs.extend(collect_orchestration_refs(sub_states));
1787 }
1788 }
1789
1790 refs
1791}
1792
1793impl Default for AgentBuilder {
1794 fn default() -> Self {
1795 Self::new()
1796 }
1797}
1798
1799#[cfg(test)]
1800mod tests {
1801 use super::*;
1802
1803 #[test]
1804 fn test_builder_new() {
1805 let builder = AgentBuilder::new();
1806 assert!(builder.spec.is_none());
1807 assert!(builder.system_prompt.is_none());
1808 }
1809
1810 #[test]
1811 fn test_builder_from_yaml() {
1812 let yaml = r#"
1813name: TestAgent
1814system_prompt: "You are helpful."
1815llm:
1816 provider: openai
1817 model: gpt-4
1818"#;
1819 let builder = AgentBuilder::from_yaml(yaml).unwrap();
1820 assert!(builder.spec.is_some());
1821 assert_eq!(builder.spec.as_ref().unwrap().name, "TestAgent");
1822 }
1823
1824 #[test]
1825 fn test_builder_from_yaml_rejects_nested_unknown_path() {
1826 let yaml = r#"
1827name: TestAgent
1828system_prompt: test
1829runtime:
1830 optimization:
1831 max_parallel_runtime_task: 4
1832"#;
1833 let error = match AgentBuilder::from_yaml(yaml) {
1834 Ok(_) => panic!("expected strict parse failure"),
1835 Err(error) => error.to_string(),
1836 };
1837 assert!(
1838 error.contains("runtime.optimization.max_parallel_runtime_task"),
1839 "{error}"
1840 );
1841 }
1842
1843 #[test]
1844 fn test_feature_override_single_llm_builder_path() {
1845 let yaml = r#"
1846name: LocalAgent
1847system_prompt: "You are helpful."
1848llm:
1849 provider: ollama
1850 model: llama3.1
1851 function_calling: true
1852"#;
1853 let builder = AgentBuilder::from_yaml(yaml)
1854 .unwrap()
1855 .auto_configure_llms()
1856 .unwrap();
1857
1858 let llm = builder.llm.as_ref().unwrap();
1859 assert!(llm.supports(LLMFeature::FunctionCalling));
1860 }
1861
1862 #[test]
1863 fn test_feature_override_named_llms_builder_path() {
1864 let yaml = r#"
1865name: LocalAgent
1866system_prompt: "You are helpful."
1867llms:
1868 default:
1869 provider: openai-compatible
1870 model: qwen3:8b
1871 base_url: http://localhost:11434/v1
1872 json_mode: true
1873llm:
1874 default: default
1875"#;
1876 let builder = AgentBuilder::from_yaml(yaml)
1877 .unwrap()
1878 .auto_configure_llms()
1879 .unwrap();
1880
1881 let registry = builder.llm_registry.as_ref().unwrap();
1882 let llm = registry.get("default").unwrap();
1883 assert!(llm.supports(LLMFeature::JsonMode));
1884 }
1885
1886 #[test]
1887 fn test_builder_from_yaml_with_tool_security() {
1888 let yaml = r#"
1889name: SecureAgent
1890system_prompt: "You are helpful."
1891llm:
1892 provider: openai
1893 model: gpt-4
1894max_context_tokens: 8192
1895error_recovery:
1896 default:
1897 max_retries: 5
1898tool_security:
1899 enabled: true
1900 tools:
1901 http:
1902 rate_limit: 10
1903"#;
1904 let builder = AgentBuilder::from_yaml(yaml).unwrap();
1905 assert!(builder.spec.is_some());
1906 let spec = builder.spec.as_ref().unwrap();
1907 assert_eq!(spec.max_context_tokens, 8192);
1908 assert_eq!(spec.error_recovery.default.max_retries, 5);
1909 assert!(spec.tool_security.enabled);
1910 }
1911
1912 #[test]
1913 fn test_builder_from_yaml_rejects_zero_max_results() {
1914 let yaml = r#"
1915name: SecureAgent
1916system_prompt: "You are helpful."
1917tool_security:
1918 enabled: true
1919 tools:
1920 web_search:
1921 max_results: 0
1922"#;
1923 let error = AgentBuilder::from_yaml(yaml).err().unwrap();
1924 assert!(
1925 error
1926 .to_string()
1927 .contains("tool_security.tools.web_search.max_results must be greater than 0")
1928 );
1929 }
1930
1931 #[test]
1932 fn test_builder_from_yaml_with_skills() {
1933 let yaml = r#"
1934name: SkillAgent
1935system_prompt: "You are helpful."
1936llm:
1937 provider: openai
1938 model: gpt-4
1939skills:
1940 - id: greeting
1941 description: "Greet users"
1942 trigger: "When user says hello"
1943 steps:
1944 - prompt: "Hello!"
1945"#;
1946 let builder = AgentBuilder::from_yaml(yaml).unwrap();
1947 assert!(builder.spec.is_some());
1948 assert!(!builder.spec.as_ref().unwrap().skills.is_empty());
1949 }
1950
1951 #[test]
1952 fn test_builder_from_spec() {
1953 let spec = AgentSpec {
1954 name: "test".to_string(),
1955 version: "1.0".to_string(),
1956 description: Some("Test agent".to_string()),
1957 system_prompt: "You are helpful".to_string(),
1958 ..Default::default()
1959 };
1960
1961 let builder = AgentBuilder::from_spec(spec);
1962 assert!(builder.spec.is_some());
1963 assert_eq!(builder.system_prompt, Some("You are helpful".to_string()));
1964 }
1965
1966 #[test]
1967 fn test_builder_rejects_invalid_programmatic_tool_security() {
1968 use ai_agents_llm::mock::MockLLMProvider;
1969
1970 let mut spec = AgentSpec::default();
1971 spec.tool_security.tools.insert(
1972 "web_search".to_string(),
1973 ai_agents_tools::ToolPolicyConfig {
1974 max_results: Some(0),
1975 ..Default::default()
1976 },
1977 );
1978
1979 let error = AgentBuilder::from_spec(spec)
1980 .llm(Arc::new(MockLLMProvider::new("test")))
1981 .build()
1982 .unwrap_err();
1983 assert!(
1984 error
1985 .to_string()
1986 .contains("max_results must be greater than 0")
1987 );
1988 }
1989
1990 #[test]
1991 fn test_builder_from_spec_with_base_dir_preserves_mutations() {
1992 let yaml = r#"
1993name: OriginalAgent
1994system_prompt: "Original prompt"
1995max_iterations: 10
1996llm:
1997 provider: openai
1998 model: gpt-4
1999"#;
2000 let mut spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();
2001 spec.name = "RewrittenAgent".to_string();
2002 spec.system_prompt = "Rewritten prompt".to_string();
2003 spec.max_iterations = 37;
2004 let base_dir = PathBuf::from("rewritten-agent-dir");
2005
2006 let builder = AgentBuilder::from_spec_with_base_dir(spec, &base_dir);
2007
2008 let stored_spec = builder.spec.as_ref().unwrap();
2009 assert_eq!(stored_spec.name, "RewrittenAgent");
2010 assert_eq!(stored_spec.system_prompt, "Rewritten prompt");
2011 assert_eq!(stored_spec.max_iterations, 37);
2012 assert_eq!(builder.system_prompt.as_deref(), Some("Rewritten prompt"));
2013 assert_eq!(builder.max_iterations, Some(37));
2014 assert_eq!(builder.yaml_dir.as_deref(), Some(base_dir.as_path()));
2015 }
2016
2017 #[tokio::test]
2018 async fn test_builder_from_spec_with_base_dir_resolves_spawner_paths() {
2019 use ai_agents_llm::mock::MockLLMProvider;
2020
2021 let base_dir = std::env::temp_dir().join(format!(
2022 "ai-agents-builder-base-dir-{}",
2023 uuid::Uuid::new_v4()
2024 ));
2025 let templates_dir = base_dir.join("templates");
2026 let agents_dir = base_dir.join("agents");
2027 std::fs::create_dir_all(&templates_dir).unwrap();
2028 std::fs::create_dir_all(&agents_dir).unwrap();
2029
2030 let template_content = "name: {{ name }}\nsystem_prompt: Template prompt\n";
2031 std::fs::write(templates_dir.join("worker.yaml"), template_content).unwrap();
2032 std::fs::write(
2033 agents_dir.join("child.yaml"),
2034 r#"
2035name: ChildAgent
2036system_prompt: "Child prompt"
2037llm:
2038 provider: definitely-not-a-provider
2039 model: unavailable
2040"#,
2041 )
2042 .unwrap();
2043
2044 let yaml = r#"
2045name: ParentAgent
2046system_prompt: "Parent prompt"
2047llm:
2048 provider: openai
2049 model: gpt-4
2050spawner:
2051 shared_llms: true
2052 templates:
2053 worker:
2054 path: templates/worker.yaml
2055 auto_spawn:
2056 - id: child
2057 agent: agents/child.yaml
2058"#;
2059 let spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();
2060
2061 let builder = AgentBuilder::from_spec_with_base_dir(spec, &base_dir)
2062 .llm(Arc::new(MockLLMProvider::new("test")))
2063 .auto_configure_spawner()
2064 .await
2065 .unwrap();
2066
2067 let template = builder
2068 .spawner
2069 .as_ref()
2070 .unwrap()
2071 .templates()
2072 .get("worker")
2073 .unwrap();
2074 assert_eq!(template.content, template_content);
2075 assert!(builder.spawner_registry.as_ref().unwrap().contains("child"));
2076
2077 std::fs::remove_dir_all(base_dir).unwrap();
2078 }
2079
2080 #[tokio::test]
2081 async fn test_builder_auto_spawn_fails_on_any_declared_child_error() {
2082 use ai_agents_llm::mock::MockLLMProvider;
2083
2084 let base_dir = std::env::temp_dir().join(format!(
2085 "ai-agents-builder-child-failure-{}",
2086 uuid::Uuid::new_v4()
2087 ));
2088 std::fs::create_dir_all(&base_dir).unwrap();
2089 std::fs::write(
2090 base_dir.join("valid.yaml"),
2091 "name: ValidChild\nsystem_prompt: valid\n",
2092 )
2093 .unwrap();
2094
2095 let yaml = r#"
2096name: ParentAgent
2097system_prompt: parent
2098llm:
2099 default: default
2100spawner:
2101 shared_llms: true
2102 auto_spawn:
2103 - id: valid
2104 agent: valid.yaml
2105 - id: missing
2106 agent: missing.yaml
2107"#;
2108 let spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();
2109 let mut registry = LLMRegistry::new();
2110 registry.register("default", Arc::new(MockLLMProvider::new("test")));
2111 registry.set_default("default");
2112
2113 let error = AgentBuilder::from_spec_with_base_dir(spec, &base_dir)
2114 .llm_registry(registry)
2115 .auto_configure_spawner()
2116 .await
2117 .err()
2118 .unwrap()
2119 .to_string();
2120 assert!(error.contains("missing"), "{error}");
2121 assert!(error.contains("missing.yaml"), "{error}");
2122
2123 std::fs::remove_dir_all(base_dir).unwrap();
2124 }
2125
2126 #[test]
2127 fn test_builder_chain() {
2128 let builder = AgentBuilder::new()
2129 .system_prompt("Test prompt")
2130 .max_iterations(5)
2131 .max_context_tokens(4096);
2132
2133 assert_eq!(builder.system_prompt, Some("Test prompt".to_string()));
2134 assert_eq!(builder.max_iterations, Some(5));
2135 assert_eq!(builder.max_context_tokens, Some(4096));
2136 }
2137
2138 #[test]
2139 fn test_builder_skills() {
2140 use ai_agents_skills::{SkillDefinition, SkillStep};
2141
2142 let skill = SkillDefinition {
2143 id: "test".to_string(),
2144 description: "Test skill".to_string(),
2145 trigger: "When testing".to_string(),
2146 steps: vec![SkillStep::Prompt {
2147 prompt: "Hello".to_string(),
2148 llm: None,
2149 }],
2150 reasoning: None,
2151 reflection: None,
2152 disambiguation: None,
2153 };
2154
2155 let builder = AgentBuilder::new().skill(skill.clone()).skills(vec![skill]);
2156
2157 assert_eq!(builder.skills.len(), 2);
2158 }
2159
2160 #[test]
2161 fn test_builder_from_yaml_with_states() {
2162 let yaml = r#"
2163name: StatefulAgent
2164system_prompt: "You are helpful."
2165llm:
2166 provider: openai
2167 model: gpt-4
2168states:
2169 initial: greeting
2170 states:
2171 greeting:
2172 prompt: "Welcome!"
2173 transitions:
2174 - to: support
2175 when: "user needs help"
2176 support:
2177 prompt: "How can I help?"
2178"#;
2179 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2180 assert!(builder.spec.is_some());
2181 let spec = builder.spec.as_ref().unwrap();
2182 assert!(spec.has_states());
2183 assert!(spec.states.is_some());
2184 let states = spec.states.as_ref().unwrap();
2185 assert_eq!(states.initial, "greeting");
2186 assert_eq!(states.states.len(), 2);
2187 }
2188
2189 #[test]
2190 fn test_builder_from_yaml_with_context() {
2191 let yaml = r#"
2192name: ContextAgent
2193system_prompt: "Hello, {{ context.user.name }}!"
2194llm:
2195 provider: openai
2196 model: gpt-4
2197context:
2198 user:
2199 type: runtime
2200 required: true
2201 time:
2202 type: builtin
2203 source: datetime
2204 refresh: per_turn
2205"#;
2206 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2207 assert!(builder.spec.is_some());
2208 let spec = builder.spec.as_ref().unwrap();
2209 assert!(spec.has_context());
2210 assert_eq!(spec.context.len(), 2);
2211 assert!(spec.context.contains_key("user"));
2212 assert!(spec.context.contains_key("time"));
2213 }
2214
2215 #[test]
2216 fn test_builder_from_yaml_with_full_v04_features() {
2217 let yaml = r#"
2218name: FullFeaturedAgent
2219version: "0.4.0"
2220system_prompt: |
2221 You are a helpful assistant.
2222 User: {{ context.user.name }}
2223 Language: {{ context.user.language }}
2224llm:
2225 provider: openai
2226 model: gpt-4
2227context:
2228 user:
2229 type: runtime
2230 required: true
2231 default:
2232 name: "Guest"
2233 language: "en"
2234 time:
2235 type: builtin
2236 source: datetime
2237 refresh: per_turn
2238states:
2239 initial: greeting
2240 states:
2241 greeting:
2242 prompt: "Welcome to our service!"
2243 prompt_mode: append
2244 transitions:
2245 - to: support
2246 when: "user needs help"
2247 auto: true
2248 priority: 10
2249 support:
2250 prompt: "I'm here to help you."
2251 max_turns: 5
2252 timeout_to: escalation
2253 transitions:
2254 - to: closing
2255 when: "issue resolved"
2256 auto: true
2257 escalation:
2258 prompt: "Let me connect you with a human agent."
2259 closing:
2260 prompt: "Thank you for using our service!"
2261"#;
2262 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2263 assert!(builder.spec.is_some());
2264 let spec = builder.spec.as_ref().unwrap();
2265
2266 assert!(spec.has_context());
2268 assert_eq!(spec.context.len(), 2);
2269
2270 assert!(spec.has_states());
2272 let states = spec.states.as_ref().unwrap();
2273 assert_eq!(states.initial, "greeting");
2274 assert_eq!(states.states.len(), 4);
2275
2276 let greeting = states.states.get("greeting").unwrap();
2278 assert!(greeting.prompt.is_some());
2279 assert_eq!(greeting.transitions.len(), 1);
2280 assert_eq!(greeting.transitions[0].to, "support");
2281 assert!(greeting.transitions[0].auto);
2282
2283 let support = states.states.get("support").unwrap();
2285 assert_eq!(support.max_turns, Some(5));
2286 assert_eq!(support.timeout_to, Some("escalation".to_string()));
2287 }
2288
2289 #[test]
2290 fn test_builder_from_yaml_with_hitl() {
2291 let yaml = r#"
2292name: HITLAgent
2293system_prompt: "You are a secure assistant."
2294llm:
2295 provider: openai
2296 model: gpt-4
2297hitl:
2298 default_timeout_seconds: 600
2299 on_timeout: reject
2300 tools:
2301 send_payment:
2302 require_approval: true
2303 approval_context:
2304 - amount
2305 - recipient
2306 approval_message: "Approve payment?"
2307 delete_record:
2308 require_approval: true
2309 conditions:
2310 - name: high_value
2311 when: "amount > 1000"
2312 require_approval: true
2313 approval_message: "High value transaction"
2314 states:
2315 escalation:
2316 on_enter: require_approval
2317 approval_message: "Escalate to human?"
2318"#;
2319 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2320 assert!(builder.spec.is_some());
2321 let spec = builder.spec.as_ref().unwrap();
2322
2323 assert!(spec.has_hitl());
2324 let hitl = spec.hitl.as_ref().unwrap();
2325 assert_eq!(hitl.default_timeout_seconds, 600);
2326 assert_eq!(hitl.tools.len(), 2);
2327 assert!(hitl.tools.get("send_payment").unwrap().require_approval);
2328 assert_eq!(hitl.conditions.len(), 1);
2329 assert_eq!(hitl.conditions[0].name, "high_value");
2330 assert_eq!(hitl.states.len(), 1);
2331 }
2332
2333 #[test]
2334 fn test_builder_from_yaml_with_compacting_memory() {
2335 let yaml = r#"
2336name: CompactingAgent
2337system_prompt: "You are a helpful assistant."
2338memory:
2339 type: compacting
2340 max_messages: 100
2341 max_recent_messages: 20
2342 compress_threshold: 15
2343 summarize_batch_size: 5
2344 summarizer_llm: router
2345llm:
2346 provider: openai
2347 model: gpt-4
2348"#;
2349 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2350 assert!(builder.spec.is_some());
2351 let spec = builder.spec.as_ref().unwrap();
2352
2353 assert!(spec.memory.is_compacting());
2354 assert_eq!(spec.memory.max_recent_messages, Some(20));
2355 assert_eq!(spec.memory.compress_threshold, Some(15));
2356 assert_eq!(spec.memory.summarize_batch_size, Some(5));
2357 assert_eq!(spec.memory.summarizer_llm, Some("router".to_string()));
2358
2359 let compacting_config = spec.memory.to_compacting_config();
2360 assert_eq!(compacting_config.max_recent_messages, 20);
2361 assert_eq!(compacting_config.compress_threshold, 15);
2362 assert_eq!(compacting_config.summarize_batch_size, 5);
2363 }
2364
2365 #[test]
2366 fn test_builder_from_yaml_with_token_budget() {
2367 let yaml = r#"
2368name: BudgetAgent
2369system_prompt: "You are a helpful assistant."
2370memory:
2371 type: compacting
2372 max_messages: 100
2373 token_budget:
2374 total: 8192
2375 allocation:
2376 summary: 2048
2377 recent_messages: 4096
2378 facts: 1024
2379 overflow_strategy: summarize_more
2380 warn_at_percent: 75
2381llm:
2382 provider: openai
2383 model: gpt-4
2384"#;
2385 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2386 assert!(builder.spec.is_some());
2387 let spec = builder.spec.as_ref().unwrap();
2388
2389 assert!(spec.memory.token_budget.is_some());
2390 let budget = spec.memory.token_budget.as_ref().unwrap();
2391 assert_eq!(budget.total, 8192);
2392 assert_eq!(budget.allocation.summary, 2048);
2393 assert_eq!(budget.allocation.recent_messages, 4096);
2394 assert_eq!(budget.allocation.facts, 1024);
2395 assert_eq!(budget.warn_at_percent, 75);
2396 }
2397
2398 #[test]
2399 fn test_builder_from_yaml_with_overflow_strategies() {
2400 use ai_agents_memory::OverflowStrategy;
2401
2402 let yaml = r#"
2403name: TruncateAgent
2404system_prompt: "You are helpful."
2405memory:
2406 type: compacting
2407 token_budget:
2408 total: 4096
2409 overflow_strategy: truncate_oldest
2410llm:
2411 provider: openai
2412 model: gpt-4
2413"#;
2414 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2415 let budget = builder
2416 .spec
2417 .as_ref()
2418 .unwrap()
2419 .memory
2420 .token_budget
2421 .as_ref()
2422 .unwrap();
2423 assert_eq!(budget.overflow_strategy, OverflowStrategy::TruncateOldest);
2424
2425 let yaml = r#"
2426name: ErrorAgent
2427system_prompt: "You are helpful."
2428memory:
2429 type: compacting
2430 token_budget:
2431 total: 4096
2432 overflow_strategy: error
2433llm:
2434 provider: openai
2435 model: gpt-4
2436"#;
2437 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2438 let budget = builder
2439 .spec
2440 .as_ref()
2441 .unwrap()
2442 .memory
2443 .token_budget
2444 .as_ref()
2445 .unwrap();
2446 assert_eq!(budget.overflow_strategy, OverflowStrategy::Error);
2447 }
2448
2449 #[test]
2450 fn test_builder_from_yaml_with_storage_file() {
2451 let yaml = r#"
2452name: PersistentAgent
2453system_prompt: "You are helpful."
2454llm:
2455 provider: openai
2456 model: gpt-4
2457storage:
2458 type: file
2459 path: "./data/sessions"
2460"#;
2461 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2462 let spec = builder.spec.as_ref().unwrap();
2463 assert!(spec.has_storage());
2464 assert!(spec.storage.is_file());
2465 assert_eq!(spec.storage.get_path(), Some("./data/sessions"));
2466 }
2467
2468 #[test]
2469 fn test_builder_from_yaml_with_storage_sqlite() {
2470 let yaml = r#"
2471name: PersistentAgent
2472system_prompt: "You are helpful."
2473llm:
2474 provider: openai
2475 model: gpt-4
2476storage:
2477 type: sqlite
2478 path: "./data/sessions.db"
2479"#;
2480 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2481 let spec = builder.spec.as_ref().unwrap();
2482 assert!(spec.has_storage());
2483 assert!(spec.storage.is_sqlite());
2484 }
2485
2486 #[test]
2487 fn test_builder_from_yaml_with_storage_redis() {
2488 let yaml = r#"
2489name: PersistentAgent
2490system_prompt: "You are helpful."
2491llm:
2492 provider: openai
2493 model: gpt-4
2494storage:
2495 type: redis
2496 url: "redis://localhost:6379"
2497 prefix: "myagent:"
2498 ttl_seconds: 86400
2499"#;
2500 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2501 let spec = builder.spec.as_ref().unwrap();
2502 assert!(spec.has_storage());
2503 assert!(spec.storage.is_redis());
2504 assert_eq!(spec.storage.get_prefix(), "myagent:");
2505 assert_eq!(spec.storage.get_ttl(), Some(86400));
2506 }
2507
2508 #[test]
2509 fn test_builder_no_storage_by_default() {
2510 let yaml = r#"
2511name: SimpleAgent
2512system_prompt: "You are helpful."
2513llm:
2514 provider: openai
2515 model: gpt-4
2516"#;
2517 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2518 let spec = builder.spec.as_ref().unwrap();
2519 assert!(!spec.has_storage());
2520 }
2521
2522 #[test]
2523 fn test_build_fails_on_missing_declared_tool() {
2524 use ai_agents_llm::mock::MockLLMProvider;
2525
2526 let yaml = r#"
2527name: ToolAgent
2528system_prompt: "You are helpful."
2529llm:
2530 provider: openai
2531 model: gpt-4
2532tools:
2533 - name: lookup_order
2534 - name: calculator
2535"#;
2536 let llm = Arc::new(MockLLMProvider::new("test"));
2537 let result = AgentBuilder::from_yaml(yaml)
2538 .unwrap()
2539 .llm(llm)
2540 .auto_configure_features()
2541 .unwrap()
2542 .build();
2544
2545 assert!(result.is_err());
2546 let err = result.unwrap_err().to_string();
2547 assert!(
2548 err.contains("lookup_order"),
2549 "error should name the missing tool: {}",
2550 err
2551 );
2552 assert!(
2554 !err.contains("calculator"),
2555 "calculator is registered, should not be missing: {}",
2556 err
2557 );
2558 }
2559
2560 #[test]
2561 fn test_build_succeeds_when_declared_tool_is_registered() {
2562 use ai_agents_core::Tool;
2563 use ai_agents_llm::mock::MockLLMProvider;
2564
2565 struct FakeTool;
2566 #[async_trait::async_trait]
2567 impl Tool for FakeTool {
2568 fn id(&self) -> &str {
2569 "lookup_order"
2570 }
2571 fn name(&self) -> &str {
2572 "Order Lookup"
2573 }
2574 fn description(&self) -> &str {
2575 "Look up an order"
2576 }
2577 fn input_schema(&self) -> serde_json::Value {
2578 serde_json::json!({})
2579 }
2580 async fn execute(
2581 &self,
2582 _args: serde_json::Value,
2583 _ctx: ai_agents_core::ToolExecutionContext,
2584 ) -> ai_agents_core::ToolResult {
2585 ai_agents_core::ToolResult::ok("ok")
2586 }
2587 }
2588
2589 let yaml = r#"
2590name: ToolAgent
2591system_prompt: "You are helpful."
2592llm:
2593 provider: openai
2594 model: gpt-4
2595tools:
2596 - name: lookup_order
2597 - name: calculator
2598"#;
2599 let llm = Arc::new(MockLLMProvider::new("test"));
2600 let result = AgentBuilder::from_yaml(yaml)
2601 .unwrap()
2602 .llm(llm)
2603 .auto_configure_features()
2604 .unwrap()
2605 .tool(Arc::new(FakeTool))
2606 .build();
2607
2608 assert!(
2609 result.is_ok(),
2610 "build should succeed when all declared tools are registered: {:?}",
2611 result.err()
2612 );
2613 }
2614
2615 #[test]
2616 fn test_spawner_config_deserializes_shared_storage() {
2617 let yaml = r#"
2618name: TestAgent
2619system_prompt: "Test"
2620llm:
2621 provider: openai
2622 model: gpt-4
2623spawner:
2624 shared_llms: true
2625 shared_storage:
2626 type: sqlite
2627 path: ./data/test.db
2628 max_agents: 5
2629"#;
2630 let builder = AgentBuilder::from_yaml(yaml).unwrap();
2631 let spec = builder.spec.as_ref().unwrap();
2632 let sc = spec.spawner.as_ref().unwrap();
2633 assert!(sc.shared_storage.is_some());
2634 assert!(sc.shared_storage.as_ref().unwrap().is_sqlite());
2635 }
2636
2637 #[test]
2638 fn test_build_succeeds_with_no_tools_declared() {
2639 use ai_agents_llm::mock::MockLLMProvider;
2640
2641 let yaml = r#"
2642name: SimpleAgent
2643system_prompt: "You are helpful."
2644llm:
2645 provider: openai
2646 model: gpt-4
2647"#;
2648 let llm = Arc::new(MockLLMProvider::new("test"));
2649 let result = AgentBuilder::from_yaml(yaml)
2650 .unwrap()
2651 .llm(llm)
2652 .auto_configure_features()
2653 .unwrap()
2654 .build();
2655
2656 assert!(
2657 result.is_ok(),
2658 "no tools: section means no validation needed: {:?}",
2659 result.err()
2660 );
2661 }
2662}