1use std::collections::HashMap;
4use std::sync::{Arc, RwLock};
5
6use crate::agent::Agent;
7
8pub type AgentRef = Arc<tokio::sync::Mutex<Agent>>;
12
13#[derive(Clone)]
20pub struct AgentRegistry {
21 agents: Arc<RwLock<HashMap<String, AgentRef>>>,
22}
23
24impl AgentRegistry {
25 #[must_use]
27 pub fn new() -> Self {
28 Self {
29 agents: Arc::new(RwLock::new(HashMap::new())),
30 }
31 }
32
33 pub fn register(&self, name: impl Into<String>, agent: Agent) -> AgentRef {
37 let agent_ref: AgentRef = Arc::new(tokio::sync::Mutex::new(agent));
38 self.agents
39 .write()
40 .unwrap_or_else(std::sync::PoisonError::into_inner)
41 .insert(name.into(), Arc::clone(&agent_ref));
42 agent_ref
43 }
44
45 #[must_use]
47 pub fn get(&self, name: &str) -> Option<AgentRef> {
48 self.agents
49 .read()
50 .unwrap_or_else(std::sync::PoisonError::into_inner)
51 .get(name)
52 .cloned()
53 }
54
55 pub fn remove(&self, name: &str) -> Option<AgentRef> {
57 self.agents
58 .write()
59 .unwrap_or_else(std::sync::PoisonError::into_inner)
60 .remove(name)
61 }
62
63 #[must_use]
65 pub fn names(&self) -> Vec<String> {
66 self.agents
67 .read()
68 .unwrap_or_else(std::sync::PoisonError::into_inner)
69 .keys()
70 .cloned()
71 .collect()
72 }
73
74 #[must_use]
76 pub fn len(&self) -> usize {
77 self.agents
78 .read()
79 .unwrap_or_else(std::sync::PoisonError::into_inner)
80 .len()
81 }
82
83 #[must_use]
85 pub fn is_empty(&self) -> bool {
86 self.agents
87 .read()
88 .unwrap_or_else(std::sync::PoisonError::into_inner)
89 .is_empty()
90 }
91}
92
93impl Default for AgentRegistry {
94 fn default() -> Self {
95 Self::new()
96 }
97}