tirea_agentos/composition/registry/
agent.rs

1use super::sorted_registry_ids;
2use super::traits::{AgentRegistry, AgentRegistryError};
3use crate::composition::AgentDefinition;
4use std::collections::HashMap;
5use std::sync::{Arc, RwLock};
6
7#[derive(Debug, Clone, Default)]
8pub struct InMemoryAgentRegistry {
9    agents: HashMap<String, AgentDefinition>,
10}
11
12impl InMemoryAgentRegistry {
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    pub fn len(&self) -> usize {
18        self.agents.len()
19    }
20
21    pub fn is_empty(&self) -> bool {
22        self.agents.is_empty()
23    }
24
25    pub fn get(&self, id: &str) -> Option<AgentDefinition> {
26        self.agents.get(id).cloned()
27    }
28
29    pub fn ids(&self) -> impl Iterator<Item = &String> {
30        self.agents.keys()
31    }
32
33    pub fn register(
34        &mut self,
35        agent_id: impl Into<String>,
36        mut def: AgentDefinition,
37    ) -> Result<(), AgentRegistryError> {
38        let agent_id = agent_id.into();
39        if self.agents.contains_key(&agent_id) {
40            return Err(AgentRegistryError::AgentIdConflict(agent_id));
41        }
42        // The registry key is canonical to avoid mismatches.
43        def.id = agent_id.clone();
44        self.agents.insert(agent_id, def);
45        Ok(())
46    }
47
48    pub fn upsert(&mut self, agent_id: impl Into<String>, mut def: AgentDefinition) {
49        let agent_id = agent_id.into();
50        def.id = agent_id.clone();
51        self.agents.insert(agent_id, def);
52    }
53
54    pub fn extend_upsert(&mut self, defs: HashMap<String, AgentDefinition>) {
55        for (id, def) in defs {
56            self.upsert(id, def);
57        }
58    }
59
60    pub fn extend_registry(&mut self, other: &dyn AgentRegistry) -> Result<(), AgentRegistryError> {
61        for (id, def) in other.snapshot() {
62            self.register(id, def)?;
63        }
64        Ok(())
65    }
66}
67
68impl AgentRegistry for InMemoryAgentRegistry {
69    fn len(&self) -> usize {
70        self.len()
71    }
72
73    fn get(&self, id: &str) -> Option<AgentDefinition> {
74        self.get(id)
75    }
76
77    fn ids(&self) -> Vec<String> {
78        sorted_registry_ids(&self.agents)
79    }
80
81    fn snapshot(&self) -> HashMap<String, AgentDefinition> {
82        self.agents.clone()
83    }
84}
85
86#[derive(Clone, Default)]
87pub struct CompositeAgentRegistry {
88    registries: Vec<Arc<dyn AgentRegistry>>,
89    cached_snapshot: Arc<RwLock<HashMap<String, AgentDefinition>>>,
90}
91
92impl std::fmt::Debug for CompositeAgentRegistry {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        let snapshot = match self.cached_snapshot.read() {
95            Ok(guard) => guard,
96            Err(poisoned) => poisoned.into_inner(),
97        };
98        f.debug_struct("CompositeAgentRegistry")
99            .field("registries", &self.registries.len())
100            .field("len", &snapshot.len())
101            .finish()
102    }
103}
104
105impl CompositeAgentRegistry {
106    pub fn try_new(
107        regs: impl IntoIterator<Item = Arc<dyn AgentRegistry>>,
108    ) -> Result<Self, AgentRegistryError> {
109        let registries: Vec<Arc<dyn AgentRegistry>> = regs.into_iter().collect();
110        let merged = Self::merge_snapshots(&registries)?;
111        Ok(Self {
112            registries,
113            cached_snapshot: Arc::new(RwLock::new(merged)),
114        })
115    }
116
117    fn merge_snapshots(
118        registries: &[Arc<dyn AgentRegistry>],
119    ) -> Result<HashMap<String, AgentDefinition>, AgentRegistryError> {
120        let mut merged = InMemoryAgentRegistry::new();
121        for reg in registries {
122            merged.extend_registry(reg.as_ref())?;
123        }
124        Ok(merged.snapshot())
125    }
126
127    fn refresh_snapshot(&self) -> Result<HashMap<String, AgentDefinition>, AgentRegistryError> {
128        Self::merge_snapshots(&self.registries)
129    }
130
131    fn read_cached_snapshot(&self) -> HashMap<String, AgentDefinition> {
132        match self.cached_snapshot.read() {
133            Ok(guard) => guard.clone(),
134            Err(poisoned) => poisoned.into_inner().clone(),
135        }
136    }
137
138    fn write_cached_snapshot(&self, snapshot: HashMap<String, AgentDefinition>) {
139        match self.cached_snapshot.write() {
140            Ok(mut guard) => *guard = snapshot,
141            Err(poisoned) => *poisoned.into_inner() = snapshot,
142        };
143    }
144}
145
146impl AgentRegistry for CompositeAgentRegistry {
147    fn len(&self) -> usize {
148        self.snapshot().len()
149    }
150
151    fn get(&self, id: &str) -> Option<AgentDefinition> {
152        self.snapshot().get(id).cloned()
153    }
154
155    fn ids(&self) -> Vec<String> {
156        let snapshot = self.snapshot();
157        sorted_registry_ids(&snapshot)
158    }
159
160    fn snapshot(&self) -> HashMap<String, AgentDefinition> {
161        match self.refresh_snapshot() {
162            Ok(snapshot) => {
163                self.write_cached_snapshot(snapshot.clone());
164                snapshot
165            }
166            Err(_) => self.read_cached_snapshot(),
167        }
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[derive(Default)]
176    struct MutableAgentRegistry {
177        agents: RwLock<HashMap<String, AgentDefinition>>,
178    }
179
180    impl MutableAgentRegistry {
181        fn replace_ids(&self, ids: &[&str]) {
182            let mut map = HashMap::new();
183            for id in ids {
184                map.insert((*id).to_string(), AgentDefinition::new("gpt-4o-mini"));
185            }
186            match self.agents.write() {
187                Ok(mut guard) => *guard = map,
188                Err(poisoned) => *poisoned.into_inner() = map,
189            }
190        }
191    }
192
193    impl AgentRegistry for MutableAgentRegistry {
194        fn len(&self) -> usize {
195            self.snapshot().len()
196        }
197
198        fn get(&self, id: &str) -> Option<AgentDefinition> {
199            self.snapshot().get(id).cloned()
200        }
201
202        fn ids(&self) -> Vec<String> {
203            let mut ids: Vec<String> = self.snapshot().keys().cloned().collect();
204            ids.sort();
205            ids
206        }
207
208        fn snapshot(&self) -> HashMap<String, AgentDefinition> {
209            match self.agents.read() {
210                Ok(guard) => guard.clone(),
211                Err(poisoned) => poisoned.into_inner().clone(),
212            }
213        }
214    }
215
216    #[test]
217    fn composite_agent_registry_reads_live_updates_from_source_registries() {
218        let dynamic = Arc::new(MutableAgentRegistry::default());
219        dynamic.replace_ids(&["agent_a"]);
220
221        let mut static_registry = InMemoryAgentRegistry::new();
222        static_registry.upsert("agent_static", AgentDefinition::new("gpt-4o-mini"));
223
224        let composite = CompositeAgentRegistry::try_new(vec![
225            dynamic.clone() as Arc<dyn AgentRegistry>,
226            Arc::new(static_registry) as Arc<dyn AgentRegistry>,
227        ])
228        .expect("compose registries");
229
230        assert!(composite.ids().contains(&"agent_a".to_string()));
231        assert!(composite.ids().contains(&"agent_static".to_string()));
232
233        dynamic.replace_ids(&["agent_a", "agent_b"]);
234        let ids = composite.ids();
235        assert!(ids.contains(&"agent_a".to_string()));
236        assert!(ids.contains(&"agent_b".to_string()));
237        assert!(ids.contains(&"agent_static".to_string()));
238    }
239
240    #[test]
241    fn composite_agent_registry_keeps_last_good_snapshot_on_runtime_conflict() {
242        let reg_a = Arc::new(MutableAgentRegistry::default());
243        reg_a.replace_ids(&["agent_a"]);
244
245        let reg_b = Arc::new(MutableAgentRegistry::default());
246        reg_b.replace_ids(&["agent_b"]);
247
248        let composite = CompositeAgentRegistry::try_new(vec![
249            reg_a.clone() as Arc<dyn AgentRegistry>,
250            reg_b.clone() as Arc<dyn AgentRegistry>,
251        ])
252        .expect("compose registries");
253
254        let initial_ids = composite.ids();
255        assert_eq!(
256            initial_ids,
257            vec!["agent_a".to_string(), "agent_b".to_string()]
258        );
259
260        reg_b.replace_ids(&["agent_a"]);
261        assert_eq!(composite.ids(), initial_ids);
262        assert!(composite.get("agent_b").is_some());
263    }
264}
tirea_agentos/composition/registry/agent.rs

tirea_agentos/composition/registry/
agent.rs
