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