ai_agents_runtime/

builder.rs

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::Arc;

use ai_agents_context::ContextManager;
use ai_agents_core::{AgentError, AgentStorage, LLMFeature, LLMProvider, Result, Tool};
use ai_agents_hitl::{ApprovalHandler, HITLEngine, RejectAllHandler};
use ai_agents_hooks::{AgentHooks, CompositeHooks};
use ai_agents_llm::LLMRegistry;
use ai_agents_llm::providers::{ProviderType, UnifiedLLMProvider};
use ai_agents_memory::{
    CompactingMemory, InMemoryStore, LLMSummarizer, Memory, NoopSummarizer, Summarizer,
};
use ai_agents_observability::{
    ObservabilityConfig, ObservabilityHooks, ObservabilityManager, ObservedLLMProvider,
    ObservedTool,
};
use ai_agents_process::ProcessProcessor;
use ai_agents_reasoning::{ReasoningConfig, ReflectionConfig};
use ai_agents_recovery::{MessageFilter, RecoveryManager};
use ai_agents_relationships::{
    RelationshipEvaluator, RelationshipEvaluatorTrait, RelationshipManager,
};
use ai_agents_skills::{SkillDefinition, SkillLoader};
use ai_agents_state::{LLMTransitionEvaluator, StateMachine, TransitionEvaluator};
use ai_agents_template::{TemplateInheritance, TemplateLoader, TemplateRenderer};
use ai_agents_tools::mcp::view::MCPViewTool;
use ai_agents_tools::mcp::wrapper::MCPWrapperTool;
use ai_agents_tools::{ToolRegistry, ToolSecurityEngine, create_builtin_registry};

use super::AgentInfo;
use super::StreamingConfig;
use super::runtime::{RuntimeAgent, ToolResourceLocks};
use crate::spec::{AgentSpec, StorageConfig};

fn feature_overrides_from_config(config: &crate::spec::LLMConfig) -> HashMap<LLMFeature, bool> {
    let mut overrides = HashMap::new();
    if let Some(enabled) = config.function_calling {
        overrides.insert(LLMFeature::FunctionCalling, enabled);
    }
    if let Some(enabled) = config.vision {
        overrides.insert(LLMFeature::Vision, enabled);
    }
    if let Some(enabled) = config.json_mode {
        overrides.insert(LLMFeature::JsonMode, enabled);
    }
    overrides
}

/// Builds alias-to-model metadata used by observed LLM wrappers.
fn model_by_alias_from_spec(spec: Option<&AgentSpec>) -> HashMap<String, String> {
    spec.map(|spec| {
        let mut models: HashMap<String, String> = spec
            .llms
            .iter()
            .map(|(alias, config)| (alias.clone(), config.model.clone()))
            .collect();
        if let Some(config) = spec.llm.as_config() {
            models.insert("default".to_string(), config.model.clone());
        }
        models
    })
    .unwrap_or_default()
}

/// Wraps every provider in a registry while preserving aliases and router/default settings.
fn wrap_registry_with_observability(
    registry: LLMRegistry,
    manager: Arc<ObservabilityManager>,
    model_by_alias: &HashMap<String, String>,
) -> LLMRegistry {
    registry.map_providers(|alias, provider| {
        let provider_name = provider.provider_name().to_string();
        let model = model_by_alias
            .get(alias)
            .cloned()
            .unwrap_or_else(|| alias.to_string());
        Arc::new(ObservedLLMProvider::new(
            provider,
            Arc::clone(&manager),
            Some(alias.to_string()),
            provider_name,
            model,
        )) as Arc<dyn LLMProvider>
    })
}

pub struct AgentBuilder {
    spec: Option<AgentSpec>,
    llm: Option<Arc<dyn LLMProvider>>,
    llm_registry: Option<LLMRegistry>,
    memory: Option<Arc<dyn Memory>>,
    tools: Option<ToolRegistry>,
    skills: Vec<SkillDefinition>,
    skill_loader: Option<SkillLoader>,
    yaml_dir: Option<PathBuf>,
    system_prompt: Option<String>,
    tools_prompt: Option<String>,
    auto_tools_prompt: bool,
    max_iterations: Option<u32>,
    max_context_tokens: Option<u32>,
    recovery_manager: Option<RecoveryManager>,
    tool_security: Option<ToolSecurityEngine>,
    process_processor: Option<ProcessProcessor>,
    message_filters: HashMap<String, Arc<dyn MessageFilter>>,
    context_manager: Option<Arc<ContextManager>>,
    state_machine: Option<Arc<StateMachine>>,
    transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
    hooks: Option<Arc<dyn AgentHooks>>,
    hitl_engine: Option<HITLEngine>,
    approval_handler: Option<Arc<dyn ApprovalHandler>>,
    storage_config: Option<StorageConfig>,
    storage: Option<Arc<dyn AgentStorage>>,
    reasoning: Option<ReasoningConfig>,
    reflection: Option<ReflectionConfig>,
    streaming: Option<StreamingConfig>,
    spawner: Option<Arc<crate::spawner::AgentSpawner>>,
    spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
    persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
    persona_templates: Option<Arc<ai_agents_persona::PersonaTemplateRegistry>>,
    observability_manager: Option<Arc<ObservabilityManager>>,
    resource_locks: Option<ToolResourceLocks>,
    llm_registry_observed: bool,
}

impl AgentBuilder {
    pub fn new() -> Self {
        Self {
            reasoning: None,
            reflection: None,
            spec: None,
            llm: None,
            llm_registry: None,
            memory: None,
            tools: None,
            skills: Vec::new(),
            skill_loader: None,
            yaml_dir: None,
            system_prompt: None,
            tools_prompt: None,
            auto_tools_prompt: true,
            max_iterations: None,
            max_context_tokens: None,
            recovery_manager: None,
            tool_security: None,
            process_processor: None,
            message_filters: HashMap::new(),
            context_manager: None,
            state_machine: None,
            transition_evaluator: None,
            hooks: None,
            hitl_engine: None,
            approval_handler: None,
            storage_config: None,
            storage: None,
            streaming: None,
            spawner: None,
            spawner_registry: None,
            persona_manager: None,
            persona_templates: None,
            observability_manager: None,
            resource_locks: None,
            llm_registry_observed: false,
        }
    }

    pub fn from_spec(spec: AgentSpec) -> Self {
        let system_prompt = spec.system_prompt.clone();
        let max_iterations = Some(spec.max_iterations);
        let max_context_tokens = Some(spec.max_context_tokens);
        let reasoning = Some(spec.reasoning.clone());
        let reflection = Some(spec.reflection.clone());

        Self {
            spec: Some(spec),
            llm: None,
            llm_registry: None,
            memory: None,
            tools: None,
            skills: Vec::new(),
            skill_loader: None,
            yaml_dir: None,
            system_prompt: Some(system_prompt),
            tools_prompt: None,
            auto_tools_prompt: true,
            max_iterations,
            max_context_tokens,
            recovery_manager: None,
            tool_security: None,
            process_processor: None,
            message_filters: HashMap::new(),
            context_manager: None,
            state_machine: None,
            transition_evaluator: None,
            hooks: None,
            hitl_engine: None,
            approval_handler: None,
            storage_config: None,
            storage: None,
            reasoning,
            reflection,
            streaming: None,
            spawner: None,
            spawner_registry: None,
            persona_manager: None,
            persona_templates: None,
            observability_manager: None,
            resource_locks: None,
            llm_registry_observed: false,
        }
    }

    /// Builds from an existing spec while resolving relative resources from the supplied directory.
    pub fn from_spec_with_base_dir(spec: AgentSpec, base_dir: impl Into<PathBuf>) -> Self {
        let mut builder = Self::from_spec(spec);
        builder.yaml_dir = Some(base_dir.into());
        builder
    }

    pub fn from_yaml(yaml_content: &str) -> Result<Self> {
        let spec = AgentSpec::from_yaml_strict(yaml_content)?;
        spec.validate()?;
        Ok(Self::from_spec(spec))
    }

    pub(crate) fn shared_resource_locks(&mut self) -> ToolResourceLocks {
        self.resource_locks
            .get_or_insert_with(|| Arc::new(parking_lot::RwLock::new(HashMap::new())))
            .clone()
    }

    pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
        self.resource_locks = Some(locks);
        self
    }

    pub fn from_yaml_file(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref();
        let content = std::fs::read_to_string(path).map_err(AgentError::IoError)?;
        let spec = AgentSpec::from_yaml_strict(&content)?;
        spec.validate()?;
        Ok(match path.parent() {
            Some(parent) => Self::from_spec_with_base_dir(spec, parent),
            None => Self::from_spec(spec),
        })
    }

    pub fn from_template(template_name: &str) -> Result<Self> {
        let loader = TemplateLoader::new();
        Self::from_template_with_loader(template_name, &loader)
    }

    pub fn from_template_with_loader(template_name: &str, loader: &TemplateLoader) -> Result<Self> {
        let renderer = TemplateRenderer::new();
        let variables = loader.variables();

        let load_and_render = |name: &str| -> Result<String> {
            let content = loader.load_template(name)?;
            renderer.render(&content, variables)
        };

        let rendered_root = load_and_render(template_name)?;
        let processed = TemplateInheritance::process(&rendered_root, load_and_render)?;
        let spec = AgentSpec::from_yaml_strict(&processed)?;
        spec.validate()?;
        Ok(Self::from_spec(spec))
    }

    pub fn auto_configure_llms(mut self) -> Result<Self> {
        let spec = self
            .spec
            .as_ref()
            .ok_or_else(|| AgentError::Config("Cannot auto-configure LLMs without spec".into()))?;

        if !spec.llms.is_empty() {
            let mut registry = LLMRegistry::new();

            for (alias, config) in &spec.llms {
                let provider_type = ProviderType::from_str(&config.provider)
                    .map_err(|e| AgentError::Config(e.to_string()))?;

                let core_config = ai_agents_core::LLMConfig {
                    temperature: Some(config.temperature),
                    max_tokens: Some(config.max_tokens),
                    top_p: config.top_p,
                    top_k: None,
                    frequency_penalty: None,
                    presence_penalty: None,
                    stop_sequences: None,
                    timeout_seconds: config.timeout_seconds,
                    reasoning: config.reasoning,
                    reasoning_effort: config.reasoning_effort.clone(),
                    reasoning_budget_tokens: config.reasoning_budget_tokens,
                    extra: config.extra.clone(),
                };
                // base_url: first-class field, fallback to extra for backward compat
                let base_url = config.base_url.clone().or_else(|| {
                    config
                        .extra
                        .get("base_url")
                        .and_then(|v| v.as_str())
                        .map(|s| s.to_string())
                });

                // api_key: resolve from api_key_env if specified
                let api_key = config
                    .api_key_env
                    .as_ref()
                    .and_then(|env_var| std::env::var(env_var).ok());

                let mut provider = UnifiedLLMProvider::from_spec_config(
                    provider_type,
                    &config.model,
                    api_key,
                    base_url,
                    core_config,
                )
                .map_err(|e| AgentError::LLM(e.to_string()))?
                .with_feature_overrides(feature_overrides_from_config(config));
                if let Some(choice) = config.tool_choice.clone() {
                    provider = provider.with_tool_choice(choice);
                }

                registry.register(alias, Arc::new(provider));
            }

            let default_alias = spec.llm.get_default_alias();
            let router_alias = spec.llm.get_router_alias();

            registry.set_default(&default_alias);
            if let Some(router) = router_alias {
                registry.set_router(&router);
            }

            self.llm_registry = Some(registry);
            self.llm_registry_observed = false;
        } else if let Some(config) = spec.llm.as_config() {
            let provider_type = ProviderType::from_str(&config.provider)
                .map_err(|e| AgentError::Config(e.to_string()))?;

            let core_config = ai_agents_core::LLMConfig {
                temperature: Some(config.temperature),
                max_tokens: Some(config.max_tokens),
                top_p: config.top_p,
                top_k: None,
                frequency_penalty: None,
                presence_penalty: None,
                stop_sequences: None,
                timeout_seconds: config.timeout_seconds,
                reasoning: config.reasoning,
                reasoning_effort: config.reasoning_effort.clone(),
                reasoning_budget_tokens: config.reasoning_budget_tokens,
                extra: config.extra.clone(),
            };
            // base_url: first-class field, fallback to extra for backward compat
            let base_url = config.base_url.clone().or_else(|| {
                config
                    .extra
                    .get("base_url")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string())
            });

            // api_key: resolve from api_key_env if specified
            let api_key = config
                .api_key_env
                .as_ref()
                .and_then(|env_var| std::env::var(env_var).ok());

            let mut provider = UnifiedLLMProvider::from_spec_config(
                provider_type,
                &config.model,
                api_key,
                base_url,
                core_config,
            )
            .map_err(|e| AgentError::LLM(e.to_string()))?
            .with_feature_overrides(feature_overrides_from_config(config));
            if let Some(choice) = config.tool_choice.clone() {
                provider = provider.with_tool_choice(choice);
            }

            self.llm = Some(Arc::new(provider));
            self.llm_registry_observed = false;
        }

        Ok(self)
    }

    /// Auto-configure recovery, tool security, process pipeline, and built-in tools from the spec.
    ///
    /// **Call order matters for tools**: this method only registers built-in tools when `self.tools` is `None`.
    /// If `.tool()` or `.tools()` was called before this, `self.tools` is already `Some` and built-ins will NOT be added.
    ///
    /// Correct:
    /// ```ignore
    /// .auto_configure_features()?   // registers built-ins (self.tools was None)
    /// .tool(Arc::new(MyTool))       // adds MyTool into the builtin registry
    /// ```
    ///
    /// Wrong — built-ins are lost:
    /// ```ignore
    /// .tool(Arc::new(MyTool))       // self.tools = Some(empty + MyTool)
    /// .auto_configure_features()?   // self.tools is Some -> skips builtin registration
    /// ```
    pub fn auto_configure_features(mut self) -> Result<Self> {
        if let Some(ref spec) = self.spec {
            self.recovery_manager = Some(RecoveryManager::new(spec.error_recovery.clone()));
            self.tool_security = Some(ToolSecurityEngine::try_new(spec.tool_security.clone())?);

            if spec.has_process() {
                let mut processor = ProcessProcessor::new(spec.process.clone());
                if let Some(ref registry) = self.llm_registry {
                    processor = processor.with_llm_registry(Arc::new(registry.clone()));
                }
                self.process_processor = Some(processor);
            }

            // Auto-register builtin tools if the user hasn't provided a custom registry.
            if self.tools.is_none() {
                self.tools = Some(create_builtin_registry());
            }
        }
        Ok(self)
    }

    /// Initialize MCP wrapper tools from `tools:` entries with `type: mcp`.
    ///
    /// Each MCP entry becomes an `MCPWrapperTool` registered as a normal builtin
    /// tool in the `ToolRegistry`. Views defined in the entry's `views:` field are registered as separate `MCPViewTool` instances sharing the parent's MCP connection.
    ///
    /// Call this after `auto_configure_features()` so the tool registry exists.
    pub async fn auto_configure_mcp(mut self) -> Result<Self> {
        if let Some(ref spec) = self.spec {
            // Collect MCP configs from tools: entries with type: mcp
            let mcp_configs: Vec<_> = spec
                .tools
                .as_ref()
                .map(|tools| {
                    tools
                        .iter()
                        .filter_map(|entry| entry.to_mcp_config())
                        .collect()
                })
                .unwrap_or_default();

            if !mcp_configs.is_empty() {
                let registry = self.tools.get_or_insert_with(create_builtin_registry);

                for config in mcp_configs {
                    let tool_name = config.name.clone();
                    let timeout_ms = config.startup_timeout_ms;
                    let views_config = config.views.clone();

                    let wrapper = MCPWrapperTool::new(config);

                    // Initialize with timeout
                    match tokio::time::timeout(
                        std::time::Duration::from_millis(timeout_ms),
                        wrapper.initialized(),
                    )
                    .await
                    {
                        Ok(Ok(initialized_tool)) => {
                            tracing::info!(
                                tool = %tool_name,
                                functions = initialized_tool.function_count(),
                                "MCP wrapper tool registered"
                            );

                            let parent = Arc::new(initialized_tool);

                            // Register the parent tool
                            registry.register(parent.clone()).map_err(|e| {
                                AgentError::Config(format!(
                                    "Failed to register MCP tool '{}': {}",
                                    tool_name, e
                                ))
                            })?;

                            // Register views as separate tools sharing the parent connection
                            for (view_name, view_config) in &views_config {
                                let view_tool = MCPViewTool::new(
                                    view_name.clone(),
                                    parent.clone(),
                                    view_config.functions.clone(),
                                    view_config.description.clone(),
                                )
                                .map_err(|e| {
                                    AgentError::Config(format!(
                                        "Failed to create MCP view '{}': {}",
                                        view_name, e
                                    ))
                                })?;

                                tracing::info!(
                                    view = %view_name,
                                    parent = %tool_name,
                                    functions = view_config.functions.len(),
                                    "MCP view tool registered"
                                );

                                registry.register(Arc::new(view_tool)).map_err(|e| {
                                    AgentError::Config(format!(
                                        "Failed to register MCP view '{}': {}",
                                        view_name, e
                                    ))
                                })?;
                            }
                        }
                        Ok(Err(e)) => {
                            return Err(AgentError::Config(format!(
                                "MCP tool '{}' initialization failed: {}",
                                tool_name, e
                            )));
                        }
                        Err(_) => {
                            return Err(AgentError::Config(format!(
                                "MCP tool '{}' timed out after {}ms",
                                tool_name, timeout_ms
                            )));
                        }
                    }
                }
            }
        }
        Ok(self)
    }

    pub fn llm(mut self, llm: Arc<dyn LLMProvider>) -> Self {
        self.llm = Some(llm);
        self
    }

    pub fn llm_alias(mut self, alias: impl Into<String>, provider: Arc<dyn LLMProvider>) -> Self {
        if self.llm_registry.is_none() {
            self.llm_registry = Some(LLMRegistry::new());
        }
        if let Some(ref mut registry) = self.llm_registry {
            registry.register(alias, provider);
        }
        self
    }

    /// Set a raw LLM registry that may still need observability wrapping.
    pub fn llm_registry(mut self, registry: LLMRegistry) -> Self {
        self.llm_registry = Some(registry);
        self.llm_registry_observed = false;
        self
    }

    pub(crate) fn authoritative_llm_registry(
        mut self,
        registry: LLMRegistry,
        observed: bool,
    ) -> Self {
        self.llm_registry = Some(registry);
        self.llm_registry_observed = observed;
        self
    }

    pub fn memory(mut self, memory: Arc<dyn Memory>) -> Self {
        self.memory = Some(memory);
        self
    }

    /// Replace the entire tool registry.
    ///
    /// If `auto_configure_features()` was called before this, the auto-registered builtins will be overwritten.
    pub fn tools(mut self, tools: ToolRegistry) -> Self {
        self.tools = Some(tools);
        self
    }

    /// Register a single tool into the existing registry.
    ///
    /// If no registry exists yet, creates an empty one first.
    /// Use this to add custom tools on top of auto-configured builtins.
    pub fn tool(mut self, tool: Arc<dyn Tool>) -> Self {
        let registry = self.tools.get_or_insert_with(ToolRegistry::new);
        let _ = registry.register(tool);
        self
    }

    /// Merge tools from another registry into the existing one.
    ///
    /// Skips tools whose ID already exists (no overwrite).
    /// If no registry exists yet, creates an empty one first.
    pub fn extend_tools(mut self, additional: ToolRegistry) -> Self {
        let registry = self.tools.get_or_insert_with(ToolRegistry::new);
        for id in additional.list_ids() {
            if registry.get(&id).is_none()
                && let Some(tool) = additional.get(&id)
            {
                let _ = registry.register(tool);
            }
        }
        self
    }

    pub fn skill(mut self, skill: SkillDefinition) -> Self {
        self.skills.push(skill);
        self
    }

    pub fn skills(mut self, skills: Vec<SkillDefinition>) -> Self {
        self.skills.extend(skills);
        self
    }

    pub fn skill_loader(mut self, loader: SkillLoader) -> Self {
        self.skill_loader = Some(loader);
        self
    }

    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(prompt.into());
        self
    }

    pub fn tools_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.tools_prompt = Some(prompt.into());
        self.auto_tools_prompt = false;
        self
    }

    pub fn auto_tools_prompt(mut self, auto: bool) -> Self {
        self.auto_tools_prompt = auto;
        self
    }

    pub fn max_iterations(mut self, max: u32) -> Self {
        self.max_iterations = Some(max);
        self
    }

    pub fn max_context_tokens(mut self, tokens: u32) -> Self {
        self.max_context_tokens = Some(tokens);
        self
    }

    pub fn recovery_manager(mut self, manager: RecoveryManager) -> Self {
        self.recovery_manager = Some(manager);
        self
    }

    pub fn tool_security(mut self, engine: ToolSecurityEngine) -> Self {
        self.tool_security = Some(engine);
        self
    }

    pub fn process_processor(mut self, processor: ProcessProcessor) -> Self {
        self.process_processor = Some(processor);
        self
    }

    pub fn message_filter(
        mut self,
        name: impl Into<String>,
        filter: Arc<dyn MessageFilter>,
    ) -> Self {
        self.message_filters.insert(name.into(), filter);
        self
    }

    pub fn context_manager(mut self, manager: Arc<ContextManager>) -> Self {
        self.context_manager = Some(manager);
        self
    }

    pub fn state_machine(mut self, machine: Arc<StateMachine>) -> Self {
        self.state_machine = Some(machine);
        self
    }

    pub fn transition_evaluator(mut self, evaluator: Arc<dyn TransitionEvaluator>) -> Self {
        self.transition_evaluator = Some(evaluator);
        self
    }

    pub fn hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
        self.hooks = Some(hooks);
        self
    }

    pub fn approval_handler(mut self, handler: Arc<dyn ApprovalHandler>) -> Self {
        self.approval_handler = Some(handler);
        self
    }

    pub fn hitl_engine(mut self, engine: HITLEngine) -> Self {
        self.hitl_engine = Some(engine);
        self
    }

    pub fn storage_config(mut self, config: StorageConfig) -> Self {
        self.storage_config = Some(config);
        self
    }

    pub fn storage(mut self, storage: Arc<dyn AgentStorage>) -> Self {
        self.storage = Some(storage);
        self
    }

    pub fn reasoning(mut self, config: ReasoningConfig) -> Self {
        self.reasoning = Some(config);
        self
    }

    pub fn reflection(mut self, config: ReflectionConfig) -> Self {
        self.reflection = Some(config);
        self
    }

    /// Set persona config directly (overrides spec).
    pub fn persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
        self.persona_manager = Some(manager);
        self
    }

    /// Provide a shared persona template registry.
    pub fn persona_templates(
        mut self,
        registry: Arc<ai_agents_persona::PersonaTemplateRegistry>,
    ) -> Self {
        self.persona_templates = Some(registry);
        self
    }

    /// Provide a shared observability manager instead of creating one from YAML.
    pub fn observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
        self.observability_manager = Some(manager);
        self
    }

    /// Creates or reuses the manager before components retain provider handles.
    fn ensure_observability_manager(&mut self) -> Result<Option<Arc<ObservabilityManager>>> {
        if let Some(manager) = self.observability_manager.as_ref() {
            return Ok(Some(Arc::clone(manager)));
        }
        let Some(ref spec) = self.spec else {
            return Ok(None);
        };
        if !spec.observability.enabled {
            return Ok(None);
        }
        let config = self.observability_config_with_pricing(&spec.observability)?;
        config
            .validate()
            .map_err(|e| AgentError::Config(e.to_string()))?;
        let manager = ObservabilityManager::new(config);
        self.observability_manager = Some(Arc::clone(&manager));
        Ok(Some(manager))
    }

    /// Loads pricing_file relative to the YAML directory before manager creation.
    fn observability_config_with_pricing(
        &self,
        config: &ObservabilityConfig,
    ) -> Result<ObservabilityConfig> {
        config
            .clone()
            .with_pricing_file_loaded(self.yaml_dir.as_deref())
            .map_err(|e| AgentError::Config(e.to_string()))
    }

    /// Wraps the builder registry once and refreshes any stored process processor registry.
    fn wrap_llm_registry_for_observability(&mut self) -> Result<()> {
        if self.llm_registry_observed {
            return Ok(());
        }
        let Some(manager) = self.ensure_observability_manager()? else {
            return Ok(());
        };
        let Some(registry) = self.llm_registry.take() else {
            return Ok(());
        };
        let model_by_alias = model_by_alias_from_spec(self.spec.as_ref());
        let wrapped = wrap_registry_with_observability(registry, manager, &model_by_alias);
        let wrapped_arc = Arc::new(wrapped.clone());
        if let Some(processor) = self.process_processor.take() {
            self.process_processor = Some(processor.with_llm_registry(wrapped_arc));
        }
        self.llm_registry = Some(wrapped);
        self.llm_registry_observed = true;
        Ok(())
    }

    pub fn streaming(mut self, enabled: bool) -> Self {
        let mut config = self.streaming.unwrap_or_default();
        config.enabled = enabled;
        self.streaming = Some(config);
        self
    }

    /// Wire spawner tools when the spec has a `spawner:` section.
    /// Call after `auto_configure_llms()` and `auto_configure_features()`.
    pub async fn auto_configure_spawner(mut self) -> Result<Self> {
        let spawner_config = match self.spec.as_ref().and_then(|s| s.spawner.as_ref()) {
            Some(c) => c.clone(),
            None => return Ok(self),
        };

        if spawner_config.shared_llms && self.llm_registry.is_none() {
            let provider = self.llm.as_ref().cloned().ok_or_else(|| {
                AgentError::Config(
                    "spawner.shared_llms requires the parent LLM provider to be configured"
                        .to_string(),
                )
            })?;
            let mut registry = LLMRegistry::new();
            registry.register("default", provider);
            registry.set_default("default");
            self.llm_registry = Some(registry);
            self.llm_registry_observed = false;
        }

        self.wrap_llm_registry_for_observability()?;
        let observability_manager = self.observability_manager.clone();

        use crate::spawner::{
            AgentRegistry, AgentSpawner,
            config::{configure_spawner_tools, resolve_templates},
        };

        let mut spawner = AgentSpawner::new();

        if let Some(ref manager) = observability_manager {
            spawner = spawner.with_observability(Arc::clone(manager));
        }
        spawner = spawner.with_resource_locks(self.shared_resource_locks());

        if spawner_config.shared_llms {
            let reg = self.llm_registry.as_ref().ok_or_else(|| {
                AgentError::Config(
                    "spawner.shared_llms requires the parent LLM registry to be configured"
                        .to_string(),
                )
            })?;
            spawner = if self.llm_registry_observed {
                spawner.with_shared_observed_llms(reg.clone())
            } else {
                spawner.with_shared_llms(reg.clone())
            };
        }

        if !spawner_config.shared_context.is_empty() {
            spawner = spawner.with_shared_context_map(spawner_config.shared_context.clone());
        }

        if let Some(max) = spawner_config.max_agents {
            spawner = spawner.with_max_agents(max);
        }

        if let Some(ref prefix) = spawner_config.name_prefix {
            spawner = spawner.with_name_prefix(prefix.clone())?;
        }

        // Resolve file-path templates against the parent YAML directory.
        if !spawner_config.templates.is_empty() {
            let resolved = resolve_templates(&spawner_config.templates, self.yaml_dir.as_deref())?;
            spawner = spawner.with_templates(resolved);
        }

        if let Some(ref allowed) = spawner_config.allowed_tools {
            spawner = spawner.with_allowed_tools(allowed.clone());
        }

        // Resolve shared storage from YAML config into a live backend.
        if let Some(ref sc) = spawner_config.shared_storage {
            let converted = crate::spec::storage::to_storage_config(sc);
            if let Some(st) = ai_agents_storage::create_storage(&converted).await? {
                spawner = spawner.with_shared_storage(Arc::clone(&st));

                // Auto-inject into parent when no explicit storage: is configured.
                let parent_has_storage = self.storage.is_some()
                    || self.storage_config.is_some()
                    || self.spec.as_ref().is_some_and(|s| s.has_storage());
                if !parent_has_storage {
                    self.storage = Some(st);
                }
            }
        }

        let spawner = Arc::new(spawner);
        let registry = Arc::new(AgentRegistry::new());

        self.spawner = Some(Arc::clone(&spawner));
        self.spawner_registry = Some(Arc::clone(&registry));
        let llm_for_tools = Arc::new(self.llm_registry.clone().unwrap_or_default());
        let agent_name = self
            .spec
            .as_ref()
            .map(|s| s.name.clone())
            .unwrap_or_default();

        let tools = configure_spawner_tools(
            Arc::clone(&spawner),
            Arc::clone(&registry),
            Arc::clone(&llm_for_tools),
            &agent_name,
        );

        let tool_registry = self.tools.get_or_insert_with(create_builtin_registry);
        for tool in tools {
            let _ = tool_registry.register(tool);
        }

        tracing::info!("Spawner tools registered");

        // Register orchestration tools if configured.
        if spawner_config.orchestration_tools.is_enabled() {
            let orch_tools = crate::orchestration::tools::configure_orchestration_tools(
                &spawner_config.orchestration_tools,
                Arc::clone(&registry),
                Arc::clone(&llm_for_tools),
            );
            let tool_registry = self.tools.get_or_insert_with(create_builtin_registry);
            for tool in orch_tools {
                let _ = tool_registry.register(tool);
            }
            tracing::info!("Orchestration tools registered");
        }

        for entry in &spawner_config.auto_spawn {
            let yaml_path = if let Some(ref dir) = self.yaml_dir {
                dir.join(&entry.agent)
            } else {
                std::path::PathBuf::from(&entry.agent)
            };

            tracing::info!(id = %entry.id, path = %yaml_path.display(), "Auto-spawning agent");
            let spawned = spawner
                .spawn_from_yaml_file_with_id(entry.id.clone(), &yaml_path)
                .await
                .map_err(|error| {
                    AgentError::Config(format!(
                        "Failed to auto-spawn agent '{}' from '{}': {}",
                        entry.id,
                        yaml_path.display(),
                        error
                    ))
                })?;
            registry.register(spawned).await.map_err(|error| {
                AgentError::Config(format!(
                    "Failed to register auto-spawned agent '{}': {}",
                    entry.id, error
                ))
            })?;
            tracing::info!(id = %entry.id, "Auto-spawned agent registered");
        }

        // Validate that all orchestration state references have matching agents.
        if let Some(ref spec) = self.spec
            && let Some(ref state_config) = spec.states
        {
            let refs = collect_orchestration_refs(&state_config.states);
            let mut missing: Vec<String> = Vec::new();

            for (agent_id, state_name, pattern) in &refs {
                if !registry.contains(agent_id) {
                    missing.push(format!(
                        "  - '{}' (referenced by state '{}' via {})",
                        agent_id, state_name, pattern
                    ));
                }
            }

            if !missing.is_empty() {
                missing.sort();
                missing.dedup();
                return Err(AgentError::Config(format!(
                    "Auto-spawn validation failed. These agents are referenced by \
                     orchestration states but were not successfully spawned:\n\n{}\n\n\
                     Check that agent YAML files exist and contain valid specs.",
                    missing.join("\n")
                )));
            }
        }

        Ok(self)
    }

    pub fn build(mut self) -> Result<RuntimeAgent> {
        let resource_locks = self.shared_resource_locks();
        // Capture actor memory and facts configs before partial moves of spec consume fields.
        let actor_memory_config = self
            .spec
            .as_ref()
            .and_then(|s| s.memory.actor_memory.clone());
        let facts_config = self.spec.as_ref().and_then(|s| s.memory.facts.clone());
        let relationships_config = self
            .spec
            .as_ref()
            .and_then(|s| s.memory.relationships.clone());

        let observability_manager = self.ensure_observability_manager()?;

        let base_prompt = self
            .system_prompt
            .ok_or_else(|| AgentError::Config("System prompt is required".into()))?;

        let mut tools = self.tools.unwrap_or_default();

        // ERROR NOTE: Don't include tools prompt here
        // - it will be added AFTER template rendering in get_effective_system_prompt() to avoid Jinja2 parsing JSON braces
        let system_prompt = base_prompt;

        let max_iterations = self.max_iterations.unwrap_or(10);

        let info = if let Some(ref spec) = self.spec {
            AgentInfo::new(&spec.name, &spec.name, &spec.version)
                .with_description(spec.description.clone().unwrap_or_default())
        } else {
            AgentInfo::new("agent", "Agent", "1.0.0")
        };

        if let Some(ref spec) = self.spec
            && !spec.skills.is_empty()
        {
            let mut loader = self.skill_loader.take().unwrap_or_default();
            if let Some(ref dir) = self.yaml_dir {
                loader.set_base_dir(dir);
            }
            let loaded_skills = loader.load_refs(&spec.skills)?;
            self.skills.extend(loaded_skills);
        }

        let mut llm_registry = self.llm_registry.unwrap_or_default();

        if let Some(llm) = self.llm
            && !llm_registry.has("default")
        {
            let provider = if self.llm_registry_observed {
                if let Some(ref manager) = observability_manager {
                    Arc::new(ObservedLLMProvider::new(
                        llm.clone(),
                        Arc::clone(manager),
                        Some("default".to_string()),
                        llm.provider_name().to_string(),
                        model_by_alias_from_spec(self.spec.as_ref())
                            .get("default")
                            .cloned()
                            .unwrap_or_else(|| "default".to_string()),
                    )) as Arc<dyn LLMProvider>
                } else {
                    llm.clone()
                }
            } else {
                llm.clone()
            };
            llm_registry.register("default", provider);
        }

        if let Some(ref spec) = self.spec {
            let default_alias = spec.llm.get_default_alias();
            let router_alias = spec.llm.get_router_alias();

            llm_registry.set_default(&default_alias);
            if let Some(router) = router_alias {
                llm_registry.set_router(&router);
            }
        }

        if llm_registry.is_empty() {
            return Err(AgentError::Config(
                "At least one LLM provider is required".into(),
            ));
        }

        if let Some(ref manager) = observability_manager
            && !self.llm_registry_observed
        {
            let model_by_alias = model_by_alias_from_spec(self.spec.as_ref());
            llm_registry = wrap_registry_with_observability(
                llm_registry,
                Arc::clone(manager),
                &model_by_alias,
            );
            self.llm_registry_observed = true;
        }

        // Create memory after LLM registry is ready (needed for CompactingMemory summarizer)
        let memory = self.memory.unwrap_or_else(|| {
            if let Some(ref spec) = self.spec {
                if spec.memory.is_compacting() {
                    let summarizer_llm = spec
                        .memory
                        .summarizer_llm
                        .as_ref()
                        .and_then(|alias| llm_registry.get(alias).ok())
                        .or_else(|| llm_registry.router().ok())
                        .or_else(|| llm_registry.default().ok());

                    let summarizer: Arc<dyn Summarizer> = match summarizer_llm {
                        Some(llm) => Arc::new(LLMSummarizer::new(llm)),
                        None => Arc::new(NoopSummarizer),
                    };
                    let config = spec.memory.to_compacting_config();
                    return Arc::new(CompactingMemory::new(summarizer, config));
                }
                Arc::new(InMemoryStore::new(spec.memory.max_messages))
            } else {
                Arc::new(InMemoryStore::new(100))
            }
        });

        // Configure persona before freezing tools (evolve tool may need registration).
        let persona_manager: Option<Arc<ai_agents_persona::PersonaManager>> =
            if let Some(pm) = self.persona_manager.take() {
                Some(pm)
            } else if let Some(ref spec) = self.spec {
                if spec.has_persona() {
                    let persona_config = spec.persona.clone().unwrap();
                    let renderer = ai_agents_context::TemplateRenderer::new();
                    let registry = self.persona_templates.clone();
                    let manager = ai_agents_persona::PersonaManager::from_config(
                        persona_config,
                        registry,
                        renderer,
                    )
                    .map_err(|e| {
                        AgentError::Config(format!("Failed to create PersonaManager: {}", e))
                    })?;
                    Some(Arc::new(manager))
                } else {
                    None
                }
            } else {
                None
            };

        // Register persona_evolve tool if allow_llm_evolve is true.
        if let Some(ref pm) = persona_manager
            && pm.should_register_evolve_tool()
        {
            let evolve_tool = ai_agents_persona::PersonaEvolveTool::new(pm.clone());
            let _ = tools.register(Arc::new(evolve_tool));
        }

        if let Some(ref manager) = observability_manager {
            tools = tools.map_tools(|tool| {
                Arc::new(ObservedTool::new(tool, Arc::clone(manager))) as Arc<dyn Tool>
            });
        }

        let relationship_manager: Option<Arc<RelationshipManager>> =
            if let Some(ref config) = relationships_config {
                if config.enabled {
                    let evaluator: Option<Arc<dyn RelationshipEvaluatorTrait>> =
                        if config.auto_update.enabled {
                            let llm = config
                                .auto_update
                                .llm
                                .as_ref()
                                .and_then(|alias| llm_registry.get(alias).ok())
                                .or_else(|| llm_registry.router().ok())
                                .or_else(|| llm_registry.default().ok());
                            llm.map(|llm| {
                                Arc::new(RelationshipEvaluator::new(llm))
                                    as Arc<dyn RelationshipEvaluatorTrait>
                            })
                        } else {
                            None
                        };
                    Some(Arc::new(RelationshipManager::from_config_with_evaluator(
                        config.clone(),
                        evaluator,
                    )?))
                } else {
                    None
                }
            } else {
                None
            };

        let tools_arc = Arc::new(tools);
        let llm_registry_arc = Arc::new(llm_registry);
        tools_arc.set_web_fetch_extractor(
            llm_registry_arc
                .router()
                .ok()
                .or_else(|| llm_registry_arc.default().ok()),
        );

        // Build the effective tool grant.
        // YAML top-level tools are explicit ordinary grants, while feature flags such as spawner management, persona evolution, and orchestration tools are explicit feature grants.
        let declared_tool_ids: Option<Vec<String>> = Some(if self.spec.is_none() {
            tools_arc.list_ids()
        } else {
            let mut ids: Vec<String> = self
                .spec
                .as_ref()
                .and_then(|s| s.tools.as_ref())
                .map(|tools| {
                    let mut ids: Vec<String> = tools
                        .iter()
                        .filter_map(|t| {
                            tools_arc
                                .canonical_id(t.name())
                                .or_else(|| Some(t.name().to_string()))
                        })
                        .collect();

                    for entry in tools {
                        if let Some(mcp_config) = entry.to_mcp_config() {
                            for view_name in mcp_config.views.keys() {
                                ids.push(
                                    tools_arc
                                        .canonical_id(view_name)
                                        .unwrap_or_else(|| view_name.clone()),
                                );
                            }
                        }
                    }

                    ids
                })
                .unwrap_or_default();

            if let Some(ref spec) = self.spec
                && let Some(ref spawner) = spec.spawner
            {
                // management_tools and orchestration_tools are explicit registration and grant signals.
                ids.extend(spawner.management_tools.granted_management_tool_ids());
                ids.extend(spawner.orchestration_tools.granted_orchestration_tool_ids());
            }

            if persona_manager
                .as_ref()
                .is_some_and(|pm| pm.should_register_evolve_tool())
            {
                // allow_llm_evolve is an explicit registration and grant signal.
                ids.push("persona_evolve".to_string());
            }

            ids.sort();
            ids.dedup();
            ids
        });

        // Validate: every effective non-MCP tool grant must exist in the registry.
        // MCP tools are excluded because they are registered via auto_configure_mcp() which
        // may or may not have been called (and MCP view names are synthetic).
        // Feature grants such as spawner management tools, orchestration tools, and persona_evolve must be registered before build.
        if let Some(ref ids) = declared_tool_ids {
            let mcp_names: Vec<String> = self
                .spec
                .as_ref()
                .and_then(|s| s.tools.as_ref())
                .map(|tools| {
                    let mut names = Vec::new();
                    for entry in tools {
                        if entry.is_mcp() {
                            names.push(entry.name().to_string());
                            if let Some(cfg) = entry.to_mcp_config() {
                                names.extend(cfg.views.keys().cloned());
                            }
                        }
                    }
                    names
                })
                .unwrap_or_default();

            let missing: Vec<&str> = ids
                .iter()
                .filter(|id| !mcp_names.contains(id))
                .filter(|id| tools_arc.get(id).is_none())
                .map(|s| s.as_str())
                .collect();

            if !missing.is_empty() {
                return Err(AgentError::Config(format!(
                    "Tools granted by YAML but not registered: [{}]. \
                     Register them via .tool(Arc::new(...)) or the matching auto_configure_* method before .build(), \
                     or remove the grant from YAML.",
                    missing.join(", ")
                )));
            }
        }

        let mut agent = RuntimeAgent::new(
            info,
            llm_registry_arc.clone(),
            memory,
            tools_arc,
            self.skills,
            system_prompt,
            max_iterations,
        )
        .with_shared_resource_locks(resource_locks)
        .with_declared_tool_ids(declared_tool_ids);

        if let Some(tokens) = self.max_context_tokens {
            agent = agent.with_max_context_tokens(tokens);
        }

        if let Some(manager) = self.recovery_manager {
            agent = agent.with_recovery_manager(manager);
        } else if let Some(ref spec) = self.spec {
            agent = agent.with_recovery_manager(RecoveryManager::new(spec.error_recovery.clone()));
        }

        if let Some(engine) = self.tool_security {
            agent = agent.with_tool_security(engine);
        } else if let Some(ref spec) = self.spec {
            agent =
                agent.with_tool_security(ToolSecurityEngine::try_new(spec.tool_security.clone())?);
        }

        if let Some(processor) = self.process_processor {
            agent =
                agent.with_process_processor(processor.with_llm_registry(llm_registry_arc.clone()));
        } else if let Some(ref spec) = self.spec
            && spec.has_process()
        {
            let processor = ProcessProcessor::new(spec.process.clone())
                .with_llm_registry(llm_registry_arc.clone());
            agent = agent.with_process_processor(processor);
        }

        for (name, filter) in self.message_filters {
            agent.register_message_filter(name, filter);
        }

        // Configure state machine from spec or builder
        if let Some(state_machine) = self.state_machine {
            let evaluator = self.transition_evaluator.unwrap_or_else(|| {
                let eval_llm = llm_registry_arc
                    .get("evaluator")
                    .or_else(|_| llm_registry_arc.router())
                    .or_else(|_| llm_registry_arc.default())
                    .expect("At least one LLM required for transition evaluator");
                Arc::new(LLMTransitionEvaluator::new(eval_llm))
            });
            agent = agent.with_state_machine(state_machine, evaluator);
        } else if let Some(ref spec) = self.spec
            && let Some(ref state_config) = spec.states
        {
            let state_machine = StateMachine::new(state_config.clone())?;
            let evaluator = self.transition_evaluator.unwrap_or_else(|| {
                let eval_llm = llm_registry_arc
                    .get("evaluator")
                    .or_else(|_| llm_registry_arc.router())
                    .or_else(|_| llm_registry_arc.default())
                    .expect("At least one LLM required for transition evaluator");
                Arc::new(LLMTransitionEvaluator::new(eval_llm))
            });
            agent = agent.with_state_machine(Arc::new(state_machine), evaluator);
        }

        // Configure context manager from spec or builder
        if let Some(context_manager) = self.context_manager {
            agent = agent.with_context_manager(context_manager);
        } else if let Some(ref spec) = self.spec
            && !spec.context.is_empty()
        {
            let context_manager = ContextManager::new(
                spec.context.clone(),
                spec.name.clone(),
                spec.version.clone(),
            );
            agent = agent.with_context_manager(Arc::new(context_manager));
        }

        // Configure parallel tools and streaming from spec
        if let Some(ref spec) = self.spec {
            agent = agent.with_parallel_tools(spec.parallel_tools.clone());
            let streaming_config = self
                .streaming
                .clone()
                .unwrap_or_else(|| spec.streaming.clone());
            agent = agent.with_streaming(streaming_config);

            agent = agent.with_runtime_config(spec.runtime.clone());

            // Configure memory token budget if specified
            if let Some(ref budget) = spec.memory.token_budget {
                agent = agent.with_memory_token_budget(budget.clone());
            }

            // Configure storage from spec if not explicitly set
            if self.storage_config.is_none() && spec.has_storage() {
                agent = agent.with_storage_config(spec.storage.clone());
            }
        }

        // Configure storage from builder
        if let Some(storage_config) = self.storage_config {
            agent = agent.with_storage_config(storage_config);
        }
        if let Some(storage) = self.storage {
            agent = agent.with_storage(storage);
        }

        // Wire spawner handles into the agent so CLI can access registry.
        if let (Some(spawner), Some(registry)) = (self.spawner, self.spawner_registry) {
            agent = agent.with_spawner_handles(spawner, registry);
        }

        // Configure hooks
        if let Some(manager) = observability_manager {
            agent = agent.with_observability(Arc::clone(&manager));
            let observability_hooks: Arc<dyn AgentHooks> =
                Arc::new(ObservabilityHooks::new(manager));
            let hooks: Arc<dyn AgentHooks> = if let Some(user_hooks) = self.hooks {
                Arc::new(
                    CompositeHooks::new()
                        .add(user_hooks)
                        .add(observability_hooks),
                )
            } else {
                observability_hooks
            };
            agent = agent.with_hooks(hooks);
        } else if let Some(hooks) = self.hooks {
            agent = agent.with_hooks(hooks);
        }

        // Configure HITL from spec or builder
        if let Some(hitl_engine) = self.hitl_engine {
            let handler = self
                .approval_handler
                .unwrap_or_else(|| Arc::new(RejectAllHandler::new()));
            agent = agent.with_hitl(hitl_engine, handler);
        } else if let Some(ref spec) = self.spec
            && let Some(ref hitl_config) = spec.hitl
        {
            let hitl_engine = HITLEngine::new(hitl_config.clone());
            let handler = self
                .approval_handler
                .unwrap_or_else(|| Arc::new(RejectAllHandler::new()));
            agent = agent.with_hitl(hitl_engine, handler);
        }

        if let Some(reasoning) = self.reasoning {
            agent = agent.with_reasoning(reasoning);
        } else if let Some(ref spec) = self.spec {
            agent = agent.with_reasoning(spec.reasoning.clone());
        }

        if let Some(reflection) = self.reflection {
            agent = agent.with_reflection(reflection);
        } else if let Some(ref spec) = self.spec {
            agent = agent.with_reflection(spec.reflection.clone());
        }

        // Configure disambiguation from spec
        if let Some(ref spec) = self.spec
            && spec.disambiguation.is_enabled()
        {
            agent = agent.with_disambiguation(spec.disambiguation.clone());
        }

        // Wire persona manager into the agent (created earlier before tools_arc).
        if let Some(pm) = persona_manager {
            agent = agent.with_persona(pm);
        }

        if let Some(relationship_manager) = relationship_manager {
            agent = agent.with_relationships(relationship_manager);
        }

        // Store actor memory and facts configs on the agent now.
        // The actual FactStore and extractor are created lazily in init_storage()
        // once the storage backend is available. This avoids the sync/async
        // mismatch that caused the builder to silently skip facts setup when
        // storage was not yet initialized at build() time.
        if actor_memory_config.is_some() || facts_config.is_some() {
            agent = agent.with_facts_config(actor_memory_config, facts_config);
        }

        Ok(agent)
    }
}

/// Collect all agent IDs referenced by orchestration state fields.
fn collect_orchestration_refs(
    states: &std::collections::HashMap<String, ai_agents_state::StateDefinition>,
) -> Vec<(String, String, &'static str)> {
    let mut refs = Vec::new();

    for (state_name, def) in states {
        if let Some(ref delegate_id) = def.delegate {
            refs.push((delegate_id.clone(), state_name.clone(), "delegate"));
        }
        if let Some(ref concurrent) = def.concurrent {
            for agent_ref in &concurrent.agents {
                refs.push((agent_ref.id().to_string(), state_name.clone(), "concurrent"));
            }
        }
        if let Some(ref gc) = def.group_chat {
            for participant in &gc.participants {
                refs.push((participant.id.clone(), state_name.clone(), "group_chat"));
            }
        }
        if let Some(ref pipeline) = def.pipeline {
            for stage in &pipeline.stages {
                refs.push((stage.id().to_string(), state_name.clone(), "pipeline"));
            }
        }
        if let Some(ref handoff) = def.handoff {
            refs.push((handoff.initial_agent.clone(), state_name.clone(), "handoff"));
            for agent_id in &handoff.available_agents {
                refs.push((agent_id.clone(), state_name.clone(), "handoff"));
            }
        }

        // Recurse into sub-states.
        if let Some(ref sub_states) = def.states {
            refs.extend(collect_orchestration_refs(sub_states));
        }
    }

    refs
}

impl Default for AgentBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builder_new() {
        let builder = AgentBuilder::new();
        assert!(builder.spec.is_none());
        assert!(builder.system_prompt.is_none());
    }

    #[test]
    fn test_builder_from_yaml() {
        let yaml = r#"
name: TestAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        assert_eq!(builder.spec.as_ref().unwrap().name, "TestAgent");
    }

    #[test]
    fn test_builder_from_yaml_rejects_nested_unknown_path() {
        let yaml = r#"
name: TestAgent
system_prompt: test
runtime:
  optimization:
    max_parallel_runtime_task: 4
"#;
        let error = match AgentBuilder::from_yaml(yaml) {
            Ok(_) => panic!("expected strict parse failure"),
            Err(error) => error.to_string(),
        };
        assert!(
            error.contains("runtime.optimization.max_parallel_runtime_task"),
            "{error}"
        );
    }

    #[test]
    fn test_feature_override_single_llm_builder_path() {
        let yaml = r#"
name: LocalAgent
system_prompt: "You are helpful."
llm:
  provider: ollama
  model: llama3.1
  function_calling: true
"#;
        let builder = AgentBuilder::from_yaml(yaml)
            .unwrap()
            .auto_configure_llms()
            .unwrap();

        let llm = builder.llm.as_ref().unwrap();
        assert!(llm.supports(LLMFeature::FunctionCalling));
    }

    #[test]
    fn test_feature_override_named_llms_builder_path() {
        let yaml = r#"
name: LocalAgent
system_prompt: "You are helpful."
llms:
  default:
    provider: openai-compatible
    model: qwen3:8b
    base_url: http://localhost:11434/v1
    json_mode: true
llm:
  default: default
"#;
        let builder = AgentBuilder::from_yaml(yaml)
            .unwrap()
            .auto_configure_llms()
            .unwrap();

        let registry = builder.llm_registry.as_ref().unwrap();
        let llm = registry.get("default").unwrap();
        assert!(llm.supports(LLMFeature::JsonMode));
    }

    #[test]
    fn test_builder_from_yaml_with_tool_security() {
        let yaml = r#"
name: SecureAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
max_context_tokens: 8192
error_recovery:
  default:
    max_retries: 5
tool_security:
  enabled: true
  tools:
    http:
      rate_limit: 10
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();
        assert_eq!(spec.max_context_tokens, 8192);
        assert_eq!(spec.error_recovery.default.max_retries, 5);
        assert!(spec.tool_security.enabled);
    }

    #[test]
    fn test_builder_from_yaml_rejects_zero_max_results() {
        let yaml = r#"
name: SecureAgent
system_prompt: "You are helpful."
tool_security:
  enabled: true
  tools:
    web_search:
      max_results: 0
"#;
        let error = AgentBuilder::from_yaml(yaml).err().unwrap();
        assert!(
            error
                .to_string()
                .contains("tool_security.tools.web_search.max_results must be greater than 0")
        );
    }

    #[test]
    fn test_builder_from_yaml_with_skills() {
        let yaml = r#"
name: SkillAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
skills:
  - id: greeting
    description: "Greet users"
    trigger: "When user says hello"
    steps:
      - prompt: "Hello!"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        assert!(!builder.spec.as_ref().unwrap().skills.is_empty());
    }

    #[test]
    fn test_builder_from_spec() {
        let spec = AgentSpec {
            name: "test".to_string(),
            version: "1.0".to_string(),
            description: Some("Test agent".to_string()),
            system_prompt: "You are helpful".to_string(),
            ..Default::default()
        };

        let builder = AgentBuilder::from_spec(spec);
        assert!(builder.spec.is_some());
        assert_eq!(builder.system_prompt, Some("You are helpful".to_string()));
    }

    #[test]
    fn test_builder_rejects_invalid_programmatic_tool_security() {
        use ai_agents_llm::mock::MockLLMProvider;

        let mut spec = AgentSpec::default();
        spec.tool_security.tools.insert(
            "web_search".to_string(),
            ai_agents_tools::ToolPolicyConfig {
                max_results: Some(0),
                ..Default::default()
            },
        );

        let error = AgentBuilder::from_spec(spec)
            .llm(Arc::new(MockLLMProvider::new("test")))
            .build()
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("max_results must be greater than 0")
        );
    }

    #[test]
    fn test_builder_from_spec_with_base_dir_preserves_mutations() {
        let yaml = r#"
name: OriginalAgent
system_prompt: "Original prompt"
max_iterations: 10
llm:
  provider: openai
  model: gpt-4
"#;
        let mut spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();
        spec.name = "RewrittenAgent".to_string();
        spec.system_prompt = "Rewritten prompt".to_string();
        spec.max_iterations = 37;
        let base_dir = PathBuf::from("rewritten-agent-dir");

        let builder = AgentBuilder::from_spec_with_base_dir(spec, &base_dir);

        let stored_spec = builder.spec.as_ref().unwrap();
        assert_eq!(stored_spec.name, "RewrittenAgent");
        assert_eq!(stored_spec.system_prompt, "Rewritten prompt");
        assert_eq!(stored_spec.max_iterations, 37);
        assert_eq!(builder.system_prompt.as_deref(), Some("Rewritten prompt"));
        assert_eq!(builder.max_iterations, Some(37));
        assert_eq!(builder.yaml_dir.as_deref(), Some(base_dir.as_path()));
    }

    #[tokio::test]
    async fn test_builder_from_spec_with_base_dir_resolves_spawner_paths() {
        use ai_agents_llm::mock::MockLLMProvider;

        let base_dir = std::env::temp_dir().join(format!(
            "ai-agents-builder-base-dir-{}",
            uuid::Uuid::new_v4()
        ));
        let templates_dir = base_dir.join("templates");
        let agents_dir = base_dir.join("agents");
        std::fs::create_dir_all(&templates_dir).unwrap();
        std::fs::create_dir_all(&agents_dir).unwrap();

        let template_content = "name: {{ name }}\nsystem_prompt: Template prompt\n";
        std::fs::write(templates_dir.join("worker.yaml"), template_content).unwrap();
        std::fs::write(
            agents_dir.join("child.yaml"),
            r#"
name: ChildAgent
system_prompt: "Child prompt"
llm:
  provider: definitely-not-a-provider
  model: unavailable
"#,
        )
        .unwrap();

        let yaml = r#"
name: ParentAgent
system_prompt: "Parent prompt"
llm:
  provider: openai
  model: gpt-4
spawner:
  shared_llms: true
  templates:
    worker:
      path: templates/worker.yaml
  auto_spawn:
    - id: child
      agent: agents/child.yaml
"#;
        let spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();

        let builder = AgentBuilder::from_spec_with_base_dir(spec, &base_dir)
            .llm(Arc::new(MockLLMProvider::new("test")))
            .auto_configure_spawner()
            .await
            .unwrap();

        let template = builder
            .spawner
            .as_ref()
            .unwrap()
            .templates()
            .get("worker")
            .unwrap();
        assert_eq!(template.content, template_content);
        assert!(builder.spawner_registry.as_ref().unwrap().contains("child"));

        std::fs::remove_dir_all(base_dir).unwrap();
    }

    #[tokio::test]
    async fn test_builder_auto_spawn_fails_on_any_declared_child_error() {
        use ai_agents_llm::mock::MockLLMProvider;

        let base_dir = std::env::temp_dir().join(format!(
            "ai-agents-builder-child-failure-{}",
            uuid::Uuid::new_v4()
        ));
        std::fs::create_dir_all(&base_dir).unwrap();
        std::fs::write(
            base_dir.join("valid.yaml"),
            "name: ValidChild\nsystem_prompt: valid\n",
        )
        .unwrap();

        let yaml = r#"
name: ParentAgent
system_prompt: parent
llm:
  default: default
spawner:
  shared_llms: true
  auto_spawn:
    - id: valid
      agent: valid.yaml
    - id: missing
      agent: missing.yaml
"#;
        let spec = AgentBuilder::from_yaml(yaml).unwrap().spec.unwrap();
        let mut registry = LLMRegistry::new();
        registry.register("default", Arc::new(MockLLMProvider::new("test")));
        registry.set_default("default");

        let error = AgentBuilder::from_spec_with_base_dir(spec, &base_dir)
            .llm_registry(registry)
            .auto_configure_spawner()
            .await
            .err()
            .unwrap()
            .to_string();
        assert!(error.contains("missing"), "{error}");
        assert!(error.contains("missing.yaml"), "{error}");

        std::fs::remove_dir_all(base_dir).unwrap();
    }

    #[test]
    fn test_builder_chain() {
        let builder = AgentBuilder::new()
            .system_prompt("Test prompt")
            .max_iterations(5)
            .max_context_tokens(4096);

        assert_eq!(builder.system_prompt, Some("Test prompt".to_string()));
        assert_eq!(builder.max_iterations, Some(5));
        assert_eq!(builder.max_context_tokens, Some(4096));
    }

    #[test]
    fn test_builder_skills() {
        use ai_agents_skills::{SkillDefinition, SkillStep};

        let skill = SkillDefinition {
            id: "test".to_string(),
            description: "Test skill".to_string(),
            trigger: "When testing".to_string(),
            steps: vec![SkillStep::Prompt {
                prompt: "Hello".to_string(),
                llm: None,
            }],
            reasoning: None,
            reflection: None,
            disambiguation: None,
        };

        let builder = AgentBuilder::new().skill(skill.clone()).skills(vec![skill]);

        assert_eq!(builder.skills.len(), 2);
    }

    #[test]
    fn test_builder_from_yaml_with_states() {
        let yaml = r#"
name: StatefulAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
states:
  initial: greeting
  states:
    greeting:
      prompt: "Welcome!"
      transitions:
        - to: support
          when: "user needs help"
    support:
      prompt: "How can I help?"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();
        assert!(spec.has_states());
        assert!(spec.states.is_some());
        let states = spec.states.as_ref().unwrap();
        assert_eq!(states.initial, "greeting");
        assert_eq!(states.states.len(), 2);
    }

    #[test]
    fn test_builder_from_yaml_with_context() {
        let yaml = r#"
name: ContextAgent
system_prompt: "Hello, {{ context.user.name }}!"
llm:
  provider: openai
  model: gpt-4
context:
  user:
    type: runtime
    required: true
  time:
    type: builtin
    source: datetime
    refresh: per_turn
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();
        assert!(spec.has_context());
        assert_eq!(spec.context.len(), 2);
        assert!(spec.context.contains_key("user"));
        assert!(spec.context.contains_key("time"));
    }

    #[test]
    fn test_builder_from_yaml_with_full_v04_features() {
        let yaml = r#"
name: FullFeaturedAgent
version: "0.4.0"
system_prompt: |
  You are a helpful assistant.
  User: {{ context.user.name }}
  Language: {{ context.user.language }}
llm:
  provider: openai
  model: gpt-4
context:
  user:
    type: runtime
    required: true
    default:
      name: "Guest"
      language: "en"
  time:
    type: builtin
    source: datetime
    refresh: per_turn
states:
  initial: greeting
  states:
    greeting:
      prompt: "Welcome to our service!"
      prompt_mode: append
      transitions:
        - to: support
          when: "user needs help"
          auto: true
          priority: 10
    support:
      prompt: "I'm here to help you."
      max_turns: 5
      timeout_to: escalation
      transitions:
        - to: closing
          when: "issue resolved"
          auto: true
    escalation:
      prompt: "Let me connect you with a human agent."
    closing:
      prompt: "Thank you for using our service!"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();

        // Check context
        assert!(spec.has_context());
        assert_eq!(spec.context.len(), 2);

        // Check states
        assert!(spec.has_states());
        let states = spec.states.as_ref().unwrap();
        assert_eq!(states.initial, "greeting");
        assert_eq!(states.states.len(), 4);

        // Check greeting state details
        let greeting = states.states.get("greeting").unwrap();
        assert!(greeting.prompt.is_some());
        assert_eq!(greeting.transitions.len(), 1);
        assert_eq!(greeting.transitions[0].to, "support");
        assert!(greeting.transitions[0].auto);

        // Check support state has timeout
        let support = states.states.get("support").unwrap();
        assert_eq!(support.max_turns, Some(5));
        assert_eq!(support.timeout_to, Some("escalation".to_string()));
    }

    #[test]
    fn test_builder_from_yaml_with_hitl() {
        let yaml = r#"
name: HITLAgent
system_prompt: "You are a secure assistant."
llm:
  provider: openai
  model: gpt-4
hitl:
  default_timeout_seconds: 600
  on_timeout: reject
  tools:
    send_payment:
      require_approval: true
      approval_context:
        - amount
        - recipient
      approval_message: "Approve payment?"
    delete_record:
      require_approval: true
  conditions:
    - name: high_value
      when: "amount > 1000"
      require_approval: true
      approval_message: "High value transaction"
  states:
    escalation:
      on_enter: require_approval
      approval_message: "Escalate to human?"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();

        assert!(spec.has_hitl());
        let hitl = spec.hitl.as_ref().unwrap();
        assert_eq!(hitl.default_timeout_seconds, 600);
        assert_eq!(hitl.tools.len(), 2);
        assert!(hitl.tools.get("send_payment").unwrap().require_approval);
        assert_eq!(hitl.conditions.len(), 1);
        assert_eq!(hitl.conditions[0].name, "high_value");
        assert_eq!(hitl.states.len(), 1);
    }

    #[test]
    fn test_builder_from_yaml_with_compacting_memory() {
        let yaml = r#"
name: CompactingAgent
system_prompt: "You are a helpful assistant."
memory:
  type: compacting
  max_messages: 100
  max_recent_messages: 20
  compress_threshold: 15
  summarize_batch_size: 5
  summarizer_llm: router
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();

        assert!(spec.memory.is_compacting());
        assert_eq!(spec.memory.max_recent_messages, Some(20));
        assert_eq!(spec.memory.compress_threshold, Some(15));
        assert_eq!(spec.memory.summarize_batch_size, Some(5));
        assert_eq!(spec.memory.summarizer_llm, Some("router".to_string()));

        let compacting_config = spec.memory.to_compacting_config();
        assert_eq!(compacting_config.max_recent_messages, 20);
        assert_eq!(compacting_config.compress_threshold, 15);
        assert_eq!(compacting_config.summarize_batch_size, 5);
    }

    #[test]
    fn test_builder_from_yaml_with_token_budget() {
        let yaml = r#"
name: BudgetAgent
system_prompt: "You are a helpful assistant."
memory:
  type: compacting
  max_messages: 100
  token_budget:
    total: 8192
    allocation:
      summary: 2048
      recent_messages: 4096
      facts: 1024
    overflow_strategy: summarize_more
    warn_at_percent: 75
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        assert!(builder.spec.is_some());
        let spec = builder.spec.as_ref().unwrap();

        assert!(spec.memory.token_budget.is_some());
        let budget = spec.memory.token_budget.as_ref().unwrap();
        assert_eq!(budget.total, 8192);
        assert_eq!(budget.allocation.summary, 2048);
        assert_eq!(budget.allocation.recent_messages, 4096);
        assert_eq!(budget.allocation.facts, 1024);
        assert_eq!(budget.warn_at_percent, 75);
    }

    #[test]
    fn test_builder_from_yaml_with_overflow_strategies() {
        use ai_agents_memory::OverflowStrategy;

        let yaml = r#"
name: TruncateAgent
system_prompt: "You are helpful."
memory:
  type: compacting
  token_budget:
    total: 4096
    overflow_strategy: truncate_oldest
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let budget = builder
            .spec
            .as_ref()
            .unwrap()
            .memory
            .token_budget
            .as_ref()
            .unwrap();
        assert_eq!(budget.overflow_strategy, OverflowStrategy::TruncateOldest);

        let yaml = r#"
name: ErrorAgent
system_prompt: "You are helpful."
memory:
  type: compacting
  token_budget:
    total: 4096
    overflow_strategy: error
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let budget = builder
            .spec
            .as_ref()
            .unwrap()
            .memory
            .token_budget
            .as_ref()
            .unwrap();
        assert_eq!(budget.overflow_strategy, OverflowStrategy::Error);
    }

    #[test]
    fn test_builder_from_yaml_with_storage_file() {
        let yaml = r#"
name: PersistentAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
storage:
  type: file
  path: "./data/sessions"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let spec = builder.spec.as_ref().unwrap();
        assert!(spec.has_storage());
        assert!(spec.storage.is_file());
        assert_eq!(spec.storage.get_path(), Some("./data/sessions"));
    }

    #[test]
    fn test_builder_from_yaml_with_storage_sqlite() {
        let yaml = r#"
name: PersistentAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
storage:
  type: sqlite
  path: "./data/sessions.db"
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let spec = builder.spec.as_ref().unwrap();
        assert!(spec.has_storage());
        assert!(spec.storage.is_sqlite());
    }

    #[test]
    fn test_builder_from_yaml_with_storage_redis() {
        let yaml = r#"
name: PersistentAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
storage:
  type: redis
  url: "redis://localhost:6379"
  prefix: "myagent:"
  ttl_seconds: 86400
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let spec = builder.spec.as_ref().unwrap();
        assert!(spec.has_storage());
        assert!(spec.storage.is_redis());
        assert_eq!(spec.storage.get_prefix(), "myagent:");
        assert_eq!(spec.storage.get_ttl(), Some(86400));
    }

    #[test]
    fn test_builder_no_storage_by_default() {
        let yaml = r#"
name: SimpleAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let spec = builder.spec.as_ref().unwrap();
        assert!(!spec.has_storage());
    }

    #[test]
    fn test_build_fails_on_missing_declared_tool() {
        use ai_agents_llm::mock::MockLLMProvider;

        let yaml = r#"
name: ToolAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
tools:
  - name: lookup_order
  - name: calculator
"#;
        let llm = Arc::new(MockLLMProvider::new("test"));
        let result = AgentBuilder::from_yaml(yaml)
            .unwrap()
            .llm(llm)
            .auto_configure_features()
            .unwrap()
            // NOT registering lookup_order — should fail
            .build();

        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("lookup_order"),
            "error should name the missing tool: {}",
            err
        );
        // calculator is built-in, so it should NOT appear in the error
        assert!(
            !err.contains("calculator"),
            "calculator is registered, should not be missing: {}",
            err
        );
    }

    #[test]
    fn test_build_succeeds_when_declared_tool_is_registered() {
        use ai_agents_core::Tool;
        use ai_agents_llm::mock::MockLLMProvider;

        struct FakeTool;
        #[async_trait::async_trait]
        impl Tool for FakeTool {
            fn id(&self) -> &str {
                "lookup_order"
            }
            fn name(&self) -> &str {
                "Order Lookup"
            }
            fn description(&self) -> &str {
                "Look up an order"
            }
            fn input_schema(&self) -> serde_json::Value {
                serde_json::json!({})
            }
            async fn execute(
                &self,
                _args: serde_json::Value,
                _ctx: ai_agents_core::ToolExecutionContext,
            ) -> ai_agents_core::ToolResult {
                ai_agents_core::ToolResult::ok("ok")
            }
        }

        let yaml = r#"
name: ToolAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
tools:
  - name: lookup_order
  - name: calculator
"#;
        let llm = Arc::new(MockLLMProvider::new("test"));
        let result = AgentBuilder::from_yaml(yaml)
            .unwrap()
            .llm(llm)
            .auto_configure_features()
            .unwrap()
            .tool(Arc::new(FakeTool))
            .build();

        assert!(
            result.is_ok(),
            "build should succeed when all declared tools are registered: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_spawner_config_deserializes_shared_storage() {
        let yaml = r#"
name: TestAgent
system_prompt: "Test"
llm:
  provider: openai
  model: gpt-4
spawner:
  shared_llms: true
  shared_storage:
    type: sqlite
    path: ./data/test.db
  max_agents: 5
"#;
        let builder = AgentBuilder::from_yaml(yaml).unwrap();
        let spec = builder.spec.as_ref().unwrap();
        let sc = spec.spawner.as_ref().unwrap();
        assert!(sc.shared_storage.is_some());
        assert!(sc.shared_storage.as_ref().unwrap().is_sqlite());
    }

    #[test]
    fn test_build_succeeds_with_no_tools_declared() {
        use ai_agents_llm::mock::MockLLMProvider;

        let yaml = r#"
name: SimpleAgent
system_prompt: "You are helpful."
llm:
  provider: openai
  model: gpt-4
"#;
        let llm = Arc::new(MockLLMProvider::new("test"));
        let result = AgentBuilder::from_yaml(yaml)
            .unwrap()
            .llm(llm)
            .auto_configure_features()
            .unwrap()
            .build();

        assert!(
            result.is_ok(),
            "no tools: section means no validation needed: {:?}",
            result.err()
        );
    }
}