1use std::collections::HashMap;
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use serde::{Deserialize, Serialize};
7
8use crate::types::{AgentId, SpaceId, Timestamp};
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Agent {
13 pub id: AgentId,
15 pub name: String,
17 pub created_at: Timestamp,
19 pub metadata: HashMap<String, String>,
21 pub default_space: SpaceId,
23}
24
25#[derive(Debug, Default)]
27pub struct AgentRegistry {
28 agents: HashMap<AgentId, Agent>,
29}
30
31fn now_micros() -> Timestamp {
32 SystemTime::now()
33 .duration_since(UNIX_EPOCH)
34 .unwrap_or_default()
35 .as_micros() as Timestamp
36}
37
38impl AgentRegistry {
39 pub fn new() -> Self {
41 Self::default()
42 }
43
44 pub fn register(&mut self, name: &str) -> Agent {
46 let agent = Agent {
47 id: AgentId::new(),
48 name: name.to_string(),
49 created_at: now_micros(),
50 metadata: HashMap::new(),
51 default_space: SpaceId::nil(),
52 };
53 self.agents.insert(agent.id, agent.clone());
54 agent
55 }
56
57 pub fn get(&self, id: AgentId) -> Option<&Agent> {
59 self.agents.get(&id)
60 }
61
62 pub fn remove(&mut self, id: AgentId) {
64 self.agents.remove(&id);
65 }
66
67 pub fn list(&self) -> Vec<&Agent> {
69 self.agents.values().collect()
70 }
71
72 pub fn update_metadata(&mut self, id: AgentId, key: &str, value: &str) {
74 if let Some(agent) = self.agents.get_mut(&id) {
75 agent.metadata.insert(key.to_string(), value.to_string());
76 }
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 #[test]
85 fn register_and_get() {
86 let mut reg = AgentRegistry::new();
87 let a = reg.register("alice");
88 assert_eq!(reg.get(a.id).unwrap().name, "alice");
89 }
90
91 #[test]
92 fn remove_agent() {
93 let mut reg = AgentRegistry::new();
94 let a = reg.register("bob");
95 reg.remove(a.id);
96 assert!(reg.get(a.id).is_none());
97 }
98
99 #[test]
100 fn list_agents() {
101 let mut reg = AgentRegistry::new();
102 reg.register("a1");
103 reg.register("a2");
104 assert_eq!(reg.list().len(), 2);
105 }
106
107 #[test]
108 fn update_metadata() {
109 let mut reg = AgentRegistry::new();
110 let a = reg.register("meta");
111 reg.update_metadata(a.id, "role", "planner");
112 assert_eq!(
113 reg.get(a.id).unwrap().metadata.get("role").unwrap(),
114 "planner"
115 );
116 }
117}