ai-agents-runtime 1.0.5

//! Spawner configuration types for YAML deserialization.

use serde::{Deserialize, Deserializer, Serialize};
use std::collections::HashMap;

use super::StorageConfig;

/// An agent to create at startup and register in the AgentRegistry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoSpawnEntry {
    /// Registry ID for this agent.
    pub id: String,
    /// Path to the agent YAML file (resolved relative to parent YAML directory).
    pub agent: String,
}

/// Tool selection: all tools or a specific subset.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum SpawnerToolGrantConfig {
    /// `management_tools: true` or `orchestration_tools: true` grants all tools in that group.
    All(bool),
    /// `management_tools: [spawn_agent, list_agents]` grants listed tools.
    Selected(Vec<String>),
}

impl Default for SpawnerToolGrantConfig {
    fn default() -> Self {
        Self::All(false)
    }
}

impl SpawnerToolGrantConfig {
    /// Returns true if any tools in this group are enabled.
    pub fn is_enabled(&self) -> bool {
        match self {
            Self::All(v) => *v,
            Self::Selected(v) => !v.is_empty(),
        }
    }

    /// Returns true if the given tool name is included.
    pub fn includes(&self, tool_name: &str) -> bool {
        match self {
            Self::All(true) => true,
            Self::All(false) => false,
            Self::Selected(v) => v.iter().any(|t| t == tool_name),
        }
    }

    /// Returns the selected tool IDs, or all IDs when this config is true.
    pub fn granted_tool_ids(&self, all_ids: &[&str]) -> Vec<String> {
        match self {
            Self::All(true) => all_ids.iter().map(|id| (*id).to_string()).collect(),
            Self::All(false) => Vec::new(),
            Self::Selected(ids) => ids.clone(),
        }
    }

    /// Returns the orchestration tool IDs granted by this config.
    pub fn granted_orchestration_tool_ids(&self) -> Vec<String> {
        const ALL: [&str; 5] = [
            "route_to_agent",
            "pipeline_process",
            "concurrent_ask",
            "group_discussion",
            "handoff_conversation",
        ];
        self.granted_tool_ids(&ALL)
    }

    /// Returns the management tool IDs granted by this config.
    pub fn granted_management_tool_ids(&self) -> Vec<String> {
        const ALL: [&str; 4] = [
            "spawn_agent",
            "send_agent_message",
            "list_agents",
            "remove_agent",
        ];
        self.granted_tool_ids(&ALL)
    }
}

pub type OrchestrationToolsConfig = SpawnerToolGrantConfig;
pub type ManagementToolsConfig = SpawnerToolGrantConfig;

/// Configuration for dynamic agent spawning declared in the `spawner:` YAML section.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SpawnerConfig {
    /// When true, spawned agents reuse the parent agent's LLM connections.
    #[serde(default)]
    pub shared_llms: bool,

    /// Shared storage backend for all spawned agents.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub shared_storage: Option<StorageConfig>,

    /// Context values injected into every spawned agent.
    #[serde(default)]
    pub shared_context: HashMap<String, serde_json::Value>,

    /// Maximum number of agents that can be spawned.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_agents: Option<usize>,

    /// Auto-naming prefix for spawned agents (e.g. "npc_" -> "npc_001").
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub name_prefix: Option<String>,

    /// Named YAML templates -- inline strings or file path references.
    #[serde(default)]
    pub templates: HashMap<String, TemplateSource>,

    /// Tool names that spawned agents are allowed to use.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub allowed_tools: Option<Vec<String>>,

    /// Agents to create at startup and register in the AgentRegistry.
    #[serde(default)]
    pub auto_spawn: Vec<AutoSpawnEntry>,

    /// Grant dynamic agent management tools (spawn_agent, send_agent_message, list_agents, remove_agent).
    #[serde(default)]
    pub management_tools: ManagementToolsConfig,

    /// Register orchestration tools (route_to_agent, group_discussion, etc.).
    #[serde(default)]
    pub orchestration_tools: OrchestrationToolsConfig,
}

/// A spawner template source: either an inline YAML string or a file path reference.
///
/// Untagged:
/// Serde tries `File` first (object with `path` key), falls back to `Inline` (plain string).
/// File paths are resolved against the parent YAML directory at config time.
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum TemplateSource {
    /// File-based template: `{ path: "./templates/npc.yaml" }`.
    File { path: String },
    /// Inline YAML template string (backward compatible).
    Inline(String),
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct FileTemplateSource {
    path: String,
}

impl<'de> Deserialize<'de> for TemplateSource {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = serde_yaml::Value::deserialize(deserializer)?;
        match value {
            serde_yaml::Value::String(inline) => Ok(Self::Inline(inline)),
            serde_yaml::Value::Mapping(_) => serde_yaml::from_value::<FileTemplateSource>(value)
                .map(|source| Self::File { path: source.path })
                .map_err(serde::de::Error::custom),
            _ => Err(serde::de::Error::custom(
                "template source must be an inline YAML string or an object with a path field",
            )),
        }
    }
}

impl TemplateSource {
    /// Returns true if this is a file path reference.
    pub fn is_file(&self) -> bool {
        matches!(self, Self::File { .. })
    }

    /// Returns true if this is an inline template string.
    pub fn is_inline(&self) -> bool {
        matches!(self, Self::Inline(_))
    }
}

impl SpawnerConfig {
    /// Returns true if any spawner configuration is present.
    pub fn is_configured(&self) -> bool {
        self.shared_llms
            || self.shared_storage.is_some()
            || !self.shared_context.is_empty()
            || self.max_agents.is_some()
            || self.name_prefix.is_some()
            || !self.templates.is_empty()
            || !self.auto_spawn.is_empty()
            || self.management_tools.is_enabled()
            || self.orchestration_tools.is_enabled()
    }
}

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

    #[test]
    fn test_default_is_not_configured() {
        let config = SpawnerConfig::default();
        assert!(!config.is_configured());
    }

    #[test]
    fn test_deserialize_inline_template() {
        let yaml = r#"
shared_llms: true
max_agents: 50
name_prefix: "npc_"
shared_context:
  world_name: "Medieval Fantasy"
  current_era: "Age of Dragons"
templates:
  npc_base: |
    name: "{{ name }}"
    system_prompt: "You are {{ name }}."
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.shared_llms);
        assert_eq!(config.max_agents, Some(50));
        assert_eq!(config.name_prefix.as_deref(), Some("npc_"));
        assert_eq!(config.shared_context.len(), 2);
        assert!(config.templates.contains_key("npc_base"));
        assert!(config.templates.get("npc_base").unwrap().is_inline());
        assert!(config.is_configured());
    }

    #[test]
    fn test_deserialize_auto_spawn() {
        let yaml = r#"
shared_llms: true
auto_spawn:
  - id: billing
    agent: agents/billing_agent.yaml
  - id: technical
    agent: agents/technical_agent.yaml
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(config.auto_spawn.len(), 2);
        assert_eq!(config.auto_spawn[0].id, "billing");
        assert_eq!(config.auto_spawn[0].agent, "agents/billing_agent.yaml");
    }

    #[test]
    fn test_deserialize_management_tools_all() {
        let yaml = r#"
management_tools: true
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.management_tools.is_enabled());
        assert!(config.management_tools.includes("spawn_agent"));
        assert_eq!(
            config.management_tools.granted_management_tool_ids(),
            vec![
                "spawn_agent".to_string(),
                "send_agent_message".to_string(),
                "list_agents".to_string(),
                "remove_agent".to_string(),
            ]
        );
    }

    #[test]
    fn test_deserialize_management_tools_selected() {
        let yaml = r#"
management_tools:
  - spawn_agent
  - send_agent_message
  - list_agents
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.management_tools.is_enabled());
        assert!(config.management_tools.includes("spawn_agent"));
        assert!(config.management_tools.includes("send_agent_message"));
        assert!(config.management_tools.includes("list_agents"));
        assert!(!config.management_tools.includes("remove_agent"));
        assert_eq!(
            config.management_tools.granted_management_tool_ids(),
            vec![
                "spawn_agent".to_string(),
                "send_agent_message".to_string(),
                "list_agents".to_string(),
            ]
        );
    }

    #[test]
    fn test_deserialize_orchestration_tools_all() {
        let yaml = r#"
orchestration_tools: true
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.orchestration_tools.is_enabled());
        assert!(config.orchestration_tools.includes("route_to_agent"));
    }

    #[test]
    fn test_deserialize_orchestration_tools_selected() {
        let yaml = r#"
orchestration_tools:
  - route_to_agent
  - group_discussion
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.orchestration_tools.is_enabled());
        assert!(config.orchestration_tools.includes("route_to_agent"));
        assert!(config.orchestration_tools.includes("group_discussion"));
        assert!(!config.orchestration_tools.includes("concurrent_ask"));
    }

    #[test]
    fn test_auto_spawn_makes_configured() {
        let config = SpawnerConfig {
            auto_spawn: vec![AutoSpawnEntry {
                id: "test".to_string(),
                agent: "test.yaml".to_string(),
            }],
            ..SpawnerConfig::default()
        };
        assert!(config.is_configured());
    }

    #[test]
    fn test_deserialize_file_template() {
        let yaml = r#"
templates:
  npc_base:
    path: ./templates/npc_base.yaml
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        match config.templates.get("npc_base") {
            Some(TemplateSource::File { path }) => {
                assert_eq!(path, "./templates/npc_base.yaml");
            }
            other => panic!("expected File variant, got {:?}", other),
        }
    }

    #[test]
    fn test_deserialize_mixed_templates() {
        let yaml = r#"
templates:
  inline_one: |
    name: "{{ name }}"
  file_one:
    path: ./templates/npc.yaml
  inline_two: "name: test"
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        assert!(config.templates.get("inline_one").unwrap().is_inline());
        assert!(config.templates.get("file_one").unwrap().is_file());
        assert!(config.templates.get("inline_two").unwrap().is_inline());
    }

    #[test]
    fn test_deserialize_absolute_path_template() {
        let yaml = r#"
templates:
  shared_guard:
    path: /opt/game/shared_templates/guard.yaml
"#;
        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
        match config.templates.get("shared_guard") {
            Some(TemplateSource::File { path }) => {
                assert_eq!(path, "/opt/game/shared_templates/guard.yaml");
            }
            other => panic!("expected File variant, got {:?}", other),
        }
    }

    #[test]
    fn test_roundtrip_serde() {
        let config = SpawnerConfig {
            shared_llms: true,
            shared_storage: None,
            shared_context: HashMap::new(),
            max_agents: Some(100),
            name_prefix: Some("test_".to_string()),
            templates: HashMap::new(),
            allowed_tools: Some(vec!["echo".to_string()]),
            auto_spawn: Vec::new(),
            management_tools: ManagementToolsConfig::default(),
            orchestration_tools: OrchestrationToolsConfig::default(),
        };
        let yaml = serde_yaml::to_string(&config).unwrap();
        let parsed: SpawnerConfig = serde_yaml::from_str(&yaml).unwrap();
        assert_eq!(parsed.max_agents, Some(100));
        assert_eq!(parsed.name_prefix.as_deref(), Some("test_"));
    }

    #[test]
    fn test_template_source_is_file() {
        let file = TemplateSource::File {
            path: "./test.yaml".to_string(),
        };
        let inline = TemplateSource::Inline("content".to_string());
        assert!(file.is_file());
        assert!(!file.is_inline());
        assert!(inline.is_inline());
        assert!(!inline.is_file());
    }
}