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, LLMProvider, Result, Tool};
use ai_agents_hitl::{ApprovalHandler, HITLEngine, RejectAllHandler};
use ai_agents_hooks::AgentHooks;
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_process::ProcessProcessor;
use ai_agents_reasoning::{ReasoningConfig, ReflectionConfig};
use ai_agents_recovery::{MessageFilter, RecoveryManager};
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;
use crate::spec::{AgentSpec, StorageConfig};

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>>,
}

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,
        }
    }

    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,
        }
    }

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

    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 mut builder = Self::from_yaml(&content)?;
        if let Some(parent) = path.parent() {
            builder.yaml_dir = Some(parent.to_path_buf());
        }
        Ok(builder)
    }

    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 = serde_yaml::from_str(&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 provider = UnifiedLLMProvider::from_spec_config(
                    provider_type,
                    &config.model,
                    api_key,
                    base_url,
                    core_config,
                )
                .map_err(|e| AgentError::LLM(e.to_string()))?;

                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);
        } 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 provider = UnifiedLLMProvider::from_spec_config(
                provider_type,
                &config.model,
                api_key,
                base_url,
                core_config,
            )
            .map_err(|e| AgentError::LLM(e.to_string()))?;

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

        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::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
    }

    pub fn llm_registry(mut self, registry: LLMRegistry) -> Self {
        self.llm_registry = Some(registry);
        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() {
                if 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
    }

    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),
        };

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

        let mut spawner = AgentSpawner::new();

        if spawner_config.shared_llms {
            if let Some(ref reg) = self.llm_registry {
                spawner = 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().map_or(false, |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");
        }

        // Auto-spawn agents from YAML files and register them.
        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");

            match AgentBuilder::from_yaml_file(&yaml_path) {
                Ok(mut sub_builder) => {
                    if spawner_config.shared_llms {
                        if let Some(ref reg) = self.llm_registry {
                            sub_builder = sub_builder.llm_registry(reg.clone());
                        }
                    }
                    match sub_builder
                        .auto_configure_llms()
                        .and_then(|b| b.auto_configure_features())
                    {
                        Ok(configured) => match configured.build() {
                            Ok(agent) => {
                                let spec_yaml =
                                    std::fs::read_to_string(&yaml_path).unwrap_or_default();
                                let spec: crate::spec::AgentSpec =
                                    serde_yaml::from_str(&spec_yaml).unwrap_or_default();

                                let spawned = crate::spawner::spawner::SpawnedAgent {
                                    id: entry.id.clone(),
                                    agent: Arc::new(agent),
                                    spec,
                                    spawned_at: chrono::Utc::now(),
                                };

                                if let Err(e) = registry.register(spawned).await {
                                    tracing::warn!(
                                        id = %entry.id,
                                        error = %e,
                                        "Failed to register auto-spawned agent"
                                    );
                                } else {
                                    tracing::info!(id = %entry.id, "Auto-spawned agent registered");
                                }
                            }
                            Err(e) => {
                                tracing::warn!(
                                    id = %entry.id,
                                    error = %e,
                                    "Failed to build auto-spawn agent"
                                );
                            }
                        },
                        Err(e) => {
                            tracing::warn!(
                                id = %entry.id,
                                error = %e,
                                "Failed to configure auto-spawn agent"
                            );
                        }
                    }
                }
                Err(e) => {
                    tracing::warn!(
                        id = %entry.id,
                        path = %yaml_path.display(),
                        error = %e,
                        "Failed to load auto-spawn YAML"
                    );
                }
            }
        }

        // Validate that all orchestration state references have matching agents.
        if let Some(ref spec) = self.spec {
            if 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 base_prompt = self
            .system_prompt
            .ok_or_else(|| AgentError::Config("System prompt is required".into()))?;

        let mut tools = self.tools.unwrap_or_else(ToolRegistry::new);

        // 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 {
            if !spec.skills.is_empty() {
                let mut loader = self.skill_loader.take().unwrap_or_else(SkillLoader::new);
                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_else(LLMRegistry::new);

        if let Some(llm) = self.llm {
            if !llm_registry.has("default") {
                llm_registry.register("default", llm.clone());
            }
        }

        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(),
            ));
        }

        // 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 {
            if pm.should_register_evolve_tool() {
                let evolve_tool = ai_agents_persona::PersonaEvolveTool::new(pm.clone());
                let _ = tools.register(Arc::new(evolve_tool));
            }
        }

        let tools_arc = Arc::new(tools);
        let llm_registry_arc = Arc::new(llm_registry);

        // Extract declared tool IDs from spec (agent-level tool scope)
        // - None  = `tools:` not specified in YAML -> all registered tools available : WILL CHANGE as no tools (Future task)
        // - Some([]) = `tools: []` in YAML -> explicitly no tools
        // - Some([...]) = specific tools listed
        let declared_tool_ids: Option<Vec<String>> = self.spec.as_ref().and_then(|s| {
            s.tools.as_ref().map(|tools| {
                let mut ids: Vec<String> = tools.iter().map(|t| t.name().to_string()).collect();

                // Include view names from MCP entries so they pass availability checks
                for entry in tools {
                    if let Some(mcp_config) = entry.to_mcp_config() {
                        for view_name in mcp_config.views.keys() {
                            ids.push(view_name.clone());
                        }
                    }
                }

                // Auto-registered tools must bypass scoping when the user has an explicit list.
                // persona_evolve is auto-registered when allow_llm_evolve is true.
                if persona_manager
                    .as_ref()
                    .is_some_and(|pm| pm.should_register_evolve_tool())
                {
                    ids.push("persona_evolve".to_string());
                }

                ids
            })
        });

        // Validate: every non-MCP tool declared in the YAML spec 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).
        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 declared in YAML but not registered: [{}]. \
                     Register them via .tool(Arc::new(...)) before .build(), \
                     or remove them from the YAML tools: list.",
                    missing.join(", ")
                )));
            }
        }

        let mut agent = RuntimeAgent::new(
            info,
            llm_registry_arc.clone(),
            memory,
            tools_arc,
            self.skills,
            system_prompt,
            max_iterations,
        )
        .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::new(spec.tool_security.clone()));
        }

        if let Some(processor) = self.process_processor {
            agent = agent.with_process_processor(processor);
        } else if let Some(ref spec) = self.spec {
            if 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 {
            if 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 {
            if !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);

            // 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(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 {
            if 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 {
            if 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);
        }

        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_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_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_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) -> 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()
        );
    }
}