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