Skip to main content

ai_agents_runtime/spec/
spawner.rs

1//! Spawner configuration types for YAML deserialization.
2
3use serde::{Deserialize, Deserializer, Serialize};
4use std::collections::HashMap;
5
6use super::StorageConfig;
7
8/// An agent to create at startup and register in the AgentRegistry.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct AutoSpawnEntry {
11    /// Registry ID for this agent.
12    pub id: String,
13    /// Path to the agent YAML file (resolved relative to parent YAML directory).
14    pub agent: String,
15}
16
17/// Tool selection: all tools or a specific subset.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19#[serde(untagged)]
20pub enum SpawnerToolGrantConfig {
21    /// `management_tools: true` or `orchestration_tools: true` grants all tools in that group.
22    All(bool),
23    /// `management_tools: [spawn_agent, list_agents]` grants listed tools.
24    Selected(Vec<String>),
25}
26
27impl Default for SpawnerToolGrantConfig {
28    fn default() -> Self {
29        Self::All(false)
30    }
31}
32
33impl SpawnerToolGrantConfig {
34    /// Returns true if any tools in this group are enabled.
35    pub fn is_enabled(&self) -> bool {
36        match self {
37            Self::All(v) => *v,
38            Self::Selected(v) => !v.is_empty(),
39        }
40    }
41
42    /// Returns true if the given tool name is included.
43    pub fn includes(&self, tool_name: &str) -> bool {
44        match self {
45            Self::All(true) => true,
46            Self::All(false) => false,
47            Self::Selected(v) => v.iter().any(|t| t == tool_name),
48        }
49    }
50
51    /// Returns the selected tool IDs, or all IDs when this config is true.
52    pub fn granted_tool_ids(&self, all_ids: &[&str]) -> Vec<String> {
53        match self {
54            Self::All(true) => all_ids.iter().map(|id| (*id).to_string()).collect(),
55            Self::All(false) => Vec::new(),
56            Self::Selected(ids) => ids.clone(),
57        }
58    }
59
60    /// Returns the orchestration tool IDs granted by this config.
61    pub fn granted_orchestration_tool_ids(&self) -> Vec<String> {
62        const ALL: [&str; 5] = [
63            "route_to_agent",
64            "pipeline_process",
65            "concurrent_ask",
66            "group_discussion",
67            "handoff_conversation",
68        ];
69        self.granted_tool_ids(&ALL)
70    }
71
72    /// Returns the management tool IDs granted by this config.
73    pub fn granted_management_tool_ids(&self) -> Vec<String> {
74        const ALL: [&str; 4] = [
75            "spawn_agent",
76            "send_agent_message",
77            "list_agents",
78            "remove_agent",
79        ];
80        self.granted_tool_ids(&ALL)
81    }
82}
83
84pub type OrchestrationToolsConfig = SpawnerToolGrantConfig;
85pub type ManagementToolsConfig = SpawnerToolGrantConfig;
86
87/// Configuration for dynamic agent spawning declared in the `spawner:` YAML section.
88#[derive(Debug, Clone, Serialize, Deserialize, Default)]
89pub struct SpawnerConfig {
90    /// When true, spawned agents reuse the parent agent's LLM connections.
91    #[serde(default)]
92    pub shared_llms: bool,
93
94    /// Shared storage backend for all spawned agents.
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub shared_storage: Option<StorageConfig>,
97
98    /// Context values injected into every spawned agent.
99    #[serde(default)]
100    pub shared_context: HashMap<String, serde_json::Value>,
101
102    /// Maximum number of agents that can be spawned.
103    #[serde(default, skip_serializing_if = "Option::is_none")]
104    pub max_agents: Option<usize>,
105
106    /// Auto-naming prefix for spawned agents (e.g. "npc_" -> "npc_001").
107    #[serde(default, skip_serializing_if = "Option::is_none")]
108    pub name_prefix: Option<String>,
109
110    /// Named YAML templates -- inline strings or file path references.
111    #[serde(default)]
112    pub templates: HashMap<String, TemplateSource>,
113
114    /// Tool names that spawned agents are allowed to use.
115    #[serde(default, skip_serializing_if = "Option::is_none")]
116    pub allowed_tools: Option<Vec<String>>,
117
118    /// Agents to create at startup and register in the AgentRegistry.
119    #[serde(default)]
120    pub auto_spawn: Vec<AutoSpawnEntry>,
121
122    /// Grant dynamic agent management tools (spawn_agent, send_agent_message, list_agents, remove_agent).
123    #[serde(default)]
124    pub management_tools: ManagementToolsConfig,
125
126    /// Register orchestration tools (route_to_agent, group_discussion, etc.).
127    #[serde(default)]
128    pub orchestration_tools: OrchestrationToolsConfig,
129}
130
131/// A spawner template source: either an inline YAML string or a file path reference.
132///
133/// Untagged:
134/// Serde tries `File` first (object with `path` key), falls back to `Inline` (plain string).
135/// File paths are resolved against the parent YAML directory at config time.
136#[derive(Debug, Clone, Serialize)]
137#[serde(untagged)]
138pub enum TemplateSource {
139    /// File-based template: `{ path: "./templates/npc.yaml" }`.
140    File { path: String },
141    /// Inline YAML template string (backward compatible).
142    Inline(String),
143}
144
145#[derive(Deserialize)]
146#[serde(deny_unknown_fields)]
147struct FileTemplateSource {
148    path: String,
149}
150
151impl<'de> Deserialize<'de> for TemplateSource {
152    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
153    where
154        D: Deserializer<'de>,
155    {
156        let value = serde_yaml::Value::deserialize(deserializer)?;
157        match value {
158            serde_yaml::Value::String(inline) => Ok(Self::Inline(inline)),
159            serde_yaml::Value::Mapping(_) => serde_yaml::from_value::<FileTemplateSource>(value)
160                .map(|source| Self::File { path: source.path })
161                .map_err(serde::de::Error::custom),
162            _ => Err(serde::de::Error::custom(
163                "template source must be an inline YAML string or an object with a path field",
164            )),
165        }
166    }
167}
168
169impl TemplateSource {
170    /// Returns true if this is a file path reference.
171    pub fn is_file(&self) -> bool {
172        matches!(self, Self::File { .. })
173    }
174
175    /// Returns true if this is an inline template string.
176    pub fn is_inline(&self) -> bool {
177        matches!(self, Self::Inline(_))
178    }
179}
180
181impl SpawnerConfig {
182    /// Returns true if any spawner configuration is present.
183    pub fn is_configured(&self) -> bool {
184        self.shared_llms
185            || self.shared_storage.is_some()
186            || !self.shared_context.is_empty()
187            || self.max_agents.is_some()
188            || self.name_prefix.is_some()
189            || !self.templates.is_empty()
190            || !self.auto_spawn.is_empty()
191            || self.management_tools.is_enabled()
192            || self.orchestration_tools.is_enabled()
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199
200    #[test]
201    fn test_default_is_not_configured() {
202        let config = SpawnerConfig::default();
203        assert!(!config.is_configured());
204    }
205
206    #[test]
207    fn test_deserialize_inline_template() {
208        let yaml = r#"
209shared_llms: true
210max_agents: 50
211name_prefix: "npc_"
212shared_context:
213  world_name: "Medieval Fantasy"
214  current_era: "Age of Dragons"
215templates:
216  npc_base: |
217    name: "{{ name }}"
218    system_prompt: "You are {{ name }}."
219"#;
220        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
221        assert!(config.shared_llms);
222        assert_eq!(config.max_agents, Some(50));
223        assert_eq!(config.name_prefix.as_deref(), Some("npc_"));
224        assert_eq!(config.shared_context.len(), 2);
225        assert!(config.templates.contains_key("npc_base"));
226        assert!(config.templates.get("npc_base").unwrap().is_inline());
227        assert!(config.is_configured());
228    }
229
230    #[test]
231    fn test_deserialize_auto_spawn() {
232        let yaml = r#"
233shared_llms: true
234auto_spawn:
235  - id: billing
236    agent: agents/billing_agent.yaml
237  - id: technical
238    agent: agents/technical_agent.yaml
239"#;
240        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
241        assert_eq!(config.auto_spawn.len(), 2);
242        assert_eq!(config.auto_spawn[0].id, "billing");
243        assert_eq!(config.auto_spawn[0].agent, "agents/billing_agent.yaml");
244    }
245
246    #[test]
247    fn test_deserialize_management_tools_all() {
248        let yaml = r#"
249management_tools: true
250"#;
251        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
252        assert!(config.management_tools.is_enabled());
253        assert!(config.management_tools.includes("spawn_agent"));
254        assert_eq!(
255            config.management_tools.granted_management_tool_ids(),
256            vec![
257                "spawn_agent".to_string(),
258                "send_agent_message".to_string(),
259                "list_agents".to_string(),
260                "remove_agent".to_string(),
261            ]
262        );
263    }
264
265    #[test]
266    fn test_deserialize_management_tools_selected() {
267        let yaml = r#"
268management_tools:
269  - spawn_agent
270  - send_agent_message
271  - list_agents
272"#;
273        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
274        assert!(config.management_tools.is_enabled());
275        assert!(config.management_tools.includes("spawn_agent"));
276        assert!(config.management_tools.includes("send_agent_message"));
277        assert!(config.management_tools.includes("list_agents"));
278        assert!(!config.management_tools.includes("remove_agent"));
279        assert_eq!(
280            config.management_tools.granted_management_tool_ids(),
281            vec![
282                "spawn_agent".to_string(),
283                "send_agent_message".to_string(),
284                "list_agents".to_string(),
285            ]
286        );
287    }
288
289    #[test]
290    fn test_deserialize_orchestration_tools_all() {
291        let yaml = r#"
292orchestration_tools: true
293"#;
294        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
295        assert!(config.orchestration_tools.is_enabled());
296        assert!(config.orchestration_tools.includes("route_to_agent"));
297    }
298
299    #[test]
300    fn test_deserialize_orchestration_tools_selected() {
301        let yaml = r#"
302orchestration_tools:
303  - route_to_agent
304  - group_discussion
305"#;
306        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
307        assert!(config.orchestration_tools.is_enabled());
308        assert!(config.orchestration_tools.includes("route_to_agent"));
309        assert!(config.orchestration_tools.includes("group_discussion"));
310        assert!(!config.orchestration_tools.includes("concurrent_ask"));
311    }
312
313    #[test]
314    fn test_auto_spawn_makes_configured() {
315        let config = SpawnerConfig {
316            auto_spawn: vec![AutoSpawnEntry {
317                id: "test".to_string(),
318                agent: "test.yaml".to_string(),
319            }],
320            ..SpawnerConfig::default()
321        };
322        assert!(config.is_configured());
323    }
324
325    #[test]
326    fn test_deserialize_file_template() {
327        let yaml = r#"
328templates:
329  npc_base:
330    path: ./templates/npc_base.yaml
331"#;
332        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
333        match config.templates.get("npc_base") {
334            Some(TemplateSource::File { path }) => {
335                assert_eq!(path, "./templates/npc_base.yaml");
336            }
337            other => panic!("expected File variant, got {:?}", other),
338        }
339    }
340
341    #[test]
342    fn test_deserialize_mixed_templates() {
343        let yaml = r#"
344templates:
345  inline_one: |
346    name: "{{ name }}"
347  file_one:
348    path: ./templates/npc.yaml
349  inline_two: "name: test"
350"#;
351        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
352        assert!(config.templates.get("inline_one").unwrap().is_inline());
353        assert!(config.templates.get("file_one").unwrap().is_file());
354        assert!(config.templates.get("inline_two").unwrap().is_inline());
355    }
356
357    #[test]
358    fn test_deserialize_absolute_path_template() {
359        let yaml = r#"
360templates:
361  shared_guard:
362    path: /opt/game/shared_templates/guard.yaml
363"#;
364        let config: SpawnerConfig = serde_yaml::from_str(yaml).unwrap();
365        match config.templates.get("shared_guard") {
366            Some(TemplateSource::File { path }) => {
367                assert_eq!(path, "/opt/game/shared_templates/guard.yaml");
368            }
369            other => panic!("expected File variant, got {:?}", other),
370        }
371    }
372
373    #[test]
374    fn test_roundtrip_serde() {
375        let config = SpawnerConfig {
376            shared_llms: true,
377            shared_storage: None,
378            shared_context: HashMap::new(),
379            max_agents: Some(100),
380            name_prefix: Some("test_".to_string()),
381            templates: HashMap::new(),
382            allowed_tools: Some(vec!["echo".to_string()]),
383            auto_spawn: Vec::new(),
384            management_tools: ManagementToolsConfig::default(),
385            orchestration_tools: OrchestrationToolsConfig::default(),
386        };
387        let yaml = serde_yaml::to_string(&config).unwrap();
388        let parsed: SpawnerConfig = serde_yaml::from_str(&yaml).unwrap();
389        assert_eq!(parsed.max_agents, Some(100));
390        assert_eq!(parsed.name_prefix.as_deref(), Some("test_"));
391    }
392
393    #[test]
394    fn test_template_source_is_file() {
395        let file = TemplateSource::File {
396            path: "./test.yaml".to_string(),
397        };
398        let inline = TemplateSource::Inline("content".to_string());
399        assert!(file.is_file());
400        assert!(!file.is_inline());
401        assert!(inline.is_inline());
402        assert!(!inline.is_file());
403    }
404}