Skip to main content

ralph_core/
hat_registry.rs

1//! Hat registry for managing agent personas.
2
3use crate::config::{HatConfig, RalphConfig};
4use ralph_proto::{Hat, HatId, Topic};
5use std::collections::{HashMap, HashSet};
6
7/// Registry for managing and creating hats from configuration.
8#[derive(Debug, Default)]
9pub struct HatRegistry {
10    hats: HashMap<HatId, Hat>,
11    configs: HashMap<HatId, HatConfig>,
12    /// Prefix index for O(1) early-exit on no-match lookups.
13    /// Contains all first segments of subscription patterns (e.g., "task" from "task.*").
14    /// Also contains "*" if any global wildcard exists.
15    prefix_index: HashSet<String>,
16}
17
18impl HatRegistry {
19    /// Creates a new empty registry.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Creates a registry from configuration.
25    ///
26    /// Empty config → empty registry (HatlessRalph is the fallback, not default hats).
27    pub fn from_config(config: &RalphConfig) -> Self {
28        let mut registry = Self::new();
29
30        for (id, hat_config) in &config.hats {
31            let hat = Self::hat_from_config(id, hat_config);
32            registry.register_with_config(hat, hat_config.clone());
33        }
34
35        registry
36    }
37
38    /// Creates a Hat from HatConfig.
39    fn hat_from_config(id: &str, config: &HatConfig) -> Hat {
40        let mut hat = Hat::new(id, &config.name);
41        hat.description = config.description.clone().unwrap_or_default();
42        hat.subscriptions = config.trigger_topics();
43        hat.publishes = config.publish_topics();
44        hat.instructions = config.instructions.clone();
45        hat
46    }
47
48    /// Registers a hat with the registry.
49    pub fn register(&mut self, hat: Hat) {
50        self.index_hat_subscriptions(&hat);
51        self.hats.insert(hat.id.clone(), hat);
52    }
53
54    /// Registers a hat with its configuration.
55    pub fn register_with_config(&mut self, hat: Hat, config: HatConfig) {
56        let id = hat.id.clone();
57        self.index_hat_subscriptions(&hat);
58        self.hats.insert(id.clone(), hat);
59        self.configs.insert(id, config);
60    }
61
62    /// Indexes a hat's subscriptions for O(1) prefix lookup.
63    fn index_hat_subscriptions(&mut self, hat: &Hat) {
64        for sub in &hat.subscriptions {
65            let pattern = sub.as_str();
66            // Global wildcard matches everything - mark it specially
67            if pattern == "*" {
68                self.prefix_index.insert("*".to_string());
69            } else {
70                // Extract first segment (e.g., "task" from "task.*" or "task.start")
71                if let Some(prefix) = pattern.split('.').next() {
72                    self.prefix_index.insert(prefix.to_string());
73                }
74            }
75        }
76    }
77
78    /// Gets a hat by ID.
79    pub fn get(&self, id: &HatId) -> Option<&Hat> {
80        self.hats.get(id)
81    }
82
83    /// Gets a hat's configuration by ID.
84    pub fn get_config(&self, id: &HatId) -> Option<&HatConfig> {
85        self.configs.get(id)
86    }
87
88    /// Returns all hats in the registry.
89    pub fn all(&self) -> impl Iterator<Item = &Hat> {
90        self.hats.values()
91    }
92
93    /// Returns all hat IDs.
94    pub fn ids(&self) -> impl Iterator<Item = &HatId> {
95        self.hats.keys()
96    }
97
98    /// Returns the number of registered hats.
99    pub fn len(&self) -> usize {
100        self.hats.len()
101    }
102
103    /// Returns true if no hats are registered.
104    pub fn is_empty(&self) -> bool {
105        self.hats.is_empty()
106    }
107
108    /// Finds all hats subscribed to a topic.
109    pub fn subscribers(&self, topic: &Topic) -> Vec<&Hat> {
110        self.hats
111            .values()
112            .filter(|hat| hat.is_subscribed(topic))
113            .collect()
114    }
115
116    /// Finds the first hat that would be triggered by a topic.
117    /// Returns the hat ID if found, used for event logging.
118    pub fn find_by_trigger(&self, topic: &str) -> Option<&HatId> {
119        let topic = Topic::new(topic);
120        self.hats
121            .values()
122            .find(|hat| hat.is_subscribed(&topic))
123            .map(|hat| &hat.id)
124    }
125
126    /// Returns true if any hat is subscribed to the given topic.
127    pub fn has_subscriber(&self, topic: &str) -> bool {
128        let topic = Topic::new(topic);
129        self.hats.values().any(|hat| hat.is_subscribed(&topic))
130    }
131
132    /// Returns the first hat subscribed to the given topic.
133    ///
134    /// Uses prefix index for O(1) early-exit when the topic prefix doesn't match
135    /// any subscription pattern.
136    pub fn get_for_topic(&self, topic: &str) -> Option<&Hat> {
137        // Fast path: Check if any subscription could possibly match this topic
138        // If we have a global wildcard "*", we must do the full scan
139        if !self.prefix_index.contains("*") {
140            // Extract prefix from topic (e.g., "task" from "task.start")
141            let topic_prefix = topic.split('.').next().unwrap_or(topic);
142            if !self.prefix_index.contains(topic_prefix) {
143                // No subscription has this prefix - early exit
144                return None;
145            }
146        }
147
148        // Fall back to full linear scan
149        self.hats.values().find(|hat| hat.is_subscribed_str(topic))
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use std::time::Instant;
157
158    #[test]
159    fn test_empty_config_creates_empty_registry() {
160        let config = RalphConfig::default();
161        let registry = HatRegistry::from_config(&config);
162
163        // Empty config → empty registry (HatlessRalph is the fallback, not default hats)
164        assert!(registry.is_empty());
165        assert_eq!(registry.len(), 0);
166    }
167
168    #[test]
169    fn test_custom_hats_from_config() {
170        let yaml = r#"
171hats:
172  implementer:
173    name: "Implementer"
174    triggers: ["task.*"]
175  reviewer:
176    name: "Reviewer"
177    triggers: ["impl.*"]
178"#;
179        let config: RalphConfig = serde_yaml::from_str(yaml).unwrap();
180        let registry = HatRegistry::from_config(&config);
181
182        assert_eq!(registry.len(), 2);
183
184        let impl_hat = registry.get(&HatId::new("implementer")).unwrap();
185        assert!(impl_hat.is_subscribed(&Topic::new("task.start")));
186        assert!(!impl_hat.is_subscribed(&Topic::new("impl.done")));
187
188        let review_hat = registry.get(&HatId::new("reviewer")).unwrap();
189        assert!(review_hat.is_subscribed(&Topic::new("impl.done")));
190    }
191
192    #[test]
193    fn test_has_subscriber() {
194        let yaml = r#"
195hats:
196  impl:
197    name: "Implementer"
198    triggers: ["task.*"]
199"#;
200        let config: RalphConfig = serde_yaml::from_str(yaml).unwrap();
201        let registry = HatRegistry::from_config(&config);
202
203        assert!(registry.has_subscriber("task.start"));
204        assert!(!registry.has_subscriber("build.task"));
205    }
206
207    #[test]
208    fn test_get_for_topic() {
209        let yaml = r#"
210hats:
211  impl:
212    name: "Implementer"
213    triggers: ["task.*"]
214"#;
215        let config: RalphConfig = serde_yaml::from_str(yaml).unwrap();
216        let registry = HatRegistry::from_config(&config);
217
218        let hat = registry.get_for_topic("task.start");
219        assert!(hat.is_some());
220        assert_eq!(hat.unwrap().id.as_str(), "impl");
221
222        let no_hat = registry.get_for_topic("build.task");
223        assert!(no_hat.is_none());
224    }
225
226    #[test]
227    fn test_empty_registry_has_no_subscribers() {
228        let config = RalphConfig::default();
229        let registry = HatRegistry::from_config(&config);
230
231        // Empty config → no subscribers (HatlessRalph handles orphaned events)
232        assert!(!registry.has_subscriber("build.task"));
233        assert!(registry.get_for_topic("build.task").is_none());
234    }
235
236    #[test]
237    fn test_find_subscribers() {
238        let yaml = r#"
239hats:
240  impl:
241    name: "Implementer"
242    triggers: ["task.*", "review.done"]
243  reviewer:
244    name: "Reviewer"
245    triggers: ["impl.*"]
246"#;
247        let config: RalphConfig = serde_yaml::from_str(yaml).unwrap();
248        let registry = HatRegistry::from_config(&config);
249
250        let task_subs = registry.subscribers(&Topic::new("task.start"));
251        assert_eq!(task_subs.len(), 1);
252        assert_eq!(task_subs[0].id.as_str(), "impl");
253
254        let impl_subs = registry.subscribers(&Topic::new("impl.done"));
255        assert_eq!(impl_subs.len(), 1);
256        assert_eq!(impl_subs[0].id.as_str(), "reviewer");
257    }
258
259    /// Benchmark test for get_for_topic() performance.
260    /// Run with: cargo test -p ralph-core bench_get_for_topic -- --nocapture
261    #[test]
262    fn bench_get_for_topic_baseline() {
263        // Create registry with 20 hats (realistic production scenario)
264        let mut yaml = String::from("hats:\n");
265        for i in 0..20 {
266            yaml.push_str(&format!(
267                "  hat{}:\n    name: \"Hat {}\"\n    triggers: [\"topic{}.*\", \"other{}.event\"]\n",
268                i, i, i, i
269            ));
270        }
271        let config: RalphConfig = serde_yaml::from_str(&yaml).unwrap();
272        let registry = HatRegistry::from_config(&config);
273
274        // Topics to test - mix of matches and non-matches
275        let topics = [
276            "topic0.start",    // First hat match
277            "topic10.event",   // Middle hat match
278            "topic19.done",    // Last hat match
279            "nomatch.topic",   // No match
280        ];
281
282        const ITERATIONS: u32 = 100_000;
283
284        let start = Instant::now();
285        for _ in 0..ITERATIONS {
286            for topic in &topics {
287                let _ = registry.get_for_topic(topic);
288            }
289        }
290        let elapsed = start.elapsed();
291
292        let ops = u64::from(ITERATIONS) * (topics.len() as u64);
293        let ns_per_op = elapsed.as_nanos() / u128::from(ops);
294
295        println!("\n=== get_for_topic() Baseline ===");
296        println!("Registry size: {} hats", registry.len());
297        println!("Operations: {}", ops);
298        println!("Total time: {:?}", elapsed);
299        println!("Time per operation: {} ns", ns_per_op);
300        println!("================================\n");
301
302        // Assert reasonable performance (sanity check)
303        assert!(ns_per_op < 10_000, "Performance degraded: {} ns/op", ns_per_op);
304    }
305}