Skip to main content

ai_agents_runtime/
builder.rs

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
51/// Builds alias-to-model metadata used by observed LLM wrappers.
52fn 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
67/// Wraps every provider in a registry while preserving aliases and router/default settings.
68fn 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    /// Builds from an existing spec while resolving relative resources from the supplied directory.
224    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    /// Constructs declared providers without consuming auxiliary routing until registrations are complete.
280    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                // base_url: first-class field, fallback to extra for backward compat
311                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                // api_key: resolve from api_key_env if specified
320                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            // base_url: first-class field, fallback to extra for backward compat
370            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            // api_key: resolve from api_key_env if specified
379            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    /// Auto-configure recovery, tool security, process pipeline, and built-in tools from the spec.
405    ///
406    /// **Call order matters for tools**: this method only registers built-in tools when `self.tools` is `None`.
407    /// If `.tool()` or `.tools()` was called before this, `self.tools` is already `Some` and built-ins will NOT be added.
408    ///
409    /// Correct:
410    /// ```ignore
411    /// .auto_configure_features()?   // registers built-ins (self.tools was None)
412    /// .tool(Arc::new(MyTool))       // adds MyTool into the builtin registry
413    /// ```
414    ///
415    /// Wrong — built-ins are lost:
416    /// ```ignore
417    /// .tool(Arc::new(MyTool))       // self.tools = Some(empty + MyTool)
418    /// .auto_configure_features()?   // self.tools is Some -> skips builtin registration
419    /// ```
420    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            // Auto-register builtin tools if the user hasn't provided a custom registry.
434            if self.tools.is_none() {
435                self.tools = Some(create_builtin_registry());
436            }
437        }
438        Ok(self)
439    }
440
441    /// Initialize MCP wrapper tools from `tools:` entries with `type: mcp`.
442    ///
443    /// Each MCP entry becomes an `MCPWrapperTool` registered as a normal builtin
444    /// tool in the `ToolRegistry`. Views defined in the entry's `views:` field are registered as separate `MCPViewTool` instances sharing the parent's MCP connection.
445    ///
446    /// Call this after `auto_configure_features()` so the tool registry exists.
447    pub async fn auto_configure_mcp(mut self) -> Result<Self> {
448        if let Some(ref spec) = self.spec {
449            // Collect MCP configs from tools: entries with type: mcp
450            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                    // Initialize with timeout
472                    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                            // Register the parent tool
488                            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                            // Register views as separate tools sharing the parent connection
496                            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    /// Registers a main provider; changes after hierarchy consumers are frozen fail at the next gate.
545    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    /// Registers an exact alias without validating an incomplete builder configuration.
554    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    /// Set a raw LLM registry that may still need observability wrapping.
568    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    /// Accepts inherited handles while retaining child-local settings at the consuming gate.
578    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    /// Replace the entire tool registry.
594    ///
595    /// If `auto_configure_features()` was called before this, the auto-registered builtins will be overwritten.
596    pub fn tools(mut self, tools: ToolRegistry) -> Self {
597        self.tools = Some(tools);
598        self
599    }
600
601    /// Register a single tool into the existing registry.
602    ///
603    /// If no registry exists yet, creates an empty one first.
604    /// Use this to add custom tools on top of auto-configured builtins.
605    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    /// Merge tools from another registry into the existing one.
612    ///
613    /// Skips tools whose ID already exists (no overwrite).
614    /// If no registry exists yet, creates an empty one first.
615    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
628    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
637    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
646    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    /// Configures recovery locals before auxiliary consumers freeze their construction snapshot.
681    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
695    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
718    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
727    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
746    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
765    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    /// Updates configured local routing inputs before the hierarchy consuming gate freezes them.
774    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    /// Set persona config directly (overrides spec).
783    pub fn persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
784        self.persona_manager = Some(manager);
785        self
786    }
787
788    /// Provide a shared persona template registry.
789    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    /// Provide a shared observability manager instead of creating one from YAML.
798    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    /// Creates or reuses the manager before components retain provider handles.
807    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    /// Loads pricing_file relative to the YAML directory before manager creation.
827    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    /// Prepares external definitions once and validates hierarchy before consumers capture providers.
838    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    /// Supplies the same prepared registry to additive construction adapters.
964    /// Reports framework wrapping status to avoid observing inherited handles twice.
965    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    /// Wraps the builder registry once and refreshes any stored process processor registry.
974    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    /// Wire spawner tools when the spec has a `spawner:` section.
1006    /// Call after `auto_configure_llms()` and `auto_configure_features()`.
1007    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        // Resolve file-path templates against the parent YAML directory.
1071        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        // Resolve shared storage from YAML config into a live backend.
1081        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                // Auto-inject into parent when no explicit storage: is configured.
1087                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(&registry));
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(&registry),
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        // Register orchestration tools if configured.
1123        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(&registry),
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        // Validate that all orchestration state references have matching agents.
1165        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    /// Finalizes one consistent hierarchy before runtime components capture selected handles.
1196    pub fn build(mut self) -> Result<RuntimeAgent> {
1197        self.prepare_routing(false)?;
1198        let resource_locks = self.shared_resource_locks();
1199        // Capture actor memory and facts configs before partial moves of spec consume fields.
1200        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        // ERROR NOTE: Don't include tools prompt here
1226        // - it will be added AFTER template rendering in get_effective_system_prompt() to avoid Jinja2 parsing JSON braces
1227        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        // Create memory after LLM registry is ready (needed for CompactingMemory summarizer)
1322        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        // Configure persona before freezing tools (evolve tool may need registration).
1354        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        // Register persona_evolve tool if allow_llm_evolve is true.
1379        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        // Build the effective tool grant.
1446        // YAML top-level tools are explicit ordinary grants, while feature flags such as spawner management, persona evolution, and orchestration tools are explicit feature grants.
1447        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                // management_tools and orchestration_tools are explicit registration and grant signals.
1484                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                // allow_llm_evolve is an explicit registration and grant signal.
1493                ids.push("persona_evolve".to_string());
1494            }
1495
1496            ids.sort();
1497            ids.dedup();
1498            ids
1499        });
1500
1501        // Validate: every effective non-MCP tool grant must exist in the registry.
1502        // MCP tools are excluded because they are registered via auto_configure_mcp() which
1503        // may or may not have been called (and MCP view names are synthetic).
1504        // Feature grants such as spawner management tools, orchestration tools, and persona_evolve must be registered before build.
1505        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        // Configure state machine from spec or builder
1598        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        // Configure context manager from spec or builder
1628        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        // Configure parallel tools and streaming from spec
1642        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            // Configure memory token budget if specified
1653            if let Some(ref budget) = spec.memory.token_budget {
1654                agent = agent.with_memory_token_budget(budget.clone());
1655            }
1656
1657            // Configure storage from spec if not explicitly set
1658            if self.storage_config.is_none() && spec.has_storage() {
1659                agent = agent.with_storage_config(spec.storage.clone());
1660            }
1661        }
1662
1663        // Configure storage from builder
1664        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        // Wire spawner handles into the agent so CLI can access registry.
1672        if let (Some(spawner), Some(registry)) = (self.spawner, self.spawner_registry) {
1673            agent = agent.with_spawner_handles(spawner, registry);
1674        }
1675
1676        // Configure hooks
1677        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        // Configure HITL from spec or builder
1696        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        // Configure disambiguation from spec
1724        if let Some(ref spec) = self.spec
1725            && spec.disambiguation.is_enabled()
1726        {
1727            agent = agent.with_disambiguation(spec.disambiguation.clone());
1728        }
1729
1730        // Wire persona manager into the agent (created earlier before tools_arc).
1731        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        // Store actor memory and facts configs on the agent now.
1740        // The actual FactStore and extractor are created lazily in init_storage()
1741        // once the storage backend is available. This avoids the sync/async
1742        // mismatch that caused the builder to silently skip facts setup when
1743        // storage was not yet initialized at build() time.
1744        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
1752/// Collect all agent IDs referenced by orchestration state fields.
1753fn 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        // Recurse into sub-states.
1785        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        // Check context
2267        assert!(spec.has_context());
2268        assert_eq!(spec.context.len(), 2);
2269
2270        // Check states
2271        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        // Check greeting state details
2277        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        // Check support state has timeout
2284        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            // NOT registering lookup_order — should fail
2543            .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        // calculator is built-in, so it should NOT appear in the error
2553        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}