1use std::collections::HashMap;
22use std::fs;
23use std::path::PathBuf;
24use std::time::{Duration, SystemTime};
25
26use serde::{Deserialize, Serialize};
27
28use super::file_memory::{FileMemoryError, FileMemoryStore};
29
30const LEDGER: &str = ".memory.json";
32
33#[derive(Debug, Clone, Copy)]
35pub struct ForgetConfig {
36 pub ttl: Duration,
38 pub min_importance: f64,
40 pub weak_grace: Duration,
42}
43
44impl Default for ForgetConfig {
45 fn default() -> Self {
46 Self {
47 ttl: Duration::from_secs(30 * 24 * 3600), min_importance: 0.5,
49 weak_grace: Duration::from_secs(30 * 24 * 3600), }
51 }
52}
53
54impl ForgetConfig {
55 pub fn new() -> Self {
57 Self::default()
58 }
59
60 pub fn with_ttl(mut self, ttl: Duration) -> Self {
62 self.ttl = ttl;
63 self
64 }
65
66 pub fn with_min_importance(mut self, min_importance: f64) -> Self {
68 self.min_importance = min_importance;
69 self
70 }
71
72 pub fn with_weak_grace(mut self, weak_grace: Duration) -> Self {
74 self.weak_grace = weak_grace;
75 self
76 }
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
81struct MemoryMeta {
82 importance: f64,
84 created_at: u64,
86 last_access_at: u64,
88}
89
90impl MemoryMeta {
91 fn new(importance: f64, now: SystemTime) -> Self {
92 Self {
93 importance,
94 created_at: epoch_ms(now),
95 last_access_at: epoch_ms(now),
96 }
97 }
98}
99
100fn epoch_ms(t: SystemTime) -> u64 {
101 t.duration_since(std::time::UNIX_EPOCH)
102 .unwrap_or_default()
103 .as_millis() as u64
104}
105
106fn from_epoch_ms(ms: u64) -> SystemTime {
107 std::time::UNIX_EPOCH + Duration::from_millis(ms)
108}
109
110pub struct ForgettingMemory {
112 files: FileMemoryStore,
113 config: ForgetConfig,
114 meta: HashMap<String, MemoryMeta>,
115}
116
117impl std::fmt::Debug for ForgettingMemory {
118 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119 f.debug_struct("ForgettingMemory")
120 .field("root", &self.files.root())
121 .field("config", &self.config)
122 .field("live", &self.meta.len())
123 .finish_non_exhaustive()
124 }
125}
126
127impl ForgettingMemory {
128 pub fn new(files: FileMemoryStore, config: ForgetConfig) -> Result<Self, FileMemoryError> {
130 let meta = Self::load_ledger(files.root())?;
131 Ok(Self {
132 files,
133 config,
134 meta,
135 })
136 }
137
138 pub fn config(&self) -> ForgetConfig {
140 self.config
141 }
142
143 fn ledger_path(root: &std::path::Path) -> PathBuf {
144 root.join(LEDGER)
145 }
146
147 fn load_ledger(root: &std::path::Path) -> Result<HashMap<String, MemoryMeta>, FileMemoryError> {
148 let path = Self::ledger_path(root);
149 if !path.exists() {
150 return Ok(HashMap::new());
151 }
152 let raw = fs::read_to_string(&path)?;
153 serde_json::from_str(&raw).map_err(|e| FileMemoryError::UnsafeName {
154 name: LEDGER.to_string(),
155 reason: format!("malformed ledger: {e}"),
156 })
157 }
158
159 fn persist_ledger(&self) -> Result<(), FileMemoryError> {
160 let path = Self::ledger_path(self.files.root());
161 let raw =
162 serde_json::to_string_pretty(&self.meta).map_err(|e| FileMemoryError::UnsafeName {
163 name: LEDGER.to_string(),
164 reason: format!("serialize ledger: {e}"),
165 })?;
166 fs::write(path, raw)?;
167 Ok(())
168 }
169
170 pub fn remember(
173 &mut self,
174 name: &str,
175 content: &str,
176 importance: f64,
177 now: SystemTime,
178 ) -> Result<(), FileMemoryError> {
179 if !(0.0..=1.0).contains(&importance) {
180 return Err(FileMemoryError::UnsafeName {
181 name: name.to_string(),
182 reason: format!("importance {importance} outside [0, 1]"),
183 });
184 }
185 self.files.write(name, content)?;
186 self.meta
187 .insert(name.to_string(), MemoryMeta::new(importance, now));
188 self.persist_ledger()?;
189 Ok(())
190 }
191
192 pub fn recall(&mut self, name: &str, now: SystemTime) -> Result<String, FileMemoryError> {
194 let content = self.files.view(name)?;
195 if let Some(meta) = self.meta.get_mut(name) {
196 meta.last_access_at = epoch_ms(now);
197 self.persist_ledger()?;
198 }
199 Ok(content)
200 }
201
202 pub fn importance(&self, name: &str) -> Option<f64> {
204 self.meta.get(name).map(|m| m.importance)
205 }
206
207 pub fn forget(&mut self, name: &str) -> Result<(), FileMemoryError> {
209 self.files.delete(name)?;
210 self.meta.remove(name);
211 self.persist_ledger()?;
212 Ok(())
213 }
214
215 fn idle(&self, name: &str, now: SystemTime) -> Option<Duration> {
217 let meta = self.meta.get(name)?;
218 let last = from_epoch_ms(meta.last_access_at);
219 now.duration_since(last).ok()
220 }
221
222 pub fn should_forget(&self, name: &str, now: SystemTime) -> bool {
224 let Some(idle) = self.idle(name, now) else {
225 return true; };
227 if idle >= self.config.ttl {
228 return true;
229 }
230 let weak = self
231 .meta
232 .get(name)
233 .is_none_or(|m| m.importance < self.config.min_importance);
234 weak && idle >= self.config.weak_grace
235 }
236
237 pub fn consolidate(&mut self, now: SystemTime) -> Result<usize, FileMemoryError> {
241 let doomed: Vec<String> = self
242 .meta
243 .keys()
244 .filter(|name| self.should_forget(name, now))
245 .cloned()
246 .collect();
247 let count = doomed.len();
248 for name in doomed {
249 let _ = self.files.delete(&name); self.meta.remove(&name);
251 }
252 if count > 0 {
253 self.persist_ledger()?;
254 }
255 Ok(count)
256 }
257
258 pub fn live_at(&self, now: SystemTime) -> Vec<String> {
260 let mut live: Vec<String> = self
261 .meta
262 .keys()
263 .filter(|name| !self.should_forget(name, now))
264 .cloned()
265 .collect();
266 live.sort();
267 live
268 }
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274
275 const T0: std::time::SystemTime = std::time::UNIX_EPOCH;
276
277 fn store_and_config(
278 ttl_secs: u64,
279 min_imp: f64,
280 grace_secs: u64,
281 ) -> (tempfile::TempDir, FileMemoryStore, ForgetConfig) {
282 let dir = tempfile::tempdir().unwrap();
283 let files = FileMemoryStore::new(dir.path()).unwrap();
284 let config = ForgetConfig::new()
285 .with_ttl(Duration::from_secs(ttl_secs))
286 .with_min_importance(min_imp)
287 .with_weak_grace(Duration::from_secs(grace_secs));
288 (dir, files, config)
289 }
290
291 #[test]
292 fn remember_recall_persists_content() {
293 let (_d, files, cfg) = store_and_config(100, 0.4, 100);
294 let mut m = ForgettingMemory::new(files, cfg).unwrap();
295 m.remember("site", "the docs live under /docs", 0.9, T0)
296 .unwrap();
297 assert_eq!(m.recall("site", T0).unwrap(), "the docs live under /docs");
298 assert_eq!(m.importance("site"), Some(0.9));
299 }
300
301 #[test]
302 fn active_memory_survives_before_ttl() {
303 let (_d, files, cfg) = store_and_config(100, 0.4, 100);
304 let mut m = ForgettingMemory::new(files, cfg).unwrap();
305 m.remember("a", "x", 0.9, T0).unwrap();
306 let later = T0 + Duration::from_secs(90); assert!(!m.should_forget("a", later));
308 assert_eq!(m.consolidate(later).unwrap(), 0);
309 }
310
311 #[test]
312 fn idle_memory_forgotten_at_ttl() {
313 let (_d, files, cfg) = store_and_config(100, 0.4, 100);
314 let mut m = ForgettingMemory::new(files, cfg).unwrap();
315 m.remember("a", "x", 0.9, T0).unwrap();
316 let past = T0 + Duration::from_secs(101);
317 assert!(m.should_forget("a", past));
318 assert_eq!(m.consolidate(past).unwrap(), 1);
319 assert!(m.live_at(past).is_empty());
320 assert!(matches!(
322 m.files.view("a"),
323 Err(FileMemoryError::NotFound(_))
324 ));
325 }
326
327 #[test]
328 fn recall_refreshes_lifespan() {
329 let (_d, files, cfg) = store_and_config(100, 0.4, 100);
330 let mut m = ForgettingMemory::new(files, cfg).unwrap();
331 m.remember("a", "x", 0.9, T0).unwrap();
332 let _ = m.recall("a", T0 + Duration::from_secs(80)).unwrap();
334 let later = T0 + Duration::from_secs(150);
336 assert!(!m.should_forget("a", later));
337 }
338
339 #[test]
340 fn weak_memory_pruned_within_ttl_after_grace() {
341 let (_d, files, cfg) = store_and_config(1000, 0.5, 50);
342 let mut m = ForgettingMemory::new(files, cfg).unwrap();
343 m.remember("weak", "trivia", 0.1, T0).unwrap();
344 let later = T0 + Duration::from_secs(60);
346 assert!(m.should_forget("weak", later));
347 assert_eq!(m.consolidate(later).unwrap(), 1);
348 }
349
350 #[test]
351 fn importance_out_of_range_rejected() {
352 let (_d, files, cfg) = store_and_config(100, 0.4, 100);
353 let mut m = ForgettingMemory::new(files, cfg).unwrap();
354 assert!(m.remember("a", "x", 1.5, T0).is_err());
355 assert!(m.remember("b", "x", -0.1, T0).is_err());
356 }
357
358 #[test]
359 fn ledger_survives_restart() {
360 let (dir, files, cfg) = store_and_config(1000, 0.5, 1000);
361 {
362 let mut m = ForgettingMemory::new(files, cfg).unwrap();
363 m.remember("a", "content", 0.8, T0).unwrap();
364 }
365 let files2 = FileMemoryStore::new(dir.path()).unwrap();
367 let m2 = ForgettingMemory::new(files2, cfg).unwrap();
368 assert!(m2.live_at(T0).contains(&"a".to_string()));
369 assert_eq!(m2.importance("a"), Some(0.8));
370 }
371}