Skip to main content

ai_agents_runtime/spawner/
spawner.rs

1//! Core agent spawner for creating agents at runtime.
2
3use std::collections::{BTreeSet, HashMap};
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
7
8use chrono::Utc;
9use minijinja::Environment;
10use tracing::info;
11
12use crate::AgentBuilder;
13use crate::RuntimeAgent;
14use crate::runtime::ToolResourceLocks;
15use crate::spec::AgentSpec;
16use ai_agents_core::{AgentError, AgentStorage, Result};
17use ai_agents_llm::LLMRegistry;
18use ai_agents_observability::ObservabilityManager;
19use ai_agents_tools::create_builtin_registry;
20
21use super::storage::NamespacedStorage;
22
23/// A spawner template with its raw content and extracted metadata.
24#[derive(Debug, Clone)]
25pub struct ResolvedTemplate {
26    /// Raw Jinja2 template string for rendering.
27    pub content: String,
28    /// Template description extracted from the `description:` field.
29    pub description: Option<String>,
30    /// Variable name -> description map extracted from `metadata.template.variables`.
31    pub variables: Option<HashMap<String, String>>,
32}
33
34impl ResolvedTemplate {
35    /// Create a ResolvedTemplate from a plain content string with no metadata.
36    pub fn from_content(content: impl Into<String>) -> Self {
37        Self {
38            content: content.into(),
39            description: None,
40            variables: None,
41        }
42    }
43}
44
45pub(crate) struct CapacityReservation {
46    agent_count: Arc<AtomicU32>,
47    released: AtomicBool,
48}
49
50impl CapacityReservation {
51    fn new(agent_count: Arc<AtomicU32>) -> Self {
52        Self {
53            agent_count,
54            released: AtomicBool::new(false),
55        }
56    }
57
58    fn release(&self) {
59        // A slot may be released by registration failure, registry removal, or final drop. The flag keeps these competing lifecycle paths exactly once.
60        if !self.released.swap(true, Ordering::AcqRel) {
61            let _ = self
62                .agent_count
63                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
64                    count.checked_sub(1)
65                });
66        }
67    }
68}
69
70impl Drop for CapacityReservation {
71    fn drop(&mut self) {
72        self.release();
73    }
74}
75
76struct AdmittedChild {
77    id: String,
78    spec: AgentSpec,
79    base_dir: Option<PathBuf>,
80}
81
82/// Metadata for a spawned agent.
83pub struct SpawnedAgent {
84    /// Unique identifier (derived from spec name or auto-generated).
85    pub id: String,
86    /// The runtime agent, wrapped in Arc for shared ownership across registry callers.
87    pub agent: Arc<RuntimeAgent>,
88    /// Retained spec for introspection and serialization.
89    pub spec: AgentSpec,
90    /// Timestamp when the agent was created.
91    pub spawned_at: chrono::DateTime<Utc>,
92    capacity_reservation: Option<CapacityReservation>,
93}
94
95impl SpawnedAgent {
96    /// Construct a detached spawned-agent record for host-managed registration.
97    pub fn from_runtime(id: String, agent: RuntimeAgent, spec: AgentSpec) -> Self {
98        Self {
99            id,
100            agent: Arc::new(agent),
101            spec,
102            spawned_at: Utc::now(),
103            capacity_reservation: None,
104        }
105    }
106
107    fn tracked(
108        id: String,
109        agent: RuntimeAgent,
110        spec: AgentSpec,
111        capacity_reservation: CapacityReservation,
112    ) -> Self {
113        Self {
114            id,
115            agent: Arc::new(agent),
116            spec,
117            spawned_at: Utc::now(),
118            capacity_reservation: Some(capacity_reservation),
119        }
120    }
121
122    pub(super) fn release_capacity(&self) {
123        if let Some(reservation) = self.capacity_reservation.as_ref() {
124            reservation.release();
125        }
126    }
127
128    #[cfg(test)]
129    pub(crate) fn untracked(id: String, agent: RuntimeAgent, spec: AgentSpec) -> Self {
130        Self::from_runtime(id, agent, spec)
131    }
132}
133
134impl std::fmt::Debug for SpawnedAgent {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        f.debug_struct("SpawnedAgent")
137            .field("id", &self.id)
138            .field("spawned_at", &self.spawned_at)
139            .finish_non_exhaustive()
140    }
141}
142
143/// Factory for creating agents at runtime from YAML, specs, or templates.
144pub struct AgentSpawner {
145    /// Shared LLM regstry - spawned agents reuse these connections.
146    llm_registry: Option<LLMRegistry>,
147
148    /// Whether shared LLM providers are already observability-wrapped.
149    llm_registry_observed: bool,
150
151    /// Shared storage backend with per-gaent `NamespacedStorage` warpping.
152    storage: Option<Arc<dyn AgentStorage>>,
153
154    /// Context values injected into every spawned agent.
155    shared_context: HashMap<String, serde_json::Value>,
156
157    /// Hard limit on the number of agents this spawner may create.
158    max_agents: Option<usize>,
159
160    /// Auto-naming prefix (e.g. "npc_" produces "npc_001", "npc_002").
161    name_prefix: Option<String>,
162
163    /// Named YAML templates with content and extracted metadata.
164    templates: HashMap<String, ResolvedTemplate>,
165
166    /// Tool names that spawned agents are allowed to declare.
167    allowed_tools: Option<Vec<String>>,
168
169    /// Shared resource locks inherited by spawned child agents.
170    resource_locks: Option<ToolResourceLocks>,
171
172    /// Shared observability manager for spawned child agents.
173    observability_manager: Option<Arc<ObservabilityManager>>,
174
175    /// Monotonic counter for auto-naming.
176    counter: AtomicU32,
177
178    /// Running count of reserved or registered agents.
179    agent_count: Arc<AtomicU32>,
180}
181
182impl AgentSpawner {
183    pub fn new() -> Self {
184        Self {
185            llm_registry: None,
186            llm_registry_observed: false,
187            storage: None,
188            shared_context: HashMap::new(),
189            max_agents: None,
190            name_prefix: None,
191            templates: HashMap::new(),
192            allowed_tools: None,
193            resource_locks: None,
194            observability_manager: None,
195            counter: AtomicU32::new(1),
196            agent_count: Arc::new(AtomicU32::new(0)),
197        }
198    }
199
200    /// Share LLM connections across all spawned agents.
201    pub fn with_shared_llms(mut self, registry: LLMRegistry) -> Self {
202        self.llm_registry = Some(registry);
203        self.llm_registry_observed = false;
204        self
205    }
206
207    /// Share an LLM registry that is already wrapped by observed providers.
208    pub fn with_shared_observed_llms(mut self, registry: LLMRegistry) -> Self {
209        self.llm_registry = Some(registry);
210        self.llm_registry_observed = true;
211        self
212    }
213
214    /// Share a storage backend (e.g. one SQLite DB for all NPCs).
215    pub fn with_shared_storage(mut self, storage: Arc<dyn AgentStorage>) -> Self {
216        self.storage = Some(storage);
217        self
218    }
219
220    /// Inject a context value available to all spawned agents.
221    pub fn with_shared_context(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
222        self.shared_context.insert(key.into(), value);
223        self
224    }
225
226    /// Inject an entire map of shared context values.
227    pub fn with_shared_context_map(mut self, ctx: HashMap<String, serde_json::Value>) -> Self {
228        self.shared_context.extend(ctx);
229        self
230    }
231
232    /// Limit total spawned agents.
233    pub fn with_max_agents(mut self, max: usize) -> Self {
234        self.max_agents = Some(max);
235        self
236    }
237
238    /// Auto-name agents with prefix + zero-padded counter.
239    pub fn with_name_prefix(mut self, prefix: impl Into<String>) -> Result<Self> {
240        let prefix = prefix.into();
241        validate_name_prefix(&prefix)?;
242        self.name_prefix = Some(prefix);
243        Ok(self)
244    }
245
246    /// Register a named template from a plain YAML string (no metadata).
247    pub fn with_template(
248        mut self,
249        name: impl Into<String>,
250        yaml_template: impl Into<String>,
251    ) -> Self {
252        self.templates
253            .insert(name.into(), ResolvedTemplate::from_content(yaml_template));
254        self
255    }
256
257    /// Bulk-register resolved templates (with metadata already extracted).
258    pub fn with_templates(mut self, templates: HashMap<String, ResolvedTemplate>) -> Self {
259        self.templates.extend(templates);
260        self
261    }
262
263    /// Set the tool allowlist for spawned agents.
264    pub fn with_allowed_tools(mut self, tools: Vec<String>) -> Self {
265        self.allowed_tools = Some(tools);
266        self
267    }
268
269    /// Share the parent's resource lock table with spawned agents.
270    pub(crate) fn with_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
271        self.resource_locks = Some(locks);
272        self
273    }
274
275    /// Share the parent's observability manager with spawned agents.
276    pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
277        self.observability_manager = Some(manager);
278        self
279    }
280
281    /// Spawn an agent from a YAML string.
282    pub async fn spawn_from_yaml(&self, yaml: &str) -> Result<SpawnedAgent> {
283        let spec = AgentSpec::from_yaml_strict(yaml)?;
284        self.spawn_admitted(None, spec, None).await
285    }
286
287    /// Spawn an agent from a pre-built AgentSpec.
288    pub async fn spawn_from_spec(&self, spec: AgentSpec) -> Result<SpawnedAgent> {
289        self.spawn_admitted(None, spec, None).await
290    }
291
292    /// Validate an explicit child ID and spec without reserving capacity or building an agent.
293    #[doc(hidden)]
294    pub fn validate_explicit_child(&self, id: &str, spec: &AgentSpec) -> Result<()> {
295        self.validate_admission(id, spec)
296    }
297
298    /// Spawn an agent with a specific ID, used for session restore.
299    pub async fn spawn_with_id(&self, id: String, spec: AgentSpec) -> Result<SpawnedAgent> {
300        self.spawn_admitted(Some(id), spec, None).await
301    }
302
303    pub(crate) async fn spawn_from_yaml_file_with_id(
304        &self,
305        id: String,
306        path: &Path,
307    ) -> Result<SpawnedAgent> {
308        let yaml = std::fs::read_to_string(path).map_err(AgentError::IoError)?;
309        let spec = AgentSpec::from_yaml_strict(&yaml)?;
310        let base_dir = path.parent().map(Path::to_path_buf);
311        self.spawn_admitted(Some(id), spec, base_dir).await
312    }
313
314    async fn spawn_admitted(
315        &self,
316        explicit_id: Option<String>,
317        spec: AgentSpec,
318        base_dir: Option<PathBuf>,
319    ) -> Result<SpawnedAgent> {
320        let admitted = self.admit_spec(explicit_id, spec, base_dir)?;
321        let reservation = self.reserve_capacity()?;
322        self.spawn_inner(admitted, reservation).await
323    }
324
325    pub(crate) async fn spawn_with_reserved_capacity(
326        &self,
327        id: String,
328        spec: AgentSpec,
329        reservation: CapacityReservation,
330    ) -> Result<SpawnedAgent> {
331        let admitted = self.admit_spec(Some(id), spec, None)?;
332        self.spawn_inner(admitted, reservation).await
333    }
334
335    fn admit_spec(
336        &self,
337        explicit_id: Option<String>,
338        spec: AgentSpec,
339        base_dir: Option<PathBuf>,
340    ) -> Result<AdmittedChild> {
341        let id = explicit_id.unwrap_or_else(|| self.generate_id(&spec.name));
342        self.validate_admission(&id, &spec)?;
343        Ok(AdmittedChild { id, spec, base_dir })
344    }
345
346    fn validate_admission(&self, id: &str, spec: &AgentSpec) -> Result<()> {
347        spec.validate()?;
348        validate_child_id(id)?;
349
350        if spec
351            .spawner
352            .as_ref()
353            .is_some_and(|config| config.is_configured())
354        {
355            return Err(AgentError::InvalidSpec(format!(
356                "Spawned agent '{}' cannot configure an active nested spawner",
357                id
358            )));
359        }
360
361        self.validate_tool_allowlist(spec)?;
362        self.validate_shared_llm_aliases(spec)
363    }
364
365    /// Internal spawn after all child admission checks have passed.
366    async fn spawn_inner(
367        &self,
368        admitted: AdmittedChild,
369        reservation: CapacityReservation,
370    ) -> Result<SpawnedAgent> {
371        let AdmittedChild { id, spec, base_dir } = admitted;
372
373        let mut builder = match base_dir {
374            Some(base_dir) => AgentBuilder::from_spec_with_base_dir(spec.clone(), base_dir),
375            None => AgentBuilder::from_spec(spec.clone()),
376        };
377
378        if let Some(ref shared_reg) = self.llm_registry {
379            // Shared providers are authoritative. Child provider declarations are never constructed or allowed to read credentials in this mode.
380            builder =
381                builder.authoritative_llm_registry(shared_reg.clone(), self.llm_registry_observed);
382        } else {
383            // Local mode configures both the single llm declaration and the multi-alias llms map.
384            builder = builder.auto_configure_llms()?;
385        }
386
387        builder = builder.auto_configure_features()?;
388
389        if let Some(ref manager) = self.observability_manager {
390            builder = builder.observability(Arc::clone(manager));
391        }
392        if let Some(ref locks) = self.resource_locks {
393            builder = builder.with_shared_resource_locks(Arc::clone(locks));
394        }
395
396        if let Some(ref shared_storage) = self.storage {
397            let namespaced = Arc::new(NamespacedStorage::new(Arc::clone(shared_storage), &id));
398            builder = builder.storage(namespaced);
399        }
400
401        let agent = builder.build()?;
402        // A child is not ready for registration until required injected or configured storage capabilities and fact initialization have completed.
403        agent.init_storage().await?;
404
405        for (key, value) in &self.shared_context {
406            agent.set_context(key, value.clone())?;
407        }
408
409        info!(agent_id = %id, name = %spec.name, "Agent spawned");
410        Ok(SpawnedAgent::tracked(id, agent, spec, reservation))
411    }
412
413    /// Spawn from a named template with caller-provided variables.
414    ///
415    /// Template rendering merges two namespaces:
416    /// - Caller variables: top-level (`{{ name }}`, `{{ role }}`)
417    /// - Shared context: under `context.` prefix (`{{ context.world_name }}`)
418    pub async fn spawn_from_template(
419        &self,
420        template_name: &str,
421        variables: HashMap<String, String>,
422    ) -> Result<SpawnedAgent> {
423        let template = self.templates.get(template_name).ok_or_else(|| {
424            AgentError::Config(format!("Spawner template not found: {}", template_name))
425        })?;
426
427        let rendered = self.render_template(&template.content, &variables)?;
428        self.spawn_from_yaml(&rendered).await
429    }
430
431    /// Returns the number of registered and in-flight reserved agents.
432    pub fn spawned_count(&self) -> u32 {
433        self.agent_count.load(Ordering::Relaxed)
434    }
435
436    /// Returns a reference to the shared LLM registry, if configured.
437    pub fn llm_registry(&self) -> Option<&LLMRegistry> {
438        self.llm_registry.as_ref()
439    }
440
441    /// Returns a reference to the shared storage, if configured.
442    pub fn shared_storage(&self) -> Option<&Arc<dyn AgentStorage>> {
443        self.storage.as_ref()
444    }
445
446    /// Returns a reference to the resolved template map.
447    pub fn templates(&self) -> &HashMap<String, ResolvedTemplate> {
448        &self.templates
449    }
450
451    fn reserve_capacity(&self) -> Result<CapacityReservation> {
452        self.reserve_restore_capacity(1, 0)?
453            .pop()
454            .ok_or_else(|| AgentError::Config("Spawn capacity reservation failed".to_string()))
455    }
456
457    pub(crate) fn reserve_restore_capacity(
458        &self,
459        additions: usize,
460        removals: usize,
461    ) -> Result<Vec<CapacityReservation>> {
462        let additions = u32::try_from(additions)
463            .map_err(|_| AgentError::Config("Spawn capacity request is too large".to_string()))?;
464        let removals = u32::try_from(removals)
465            .map_err(|_| AgentError::Config("Spawn removal count is too large".to_string()))?;
466
467        loop {
468            let current = self.agent_count.load(Ordering::Acquire);
469            let after_removals = current.checked_sub(removals).ok_or_else(|| {
470                AgentError::Config("Restore removal count exceeds reserved capacity".to_string())
471            })?;
472            let final_count = after_removals.checked_add(additions).ok_or_else(|| {
473                AgentError::Config("Spawn capacity counter overflowed".to_string())
474            })?;
475            if self
476                .max_agents
477                .is_some_and(|max| final_count as usize > max)
478            {
479                return Err(AgentError::Config(format!(
480                    "Spawn limit exceeded by restored topology: {}/{}",
481                    final_count,
482                    self.max_agents.unwrap()
483                )));
484            }
485            if additions == 0 {
486                return Ok(Vec::new());
487            }
488
489            //
490            // Restore additions are reserved before staging. Temporary over-cap counts are allowed only when committed removals make the final topology fit.
491            //
492            let reserved_count = current.checked_add(additions).ok_or_else(|| {
493                AgentError::Config("Spawn capacity counter overflowed".to_string())
494            })?;
495            match self.agent_count.compare_exchange_weak(
496                current,
497                reserved_count,
498                Ordering::AcqRel,
499                Ordering::Acquire,
500            ) {
501                Ok(_) => {
502                    return Ok((0..additions)
503                        .map(|_| CapacityReservation::new(Arc::clone(&self.agent_count)))
504                        .collect());
505                }
506                Err(_) => continue,
507            }
508        }
509    }
510
511    fn generate_id(&self, spec_name: &str) -> String {
512        if let Some(ref prefix) = self.name_prefix {
513            let n = self.counter.fetch_add(1, Ordering::Relaxed);
514            return format!("{}{:03}", prefix, n);
515        }
516
517        let mut generated = String::with_capacity(spec_name.len());
518        for character in spec_name.chars() {
519            if character.is_ascii_alphanumeric() || matches!(character, '_' | '-' | '.') {
520                generated.push(character.to_ascii_lowercase());
521            } else if !generated.ends_with('_') {
522                generated.push('_');
523            }
524        }
525        let generated = generated.trim_matches(['_', '-', '.']).to_string();
526        if generated.is_empty() {
527            let n = self.counter.fetch_add(1, Ordering::Relaxed);
528            format!("agent_{n:03}")
529        } else {
530            generated
531        }
532    }
533
534    fn validate_tool_allowlist(&self, spec: &AgentSpec) -> Result<()> {
535        let Some(allowed) = self.allowed_tools.as_ref() else {
536            return Ok(());
537        };
538        let registry = create_builtin_registry();
539        let canonicalize = |tool: &str| {
540            registry
541                .canonical_id(tool)
542                .unwrap_or_else(|| tool.to_string())
543        };
544        let allowed: BTreeSet<String> = allowed.iter().map(|tool| canonicalize(tool)).collect();
545        let disallowed: BTreeSet<String> = spec
546            .tools
547            .iter()
548            .flatten()
549            .map(|tool| tool.name().to_string())
550            .filter(|tool| !allowed.contains(&canonicalize(tool)))
551            .collect();
552        if disallowed.is_empty() {
553            Ok(())
554        } else {
555            Err(AgentError::InvalidSpec(format!(
556                "Spawned agent declares tools outside the spawner allowlist: {}",
557                disallowed.into_iter().collect::<Vec<_>>().join(", ")
558            )))
559        }
560    }
561
562    fn validate_shared_llm_aliases(&self, spec: &AgentSpec) -> Result<()> {
563        let Some(registry) = self.llm_registry.as_ref() else {
564            return Ok(());
565        };
566        registry.default().map_err(|error| {
567            AgentError::Config(format!(
568                "Inherited LLM registry has no usable default: {error}"
569            ))
570        })?;
571
572        let mut required: BTreeSet<String> = spec.llms.keys().cloned().collect();
573        required.insert(spec.llm.get_default_alias());
574        if let Some(router) = spec.llm.get_router_alias() {
575            required.insert(router);
576        }
577        required.extend(spec.referenced_llm_aliases());
578        let missing: Vec<String> = required
579            .into_iter()
580            .filter(|alias| !registry.has(alias))
581            .collect();
582        if missing.is_empty() {
583            Ok(())
584        } else {
585            Err(AgentError::InvalidSpec(format!(
586                "Spawned agent requires inherited LLM alias(es) not present in the shared registry: {}",
587                missing.join(", ")
588            )))
589        }
590    }
591
592    /// Render a template string with caller variables and shared context.
593    fn render_template(
594        &self,
595        template_str: &str,
596        variables: &HashMap<String, String>,
597    ) -> Result<String> {
598        let mut env = Environment::new();
599        env.add_template("_spawn", template_str)
600            .map_err(|e| AgentError::TemplateError(format!("template parse error: {}", e)))?;
601
602        let tmpl = env
603            .get_template("_spawn")
604            .map_err(|e| AgentError::TemplateError(format!("template load error: {}", e)))?;
605
606        // Caller variables are top-level; shared context lives under "context".
607        let mut ctx = serde_json::Map::new();
608
609        for (k, v) in variables {
610            ctx.insert(k.clone(), serde_json::Value::String(v.clone()));
611        }
612
613        // Shared context as a nested object so {{ context.world_name }} works.
614        let context_obj = serde_json::Value::Object(
615            self.shared_context
616                .iter()
617                .map(|(k, v)| (k.clone(), v.clone()))
618                .collect(),
619        );
620        ctx.insert("context".to_string(), context_obj);
621
622        let ctx_value = serde_json::Value::Object(ctx);
623        let mj_value = minijinja::Value::from_serialize(&ctx_value);
624
625        tmpl.render(mj_value)
626            .map_err(|e| AgentError::TemplateError(format!("template render error: {}", e)))
627    }
628}
629
630fn validate_child_id(id: &str) -> Result<()> {
631    if id.is_empty() {
632        return Err(AgentError::InvalidSpec(
633            "Spawned agent ID cannot be empty".to_string(),
634        ));
635    }
636    if matches!(id, "." | "..") {
637        return Err(AgentError::InvalidSpec(format!(
638            "Spawned agent ID '{id}' is reserved"
639        )));
640    }
641    if id.len() > 128 {
642        return Err(AgentError::InvalidSpec(
643            "Spawned agent ID cannot exceed 128 bytes".to_string(),
644        ));
645    }
646    if !id
647        .bytes()
648        .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))
649    {
650        return Err(AgentError::InvalidSpec(format!(
651            "Spawned agent ID '{}' must contain only ASCII letters, digits, '_', '-' or '.'",
652            id
653        )));
654    }
655    if id.ends_with('.') {
656        return Err(AgentError::InvalidSpec(
657            "Spawned agent ID cannot end with a dot".to_string(),
658        ));
659    }
660    let windows_stem = id.split('.').next().unwrap_or(id).to_ascii_uppercase();
661    if matches!(
662        windows_stem.as_str(),
663        "CON" | "PRN" | "AUX" | "NUL" | "CLOCK$"
664    ) || windows_stem
665        .strip_prefix("COM")
666        .is_some_and(|suffix| matches!(suffix, "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9"))
667        || windows_stem.strip_prefix("LPT").is_some_and(|suffix| {
668            matches!(suffix, "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9")
669        })
670    {
671        return Err(AgentError::InvalidSpec(format!(
672            "Spawned agent ID '{id}' is reserved on Windows"
673        )));
674    }
675    Ok(())
676}
677
678fn validate_name_prefix(prefix: &str) -> Result<()> {
679    if prefix.is_empty() {
680        return Err(AgentError::InvalidSpec(
681            "Spawner name_prefix cannot be empty".to_string(),
682        ));
683    }
684    validate_child_id(&format!("{prefix}001")).map_err(|error| {
685        AgentError::InvalidSpec(format!("Invalid spawner name_prefix '{prefix}': {error}"))
686    })
687}
688
689impl Default for AgentSpawner {
690    fn default() -> Self {
691        Self::new()
692    }
693}
694
695#[cfg(test)]
696mod tests {
697    use super::super::registry::AgentRegistry;
698    use super::*;
699    use ai_agents_llm::mock::MockLLMProvider;
700
701    fn shared_registry() -> LLMRegistry {
702        let mut registry = LLMRegistry::new();
703        registry.register("default", Arc::new(MockLLMProvider::new("shared")));
704        registry.set_default("default");
705        registry
706    }
707
708    fn shared_spawner() -> AgentSpawner {
709        AgentSpawner::new().with_shared_llms(shared_registry())
710    }
711
712    #[test]
713    fn test_generate_id_with_prefix() {
714        let spawner = AgentSpawner::new().with_name_prefix("npc_").unwrap();
715        assert_eq!(spawner.generate_id("Gormund"), "npc_001");
716        assert_eq!(spawner.generate_id("Elena"), "npc_002");
717    }
718
719    #[test]
720    fn test_generate_id_without_prefix() {
721        let spawner = AgentSpawner::new();
722        assert_eq!(spawner.generate_id("My Agent"), "my_agent");
723        assert_eq!(spawner.generate_id("Test.Bot"), "test.bot");
724        assert_eq!(spawner.generate_id("작업자"), "agent_001");
725        assert_eq!(spawner.generate_id("TestBot"), "testbot");
726    }
727
728    #[test]
729    fn test_capacity_reservation_is_atomic_and_released_on_drop() {
730        let spawner = AgentSpawner::new().with_max_agents(1);
731        let reservation = spawner.reserve_capacity().unwrap();
732        assert_eq!(spawner.spawned_count(), 1);
733        assert!(spawner.reserve_capacity().is_err());
734        drop(reservation);
735        assert_eq!(spawner.spawned_count(), 0);
736        assert!(spawner.reserve_capacity().is_ok());
737    }
738
739    #[test]
740    fn restore_reservation_uses_committed_removal_credit_without_leaking() {
741        let spawner = AgentSpawner::new().with_max_agents(2);
742        let current = spawner.reserve_restore_capacity(2, 0).unwrap();
743        assert_eq!(spawner.spawned_count(), 2);
744
745        let replacement = spawner.reserve_restore_capacity(1, 1).unwrap();
746        assert_eq!(spawner.spawned_count(), 3);
747        drop(replacement);
748        assert_eq!(spawner.spawned_count(), 2);
749        drop(current);
750        assert_eq!(spawner.spawned_count(), 0);
751    }
752
753    #[tokio::test]
754    async fn test_concurrent_spawns_cannot_exceed_capacity() {
755        let spawner = Arc::new(shared_spawner().with_max_agents(3));
756        let barrier = Arc::new(tokio::sync::Barrier::new(9));
757        let mut tasks = Vec::new();
758        for index in 0..8 {
759            let spawner = Arc::clone(&spawner);
760            let barrier = Arc::clone(&barrier);
761            tasks.push(tokio::spawn(async move {
762                let spec = AgentSpec {
763                    name: format!("Worker {index}"),
764                    system_prompt: "worker".to_string(),
765                    ..AgentSpec::default()
766                };
767                barrier.wait().await;
768                spawner.spawn_with_id(format!("worker_{index}"), spec).await
769            }));
770        }
771        barrier.wait().await;
772
773        let mut spawned = Vec::new();
774        let mut rejected = 0;
775        for task in tasks {
776            match task.await.unwrap() {
777                Ok(agent) => spawned.push(agent),
778                Err(error) => {
779                    assert!(error.to_string().contains("Spawn limit exceeded"));
780                    rejected += 1;
781                }
782            }
783        }
784        assert_eq!(spawned.len(), 3);
785        assert_eq!(rejected, 5);
786        assert_eq!(spawner.spawned_count(), 3);
787        drop(spawned);
788        assert_eq!(spawner.spawned_count(), 0);
789    }
790
791    #[test]
792    fn test_render_template_basic() {
793        let spawner = AgentSpawner::new()
794            .with_shared_context("world_name", serde_json::json!("Fantasy Land"));
795
796        let template =
797            "name: {{ name }}\nsystem_prompt: You are {{ name }} in {{ context.world_name }}.";
798        let mut vars = HashMap::new();
799        vars.insert("name".to_string(), "Gormund".to_string());
800
801        let rendered = spawner.render_template(template, &vars).unwrap();
802        assert!(rendered.contains("name: Gormund"));
803        assert!(rendered.contains("Fantasy Land"));
804    }
805
806    #[tokio::test]
807    async fn test_tool_allowlist_rejects_instead_of_stripping() {
808        let spawner =
809            shared_spawner().with_allowed_tools(vec!["echo".to_string(), "calculator".to_string()]);
810        let yaml = r#"
811name: Test
812system_prompt: test
813tools:
814  - echo
815  - file
816  - http
817"#;
818        let error = spawner.spawn_from_yaml(yaml).await.unwrap_err().to_string();
819        assert!(error.contains("outside the spawner allowlist"), "{error}");
820        assert!(error.contains("file"), "{error}");
821        assert!(error.contains("http"), "{error}");
822        assert_eq!(spawner.spawned_count(), 0);
823    }
824
825    #[test]
826    fn test_validate_child_id_rejects_non_portable_values() {
827        assert!(validate_child_id("worker-01").is_ok());
828        assert!(validate_child_id("worker_01").is_ok());
829        assert!(validate_child_id("worker.01").is_ok());
830        assert!(validate_child_id("").is_err());
831        assert!(validate_child_id(".").is_err());
832        assert!(validate_child_id("..").is_err());
833        assert!(validate_child_id("../worker").is_err());
834        assert!(validate_child_id("worker name").is_err());
835        assert!(validate_child_id("worker/01").is_err());
836        assert!(validate_child_id("wörker").is_err());
837        assert!(validate_child_id("worker.").is_err());
838        for reserved in ["CON", "con.txt", "NUL", "COM1", "LPT9.log"] {
839            assert!(validate_child_id(reserved).is_err(), "{reserved}");
840        }
841    }
842
843    #[test]
844    fn name_prefix_is_validated_when_configured() {
845        assert!(AgentSpawner::new().with_name_prefix("worker_").is_ok());
846        assert!(AgentSpawner::new().with_name_prefix("").is_err());
847        assert!(AgentSpawner::new().with_name_prefix("bad/path").is_err());
848        assert!(
849            AgentSpawner::new()
850                .with_name_prefix("x".repeat(126))
851                .is_err()
852        );
853    }
854
855    #[test]
856    fn test_tool_allowlist_compares_builtin_canonical_identity() {
857        let spawner = shared_spawner().with_allowed_tools(vec!["Copy Path".to_string()]);
858        let spec = AgentSpec::from_yaml_strict(
859            "name: Worker\nsystem_prompt: worker\ntools:\n  - copy_path\n",
860        )
861        .unwrap();
862
863        spawner.validate_explicit_child("worker", &spec).unwrap();
864    }
865
866    #[tokio::test]
867    async fn test_shared_llms_ignore_child_provider_declarations() {
868        let spawner = shared_spawner();
869        let yaml = r#"
870name: SharedChild
871system_prompt: shared
872llm:
873  provider: definitely-not-a-provider
874  model: unavailable
875"#;
876        let spawned = spawner.spawn_from_yaml(yaml).await.unwrap();
877        assert_eq!(spawned.agent.llm_registry().default_alias(), "default");
878        drop(spawned);
879        assert_eq!(spawner.spawned_count(), 0);
880    }
881
882    #[tokio::test]
883    async fn test_shared_llms_require_declared_child_aliases() {
884        let spawner = shared_spawner();
885        let yaml = r#"
886name: AliasChild
887system_prompt: shared
888llm:
889  default: specialist
890llms:
891  specialist:
892    provider: definitely-not-a-provider
893    model: unavailable
894"#;
895        let error = spawner.spawn_from_yaml(yaml).await.unwrap_err().to_string();
896        assert!(error.contains("specialist"), "{error}");
897        assert_eq!(spawner.spawned_count(), 0);
898    }
899
900    #[test]
901    fn test_shared_llms_require_explicit_subsystem_aliases() {
902        let spawner = shared_spawner();
903        let spec = AgentSpec::from_yaml_strict(
904            "name: Worker\nsystem_prompt: worker\nmemory:\n  type: compacting\n  summarizer_llm: specialist\n",
905        )
906        .unwrap();
907
908        let error = spawner
909            .validate_explicit_child("worker", &spec)
910            .unwrap_err()
911            .to_string();
912        assert!(error.contains("specialist"), "{error}");
913    }
914
915    #[tokio::test]
916    async fn test_local_mode_configures_single_and_multi_declarations() {
917        let single = r#"
918name: SingleChild
919system_prompt: local
920llm:
921  provider: ollama
922  model: local-single
923"#;
924        let multi = r#"
925name: MultiChild
926system_prompt: local
927llm:
928  default: specialist
929llms:
930  specialist:
931    provider: ollama
932    model: local-multi
933"#;
934        let spawner = AgentSpawner::new();
935        let single = spawner.spawn_from_yaml(single).await.unwrap();
936        assert_eq!(
937            single
938                .agent
939                .llm_registry()
940                .default()
941                .unwrap()
942                .provider_name(),
943            "ollama"
944        );
945        drop(single);
946
947        let multi = spawner.spawn_from_yaml(multi).await.unwrap();
948        assert!(multi.agent.llm_registry().has("specialist"));
949        assert_eq!(multi.agent.llm_registry().default_alias(), "specialist");
950        drop(multi);
951        assert_eq!(spawner.spawned_count(), 0);
952    }
953
954    #[test]
955    fn test_explicit_admission_validation_does_not_reserve_capacity() {
956        let spawner = shared_spawner().with_max_agents(1);
957        let spec = AgentSpec {
958            name: "Child".to_string(),
959            system_prompt: "child".to_string(),
960            ..AgentSpec::default()
961        };
962
963        spawner.validate_explicit_child("child", &spec).unwrap();
964        assert_eq!(spawner.spawned_count(), 0);
965    }
966
967    #[tokio::test]
968    async fn test_admission_rejects_nested_spawner_and_invalid_explicit_id() {
969        let nested = r#"
970name: NestedChild
971system_prompt: nested
972spawner:
973  management_tools: true
974"#;
975        let spawner = shared_spawner();
976        let error = spawner
977            .spawn_from_yaml(nested)
978            .await
979            .unwrap_err()
980            .to_string();
981        assert!(error.contains("nested spawner"), "{error}");
982
983        let inert = "name: InertChild\nsystem_prompt: inert\nspawner: {}\n";
984        let spawned = spawner.spawn_from_yaml(inert).await.unwrap();
985        drop(spawned);
986
987        let spec = AgentSpec {
988            name: "Child".to_string(),
989            system_prompt: "child".to_string(),
990            ..AgentSpec::default()
991        };
992        let error = spawner
993            .spawn_with_id("../child".to_string(), spec)
994            .await
995            .unwrap_err()
996            .to_string();
997        assert!(error.contains("Spawned agent ID"), "{error}");
998    }
999
1000    #[tokio::test]
1001    async fn test_build_failure_releases_reserved_capacity_once() {
1002        let spawner = AgentSpawner::new().with_max_agents(1);
1003        let invalid = r#"
1004name: InvalidChild
1005system_prompt: invalid
1006llm:
1007  provider: definitely-not-a-provider
1008  model: unavailable
1009"#;
1010        assert!(spawner.spawn_from_yaml(invalid).await.is_err());
1011        assert_eq!(spawner.spawned_count(), 0);
1012
1013        let valid = r#"
1014name: ValidChild
1015system_prompt: valid
1016llm:
1017  provider: ollama
1018  model: local
1019"#;
1020        let spawned = spawner.spawn_from_yaml(valid).await.unwrap();
1021        assert_eq!(spawner.spawned_count(), 1);
1022        drop(spawned);
1023        assert_eq!(spawner.spawned_count(), 0);
1024    }
1025
1026    #[tokio::test]
1027    async fn test_registration_failure_releases_reserved_capacity_once() {
1028        let spawner = shared_spawner().with_max_agents(2);
1029        let registry = AgentRegistry::new();
1030        let spec = AgentSpec {
1031            name: "Child".to_string(),
1032            system_prompt: "child".to_string(),
1033            ..AgentSpec::default()
1034        };
1035
1036        let first = spawner
1037            .spawn_with_id("child".to_string(), spec.clone())
1038            .await
1039            .unwrap();
1040        registry.register(first).await.unwrap();
1041        assert_eq!(spawner.spawned_count(), 1);
1042
1043        let duplicate = spawner
1044            .spawn_with_id("child".to_string(), spec)
1045            .await
1046            .unwrap();
1047        assert_eq!(spawner.spawned_count(), 2);
1048        assert!(registry.register(duplicate).await.is_err());
1049        assert_eq!(spawner.spawned_count(), 1);
1050
1051        assert!(registry.remove("child").await.is_some());
1052        assert_eq!(spawner.spawned_count(), 0);
1053        assert!(registry.remove("child").await.is_none());
1054        assert_eq!(spawner.spawned_count(), 0);
1055    }
1056
1057    #[test]
1058    fn test_with_template_plain_string() {
1059        let spawner =
1060            AgentSpawner::new().with_template("basic", "name: {{ name }}\nsystem_prompt: hi");
1061        let tpl = spawner.templates().get("basic").unwrap();
1062        assert_eq!(tpl.content, "name: {{ name }}\nsystem_prompt: hi");
1063        assert!(tpl.description.is_none());
1064        assert!(tpl.variables.is_none());
1065    }
1066
1067    #[test]
1068    fn test_with_templates_resolved() {
1069        let mut templates = HashMap::new();
1070        templates.insert(
1071            "base".to_string(),
1072            ResolvedTemplate {
1073                content: "name: {{ name }}".to_string(),
1074                description: Some("Test template".to_string()),
1075                variables: Some({
1076                    let mut v = HashMap::new();
1077                    v.insert("role".to_string(), "occupation".to_string());
1078                    v
1079                }),
1080            },
1081        );
1082        let spawner = AgentSpawner::new().with_templates(templates);
1083        let tpl = spawner.templates().get("base").unwrap();
1084        assert_eq!(tpl.description.as_deref(), Some("Test template"));
1085        assert_eq!(
1086            tpl.variables.as_ref().unwrap().get("role").unwrap(),
1087            "occupation"
1088        );
1089    }
1090}