Skip to main content

mentedb_core/
agent.rs

1//! Agent Registry: tracks agents that participate in the memory system.
2
3use std::collections::HashMap;
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use serde::{Deserialize, Serialize};
7
8use crate::types::{AgentId, SpaceId, Timestamp};
9
10/// An agent that can read/write memories.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Agent {
13    /// Unique agent identifier.
14    pub id: AgentId,
15    /// Human-readable name.
16    pub name: String,
17    /// When the agent was registered (µs since epoch).
18    pub created_at: Timestamp,
19    /// Arbitrary key-value metadata.
20    pub metadata: HashMap<String, String>,
21    /// The default space this agent operates in.
22    pub default_space: SpaceId,
23}
24
25/// In-memory registry of all known agents.
26#[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    /// Create an empty registry.
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    /// Register a new agent and return it.
45    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    /// Look up an agent by ID.
58    pub fn get(&self, id: AgentId) -> Option<&Agent> {
59        self.agents.get(&id)
60    }
61
62    /// Remove an agent from the registry.
63    pub fn remove(&mut self, id: AgentId) {
64        self.agents.remove(&id);
65    }
66
67    /// List all registered agents.
68    pub fn list(&self) -> Vec<&Agent> {
69        self.agents.values().collect()
70    }
71
72    /// Update a single metadata key for an agent.
73    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}