Skip to main content

punch_kernel/
registry.rs

1//! Agent template registry.
2//!
3//! Loads [`FighterManifest`] templates from `agent.toml` files found in a
4//! designated agents directory. Templates can then be looked up by name to
5//! quickly spawn pre-configured fighters.
6
7use std::collections::HashMap;
8use std::path::Path;
9
10use tracing::{debug, info, instrument, warn};
11
12use punch_types::{FighterManifest, PunchError, PunchResult};
13
14/// A registry of named agent templates.
15///
16/// Templates are loaded from TOML files and cached in memory. The registry is
17/// **not** `Sync` by itself — wrap it in an `Arc<RwLock<_>>` if concurrent
18/// access is needed (the [`Ring`](crate::ring::Ring) does this internally via
19/// its own synchronization).
20#[derive(Debug, Default)]
21pub struct AgentRegistry {
22    templates: HashMap<String, FighterManifest>,
23}
24
25impl AgentRegistry {
26    /// Create an empty registry.
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    /// Scan `agents_dir` for agent templates and load them.
32    ///
33    /// The expected directory structure is:
34    ///
35    /// ```text
36    /// agents/
37    ///   coder/
38    ///     agent.toml
39    ///   reviewer/
40    ///     agent.toml
41    /// ```
42    ///
43    /// Each `agent.toml` is deserialized into a [`FighterManifest`]. The
44    /// directory name is used as the template key (lowercased).
45    #[instrument(skip(self), fields(dir = %agents_dir.display()))]
46    pub fn load_templates(&mut self, agents_dir: &Path) -> PunchResult<()> {
47        if !agents_dir.is_dir() {
48            return Err(PunchError::Config(format!(
49                "agents directory does not exist: {}",
50                agents_dir.display()
51            )));
52        }
53
54        let entries = std::fs::read_dir(agents_dir).map_err(|e| {
55            PunchError::Config(format!(
56                "failed to read agents directory {}: {}",
57                agents_dir.display(),
58                e
59            ))
60        })?;
61
62        let mut loaded = 0usize;
63
64        for entry in entries {
65            let entry = match entry {
66                Ok(e) => e,
67                Err(e) => {
68                    warn!(error = %e, "failed to read directory entry");
69                    continue;
70                }
71            };
72
73            let path = entry.path();
74            if !path.is_dir() {
75                continue;
76            }
77
78            let toml_path = path.join("agent.toml");
79            if !toml_path.exists() {
80                debug!(dir = %path.display(), "skipping directory without agent.toml");
81                continue;
82            }
83
84            let template_name = path
85                .file_name()
86                .and_then(|n| n.to_str())
87                .map(|s| s.to_lowercase())
88                .unwrap_or_default();
89
90            if template_name.is_empty() {
91                warn!(dir = %path.display(), "could not determine template name");
92                continue;
93            }
94
95            match self.load_single_template(&toml_path) {
96                Ok(manifest) => {
97                    info!(template = %template_name, name = %manifest.name, "loaded agent template");
98                    self.templates.insert(template_name, manifest);
99                    loaded += 1;
100                }
101                Err(e) => {
102                    warn!(
103                        template = %template_name,
104                        error = %e,
105                        "failed to load agent template"
106                    );
107                }
108            }
109        }
110
111        info!(loaded, "agent template scan complete");
112        Ok(())
113    }
114
115    /// Load a single `agent.toml` file into a [`FighterManifest`].
116    fn load_single_template(&self, path: &Path) -> PunchResult<FighterManifest> {
117        let content = std::fs::read_to_string(path)
118            .map_err(|e| PunchError::Config(format!("failed to read {}: {}", path.display(), e)))?;
119
120        let manifest: FighterManifest = toml::from_str(&content).map_err(|e| {
121            PunchError::Config(format!("failed to parse {}: {}", path.display(), e))
122        })?;
123
124        Ok(manifest)
125    }
126
127    /// Retrieve a template by name (case-insensitive lookup).
128    pub fn get_template(&self, name: &str) -> Option<&FighterManifest> {
129        self.templates.get(&name.to_lowercase())
130    }
131
132    /// List all registered template names.
133    pub fn list_templates(&self) -> Vec<String> {
134        let mut names: Vec<String> = self.templates.keys().cloned().collect();
135        names.sort();
136        names
137    }
138
139    /// Register a template manually (e.g. from an API call).
140    pub fn register(&mut self, name: String, manifest: FighterManifest) {
141        self.templates.insert(name.to_lowercase(), manifest);
142    }
143
144    /// Remove a template by name.
145    pub fn unregister(&mut self, name: &str) -> Option<FighterManifest> {
146        self.templates.remove(&name.to_lowercase())
147    }
148
149    /// Number of registered templates.
150    pub fn len(&self) -> usize {
151        self.templates.len()
152    }
153
154    /// Whether the registry is empty.
155    pub fn is_empty(&self) -> bool {
156        self.templates.is_empty()
157    }
158}