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