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, Copy, PartialEq, Eq, Serialize, Deserialize)]
12pub enum Permission {
13 Read,
14 Write,
15 ReadWrite,
16 Admin,
17}
18
19impl Permission {
20 fn satisfies(self, required: Permission) -> bool {
22 matches!(
23 (self, required),
24 (Permission::Admin, _)
25 | (Permission::ReadWrite, Permission::Read)
26 | (Permission::ReadWrite, Permission::Write)
27 | (Permission::ReadWrite, Permission::ReadWrite)
28 | (Permission::Read, Permission::Read)
29 | (Permission::Write, Permission::Write)
30 )
31 }
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct AccessEntry {
37 pub agent_id: AgentId,
39 pub permission: Permission,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct MemorySpace {
46 pub id: SpaceId,
48 pub name: String,
50 pub owner: AgentId,
52 pub access_list: Vec<AccessEntry>,
54 pub created_at: Timestamp,
56 pub max_memories: Option<usize>,
58 pub current_count: usize,
60}
61
62#[derive(Debug, Default)]
64pub struct SpaceManager {
65 spaces: HashMap<SpaceId, MemorySpace>,
66}
67
68fn now_micros() -> Timestamp {
69 SystemTime::now()
70 .duration_since(UNIX_EPOCH)
71 .unwrap_or_default()
72 .as_micros() as Timestamp
73}
74
75impl SpaceManager {
76 pub fn new() -> Self {
78 Self::default()
79 }
80
81 pub fn create_space(&mut self, name: &str, owner: AgentId) -> MemorySpace {
83 let space = MemorySpace {
84 id: SpaceId::new(),
85 name: name.to_string(),
86 owner,
87 access_list: vec![AccessEntry {
88 agent_id: owner,
89 permission: Permission::Admin,
90 }],
91 created_at: now_micros(),
92 max_memories: None,
93 current_count: 0,
94 };
95 self.spaces.insert(space.id, space.clone());
96 space
97 }
98
99 pub fn get_space(&self, id: SpaceId) -> Option<&MemorySpace> {
101 self.spaces.get(&id)
102 }
103
104 pub fn delete_space(&mut self, id: SpaceId) {
106 self.spaces.remove(&id);
107 }
108
109 pub fn grant_access(&mut self, space: SpaceId, agent: AgentId, perm: Permission) {
111 if let Some(s) = self.spaces.get_mut(&space) {
112 if let Some(entry) = s.access_list.iter_mut().find(|e| e.agent_id == agent) {
114 entry.permission = perm;
115 } else {
116 s.access_list.push(AccessEntry {
117 agent_id: agent,
118 permission: perm,
119 });
120 }
121 }
122 }
123
124 pub fn revoke_access(&mut self, space: SpaceId, agent: AgentId) {
126 if let Some(s) = self.spaces.get_mut(&space) {
127 s.access_list.retain(|e| e.agent_id != agent);
128 }
129 }
130
131 pub fn check_access(&self, space: SpaceId, agent: AgentId, required: Permission) -> bool {
133 self.spaces.get(&space).is_some_and(|s| {
134 if s.owner == agent {
135 return true;
136 }
137 s.access_list
138 .iter()
139 .any(|e| e.agent_id == agent && e.permission.satisfies(required))
140 })
141 }
142
143 pub fn list_spaces_for_agent(&self, agent: AgentId) -> Vec<&MemorySpace> {
145 self.spaces
146 .values()
147 .filter(|s| s.owner == agent || s.access_list.iter().any(|e| e.agent_id == agent))
148 .collect()
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155
156 fn agent() -> AgentId {
157 AgentId::new()
158 }
159
160 #[test]
161 fn create_and_get() {
162 let mut mgr = SpaceManager::new();
163 let owner = agent();
164 let sp = mgr.create_space("test", owner);
165 assert_eq!(mgr.get_space(sp.id).unwrap().name, "test");
166 }
167
168 #[test]
169 fn owner_has_admin() {
170 let mut mgr = SpaceManager::new();
171 let owner = agent();
172 let sp = mgr.create_space("s", owner);
173 assert!(mgr.check_access(sp.id, owner, Permission::Admin));
174 }
175
176 #[test]
177 fn grant_and_check() {
178 let mut mgr = SpaceManager::new();
179 let owner = agent();
180 let reader = agent();
181 let sp = mgr.create_space("s", owner);
182 mgr.grant_access(sp.id, reader, Permission::Read);
183 assert!(mgr.check_access(sp.id, reader, Permission::Read));
184 assert!(!mgr.check_access(sp.id, reader, Permission::Write));
185 }
186
187 #[test]
188 fn revoke_access() {
189 let mut mgr = SpaceManager::new();
190 let owner = agent();
191 let a = agent();
192 let sp = mgr.create_space("s", owner);
193 mgr.grant_access(sp.id, a, Permission::ReadWrite);
194 mgr.revoke_access(sp.id, a);
195 assert!(!mgr.check_access(sp.id, a, Permission::Read));
196 }
197
198 #[test]
199 fn list_spaces_for_agent() {
200 let mut mgr = SpaceManager::new();
201 let o1 = agent();
202 let o2 = agent();
203 mgr.create_space("s1", o1);
204 mgr.create_space("s2", o2);
205 assert_eq!(mgr.list_spaces_for_agent(o1).len(), 1);
206 }
207}