1use chrono::Utc;
7
8use crate::memory::types::{MemoryEntry, ProtectionLevel};
9
10#[derive(Debug, Clone)]
24pub struct AutoProtector {
25 pub protection_low_access: u32,
27 pub protection_medium_access: u32,
29 pub protection_high_access: u32,
31 pub protection_medium_sessions: u32,
33 pub protection_high_sessions: u32,
35 pub demotion_stale_days: u32,
37}
38
39impl AutoProtector {
40 pub fn new(
42 low_access: u32,
43 medium_access: u32,
44 high_access: u32,
45 medium_sessions: u32,
46 high_sessions: u32,
47 demotion_stale_days: u32,
48 ) -> Self {
49 Self {
50 protection_low_access: low_access,
51 protection_medium_access: medium_access,
52 protection_high_access: high_access,
53 protection_medium_sessions: medium_sessions,
54 protection_high_sessions: high_sessions,
55 demotion_stale_days,
56 }
57 }
58
59 pub fn default_protector() -> Self {
61 Self::new(2, 3, 5, 2, 3, 30)
62 }
63
64 pub fn compute_protection(&self, entry: &MemoryEntry) -> ProtectionLevel {
68 if entry.memory_type.is_auto_protected() {
70 return ProtectionLevel::Permanent;
71 }
72
73 if entry.pinned {
75 return ProtectionLevel::Permanent;
76 }
77
78 if entry.user_corrected {
80 return ProtectionLevel::High;
81 }
82
83 let access_count = entry.access_count;
85 let session_span = entry.session_appearances;
86
87 if access_count >= self.protection_high_access
89 || session_span >= self.protection_high_sessions
90 {
91 return ProtectionLevel::High;
92 }
93
94 if access_count >= self.protection_medium_access
96 || session_span >= self.protection_medium_sessions
97 {
98 return ProtectionLevel::Medium;
99 }
100
101 if access_count >= self.protection_low_access {
103 return ProtectionLevel::Low;
104 }
105
106 ProtectionLevel::None
108 }
109
110 pub fn should_demote_protection(
115 &self,
116 entry: &MemoryEntry,
117 current: ProtectionLevel,
118 ) -> Option<ProtectionLevel> {
119 if entry.pinned || current == ProtectionLevel::Permanent {
121 return None;
122 }
123
124 let days_since_access = (Utc::now() - entry.accessed_at).num_days() as u32;
125 let stale = self.demotion_stale_days;
126
127 if current == ProtectionLevel::High
129 && days_since_access > stale
130 && entry.access_count < self.protection_medium_access
131 {
132 return Some(ProtectionLevel::Medium);
133 }
134
135 if current == ProtectionLevel::Medium && days_since_access > stale * 2 {
137 return Some(ProtectionLevel::Low);
138 }
139
140 if current == ProtectionLevel::Low && days_since_access > stale * 3 {
142 return Some(ProtectionLevel::None);
143 }
144
145 None
146 }
147
148 pub fn record_access(entry: &mut MemoryEntry, current_session_id: &str) {
152 entry.access_count += 1;
153 entry.accessed_at = Utc::now();
154
155 if !entry
157 .seen_in_sessions
158 .contains(¤t_session_id.to_string())
159 {
160 entry.session_appearances += 1;
161 entry.seen_in_sessions.push(current_session_id.to_string());
162 if entry.seen_in_sessions.len() > 100 {
164 entry.seen_in_sessions.remove(0);
165 }
166 }
167
168 let boosted = 0.5 + 0.5 * entry.decay_score;
170 entry.decay_score = entry.decay_score.max(boosted);
171 }
172}
173
174#[cfg(test)]
179mod tests {
180 use super::*;
181 use crate::memory::{MemoryEntry, MemoryTier, MemoryType};
182 use chrono::Duration;
183
184 fn make_entry_with_access(access_count: u32, sessions: u32) -> MemoryEntry {
185 let mut entry = make_base_entry();
186 entry.access_count = access_count;
187 entry.session_appearances = sessions;
188 entry.seen_in_sessions = (0..sessions).map(|i| format!("session-{}", i)).collect();
189 entry
190 }
191
192 fn make_base_entry() -> MemoryEntry {
193 MemoryEntry {
194 id: "test".to_string(),
195 memory_type: MemoryType::Fact,
196 tier: MemoryTier::Warm,
197 content: "test".to_string(),
198 content_hash: 0,
199 tags: vec![],
200 source: "test".to_string(),
201 session_id: None,
202 importance: 0.5,
203 pinned: false,
204 protection: ProtectionLevel::None,
205 auto_classified: false,
206 session_appearances: 0,
207 user_corrected: false,
208 seen_in_sessions: vec![],
209 created_at: Utc::now(),
210 accessed_at: Utc::now(),
211 modified_at: Utc::now(),
212 access_count: 0,
213 decay_score: 1.0,
214 compaction_level: 0,
215 compacted_from: vec![],
216 related_ids: vec![],
217 contradicts: None,
218 }
219 }
220
221 #[test]
222 fn test_protection_none_default() {
223 let protector = AutoProtector::default_protector();
224 let entry = make_entry_with_access(0, 0);
225 assert_eq!(protector.compute_protection(&entry), ProtectionLevel::None);
226 }
227
228 #[test]
229 fn test_protection_low() {
230 let protector = AutoProtector::default_protector();
231 let entry = make_entry_with_access(2, 0);
232 assert_eq!(protector.compute_protection(&entry), ProtectionLevel::Low);
233 }
234
235 #[test]
236 fn test_protection_medium_access() {
237 let protector = AutoProtector::default_protector();
238 let entry = make_entry_with_access(3, 0);
239 assert_eq!(
240 protector.compute_protection(&entry),
241 ProtectionLevel::Medium
242 );
243 }
244
245 #[test]
246 fn test_protection_medium_sessions() {
247 let protector = AutoProtector::default_protector();
248 let entry = make_entry_with_access(0, 2);
249 assert_eq!(
250 protector.compute_protection(&entry),
251 ProtectionLevel::Medium
252 );
253 }
254
255 #[test]
256 fn test_protection_high_access() {
257 let protector = AutoProtector::default_protector();
258 let entry = make_entry_with_access(5, 0);
259 assert_eq!(protector.compute_protection(&entry), ProtectionLevel::High);
260 }
261
262 #[test]
263 fn test_protection_high_sessions() {
264 let protector = AutoProtector::default_protector();
265 let entry = make_entry_with_access(0, 3);
266 assert_eq!(protector.compute_protection(&entry), ProtectionLevel::High);
267 }
268
269 #[test]
270 fn test_protection_permanent_for_profile() {
271 let protector = AutoProtector::default_protector();
272 let mut entry = make_base_entry();
273 entry.memory_type = MemoryType::UserProfile;
274 assert_eq!(
275 protector.compute_protection(&entry),
276 ProtectionLevel::Permanent
277 );
278 }
279
280 #[test]
281 fn test_protection_permanent_for_preference() {
282 let protector = AutoProtector::default_protector();
283 let mut entry = make_base_entry();
284 entry.memory_type = MemoryType::Preference;
285 assert_eq!(
286 protector.compute_protection(&entry),
287 ProtectionLevel::Permanent
288 );
289 }
290
291 #[test]
292 fn test_protection_user_correction() {
293 let protector = AutoProtector::default_protector();
294 let mut entry = make_base_entry();
295 entry.user_corrected = true;
296 assert_eq!(protector.compute_protection(&entry), ProtectionLevel::High);
297 }
298
299 #[test]
300 fn test_protection_pinned() {
301 let protector = AutoProtector::default_protector();
302 let mut entry = make_base_entry();
303 entry.pinned = true;
304 assert_eq!(
305 protector.compute_protection(&entry),
306 ProtectionLevel::Permanent
307 );
308 }
309
310 #[test]
311 fn test_demote_high_to_medium() {
312 let protector = AutoProtector::default_protector();
313 let mut entry = make_entry_with_access(2, 0); entry.accessed_at = Utc::now() - Duration::days(35); let result = protector.should_demote_protection(&entry, ProtectionLevel::High);
316 assert_eq!(result, Some(ProtectionLevel::Medium));
317 }
318
319 #[test]
320 fn test_demote_medium_to_low() {
321 let protector = AutoProtector::default_protector();
322 let mut entry = make_entry_with_access(3, 1);
323 entry.accessed_at = Utc::now() - Duration::days(65); let result = protector.should_demote_protection(&entry, ProtectionLevel::Medium);
325 assert_eq!(result, Some(ProtectionLevel::Low));
326 }
327
328 #[test]
329 fn test_demote_low_to_none() {
330 let protector = AutoProtector::default_protector();
331 let mut entry = make_entry_with_access(2, 0);
332 entry.accessed_at = Utc::now() - Duration::days(95); let result = protector.should_demote_protection(&entry, ProtectionLevel::Low);
334 assert_eq!(result, Some(ProtectionLevel::None));
335 }
336
337 #[test]
338 fn test_no_demote_permanent() {
339 let protector = AutoProtector::default_protector();
340 let mut entry = make_base_entry();
341 entry.accessed_at = Utc::now() - Duration::days(365);
342 let result = protector.should_demote_protection(&entry, ProtectionLevel::Permanent);
343 assert_eq!(result, None);
344 }
345
346 #[test]
347 fn test_no_demote_pinned() {
348 let protector = AutoProtector::default_protector();
349 let mut entry = make_base_entry();
350 entry.pinned = true;
351 entry.accessed_at = Utc::now() - Duration::days(365);
352 let result = protector.should_demote_protection(&entry, ProtectionLevel::High);
353 assert_eq!(result, None);
354 }
355
356 #[test]
357 fn test_record_access() {
358 let mut entry = make_base_entry();
359 entry.decay_score = 0.2;
360 AutoProtector::record_access(&mut entry, "session-1");
361
362 assert_eq!(entry.access_count, 1);
363 assert_eq!(entry.session_appearances, 1);
364 assert!(entry.seen_in_sessions.contains(&"session-1".to_string()));
365 assert!(entry.decay_score > 0.2, "Should recover decay on access");
366 }
367
368 #[test]
369 fn test_record_access_dedup_session() {
370 let mut entry = make_base_entry();
371 AutoProtector::record_access(&mut entry, "session-1");
372 AutoProtector::record_access(&mut entry, "session-1");
373 assert_eq!(entry.access_count, 2);
374 assert_eq!(
375 entry.session_appearances, 1,
376 "Same session should not increment appearances"
377 );
378 }
379}