telos_agent/knowledge/memory/
index.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3
4use crate::knowledge::memory::format::{MemoryCategory, MemoryEntry, MemoryFormat, MemoryStatus};
5
6mod maintenance;
7mod query;
8
9pub use maintenance::{
10    MemoryMaintenanceAction, MemoryMaintenanceActionKind, MemoryMaintenancePolicy,
11    MemoryMaintenanceReport,
12};
13pub use query::{MemoryQuery, MemorySort};
14
15/// Result of inserting or merging a memory entry.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum UpsertOutcome {
18    Created,
19    Updated,
20}
21
22/// Persistent store for agent memories, organized by category into subdirectories.
23/// Maintains a MEMORY.md index file.
24pub struct MemoryStore {
25    root: PathBuf,
26    /// name → full path on disk
27    index: HashMap<String, PathBuf>,
28    /// name → parsed entry (in-memory cache to avoid O(n) disk reads)
29    cache: HashMap<String, MemoryEntry>,
30}
31
32impl MemoryStore {
33    /// Open or create a memory store at the given root directory.
34    pub fn new(root: PathBuf) -> Self {
35        std::fs::create_dir_all(&root).ok();
36        let mut store = Self { root, index: HashMap::new(), cache: HashMap::new() };
37        store.rebuild_index();
38        store
39    }
40
41    /// Re-scan the root directory and rebuild the in-memory index and cache.
42    fn rebuild_index(&mut self) {
43        self.index.clear();
44        self.cache.clear();
45        let subdirs = ["scripts", "commands", "patterns", "facts", "workflows"];
46        for subdir in subdirs {
47            let dir = self.root.join(subdir);
48            if !dir.exists() {
49                continue;
50            }
51            if let Ok(entries) = std::fs::read_dir(&dir) {
52                for entry in entries.flatten() {
53                    let path = entry.path();
54                    if path.extension().is_none_or(|e| e != "md") {
55                        continue;
56                    }
57                    if let Ok(content) = std::fs::read_to_string(&path)
58                        && let Some(mem) = MemoryFormat::parse(&content)
59                    {
60                        self.cache.insert(mem.name.clone(), mem.clone());
61                        self.index.insert(mem.name.clone(), path);
62                    }
63                }
64            }
65        }
66    }
67
68    /// Write a memory entry to disk and update the index.
69    pub fn write(&mut self, entry: MemoryEntry) -> std::io::Result<()> {
70        let filename = sanitize_name(&entry.name);
71        let subdir = category_subdir(&entry.category);
72        let dir = self.root.join(subdir);
73        std::fs::create_dir_all(&dir)?;
74        let path = dir.join(format!("{}.md", filename));
75        if let Some(old_path) = self.index.get(&entry.name)
76            && old_path != &path
77        {
78            match std::fs::remove_file(old_path) {
79                Ok(()) => {}
80                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
81                Err(err) => return Err(err),
82            }
83        }
84        let content = MemoryFormat::serialize(&entry);
85        std::fs::write(&path, content)?;
86        self.cache.insert(entry.name.clone(), entry.clone());
87        self.index.insert(entry.name.clone(), path);
88        self.write_index_md()?;
89        Ok(())
90    }
91
92    /// Write a new memory or merge into an existing one with the same name or description.
93    pub fn upsert(&mut self, entry: MemoryEntry) -> std::io::Result<UpsertOutcome> {
94        let existing_name = if self.index.contains_key(&entry.name) {
95            Some(entry.name.clone())
96        } else {
97            self.index.keys().find_map(|name| {
98                let existing = self.read(name)?;
99                if existing.description.eq_ignore_ascii_case(&entry.description) {
100                    Some(existing.name)
101                } else {
102                    None
103                }
104            })
105        };
106
107        if let Some(name) = existing_name
108            && let Some(mut existing) = self.read(&name)
109        {
110            existing.description =
111                if entry.description.is_empty() { existing.description } else { entry.description };
112            existing.category = entry.category;
113            existing.tags = merge_strings(existing.tags, entry.tags);
114            existing.updated = entry.updated;
115            existing.status = entry.status;
116            existing.confidence = entry.confidence.or(existing.confidence);
117            existing.related = merge_strings(existing.related, entry.related);
118            existing.source_session = entry.source_session.or(existing.source_session);
119            existing.body = merge_body(&existing.body, &entry.body);
120            self.write(existing)?;
121            return Ok(UpsertOutcome::Updated);
122        }
123
124        self.write(entry)?;
125        Ok(UpsertOutcome::Created)
126    }
127
128    /// Read a memory entry by name.
129    pub fn read(&self, name: &str) -> Option<MemoryEntry> {
130        self.cache.get(name).cloned()
131    }
132
133    /// Return all memory names.
134    pub fn list(&self) -> Vec<String> {
135        let mut names: Vec<String> = self.index.keys().cloned().collect();
136        names.sort();
137        names
138    }
139
140    /// Update the status of a memory entry.
141    pub fn update_status(&mut self, name: &str, status: MemoryStatus) -> std::io::Result<()> {
142        if let Some(mut entry) = self.read(name) {
143            entry.status = status;
144            self.write(entry)?;
145        }
146        Ok(())
147    }
148
149    /// Record that a memory was explicitly used.
150    pub fn record_use(&mut self, name: &str) -> std::io::Result<()> {
151        if let Some(mut entry) = self.read(name) {
152            entry.times_used = entry.times_used.saturating_add(1);
153            entry.updated = unix_timestamp();
154            self.write(entry)?;
155        }
156        Ok(())
157    }
158
159    /// Move a memory to the _archived directory (never delete).
160    pub fn archive(&mut self, name: &str) -> std::io::Result<()> {
161        self.cache.remove(name);
162        if let Some(path) = self.index.remove(name) {
163            let archive_dir = self.root.join("_archived");
164            std::fs::create_dir_all(&archive_dir)?;
165            let dest = archive_dir.join(path.file_name().unwrap());
166            std::fs::rename(&path, &dest)?;
167            self.write_index_md()?;
168        }
169        Ok(())
170    }
171
172    /// Write the MEMORY.md index file.
173    fn write_index_md(&self) -> std::io::Result<()> {
174        let mut lines = Vec::new();
175        let mut names: Vec<&String> = self.index.keys().collect();
176        names.sort();
177        for name in names {
178            if let Some(entry) = self.read(name) {
179                let fname = sanitize_name(name);
180                let subdir = category_subdir(&entry.category);
181                lines.push(format!(
182                    "- [{}]({}/{}.md) — {}",
183                    entry.name, subdir, fname, entry.description
184                ));
185            }
186        }
187        std::fs::write(self.root.join("MEMORY.md"), lines.join("\n"))
188    }
189}
190
191/// Map category to storage subdirectory.
192fn category_subdir(cat: &MemoryCategory) -> &'static str {
193    match cat {
194        MemoryCategory::Script => "scripts",
195        MemoryCategory::Command => "commands",
196        MemoryCategory::Pattern => "patterns",
197        MemoryCategory::Fact => "facts",
198        MemoryCategory::Workflow => "workflows",
199    }
200}
201
202/// Sanitize a name for use as a filename.
203fn sanitize_name(name: &str) -> String {
204    name.chars()
205        .map(|c| if c.is_alphanumeric() || c == '-' || c == '_' { c } else { '-' })
206        .collect::<String>()
207        .to_lowercase()
208}
209
210fn merge_strings(mut existing: Vec<String>, incoming: Vec<String>) -> Vec<String> {
211    for item in incoming {
212        if !existing.iter().any(|e| e.eq_ignore_ascii_case(&item)) {
213            existing.push(item);
214        }
215    }
216    existing
217}
218
219fn merge_body(existing: &str, incoming: &str) -> String {
220    if incoming.trim().is_empty() || existing.contains(incoming.trim()) {
221        return existing.to_string();
222    }
223    if existing.trim().is_empty() {
224        return incoming.to_string();
225    }
226    format!("{}\n\n---\n\n{}", existing.trim_end(), incoming.trim_start())
227}
228
229pub fn unix_timestamp() -> String {
230    std::time::SystemTime::now()
231        .duration_since(std::time::UNIX_EPOCH)
232        .map(|d| d.as_secs().to_string())
233        .unwrap_or_default()
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use crate::knowledge::memory::format::{MemoryCategory, MemoryEntry, MemoryStatus};
240
241    fn test_entry(name: &str, desc: &str, cat: MemoryCategory) -> MemoryEntry {
242        MemoryEntry {
243            name: name.into(),
244            description: desc.into(),
245            category: cat,
246            tags: vec!["test".into()],
247            created: "2026-06-18".into(),
248            updated: "2026-06-18".into(),
249            status: MemoryStatus::Working,
250            times_used: 1,
251            confidence: None,
252            related: vec![],
253            source_session: None,
254            body: "Test body".into(),
255        }
256    }
257
258    #[test]
259    fn write_and_read_roundtrip() {
260        let dir = tempfile::tempdir().unwrap();
261        let mut store = MemoryStore::new(dir.path().to_path_buf());
262
263        store.write(test_entry("test-mem", "A test memory", MemoryCategory::Fact)).unwrap();
264        let entry = store.read("test-mem").unwrap();
265        assert_eq!(entry.name, "test-mem");
266        assert_eq!(entry.category, MemoryCategory::Fact);
267    }
268
269    #[test]
270    fn search_finds_by_tag() {
271        let dir = tempfile::tempdir().unwrap();
272        let mut store = MemoryStore::new(dir.path().to_path_buf());
273
274        let mut entry = test_entry("deploy-script", "Deploy script", MemoryCategory::Script);
275        entry.tags = vec!["deploy".into(), "staging".into()];
276        store.write(entry).unwrap();
277
278        let results = store.search("staging");
279        assert_eq!(results.len(), 1);
280        assert_eq!(results[0].name, "deploy-script");
281    }
282
283    #[test]
284    fn list_returns_sorted_names() {
285        let dir = tempfile::tempdir().unwrap();
286        let mut store = MemoryStore::new(dir.path().to_path_buf());
287
288        store.write(test_entry("b", "B", MemoryCategory::Fact)).unwrap();
289        store.write(test_entry("a", "A", MemoryCategory::Fact)).unwrap();
290
291        let names = store.list();
292        assert_eq!(names, vec!["a", "b"]);
293    }
294
295    #[test]
296    fn archive_moves_to_archived_dir() {
297        let dir = tempfile::tempdir().unwrap();
298        let mut store = MemoryStore::new(dir.path().to_path_buf());
299
300        store.write(test_entry("old-mem", "Old", MemoryCategory::Fact)).unwrap();
301        store.archive("old-mem").unwrap();
302
303        assert!(store.read("old-mem").is_none());
304        assert!(dir.path().join("_archived").exists());
305    }
306
307    #[test]
308    fn update_status_changes_entry() {
309        let dir = tempfile::tempdir().unwrap();
310        let mut store = MemoryStore::new(dir.path().to_path_buf());
311
312        store.write(test_entry("flaky-script", "Flaky", MemoryCategory::Script)).unwrap();
313        store.update_status("flaky-script", MemoryStatus::NeedsFix).unwrap();
314
315        let entry = store.read("flaky-script").unwrap();
316        assert_eq!(entry.status, MemoryStatus::NeedsFix);
317    }
318
319    #[test]
320    fn top_by_usage_returns_most_used() {
321        let dir = tempfile::tempdir().unwrap();
322        let mut store = MemoryStore::new(dir.path().to_path_buf());
323
324        let mut high = test_entry("high-usage", "High", MemoryCategory::Fact);
325        high.times_used = 10;
326        let mut low = test_entry("low-usage", "Low", MemoryCategory::Fact);
327        low.times_used = 1;
328        store.write(high).unwrap();
329        store.write(low).unwrap();
330
331        let top = store.top_by_usage(1);
332        assert_eq!(top.len(), 1);
333        assert_eq!(top[0].name, "high-usage");
334    }
335
336    #[test]
337    fn upsert_updates_existing_memory_by_name() {
338        let dir = tempfile::tempdir().unwrap();
339        let mut store = MemoryStore::new(dir.path().to_path_buf());
340
341        store.write(test_entry("same", "Old", MemoryCategory::Fact)).unwrap();
342        let mut replacement = test_entry("same", "New", MemoryCategory::Command);
343        replacement.tags = vec!["new".into()];
344        replacement.body = "New body".into();
345
346        let outcome = store.upsert(replacement).unwrap();
347        assert_eq!(outcome, UpsertOutcome::Updated);
348
349        let entry = store.read("same").unwrap();
350        assert_eq!(entry.description, "New");
351        assert_eq!(entry.category, MemoryCategory::Command);
352        assert!(entry.tags.contains(&"new".to_string()));
353        assert!(entry.body.contains("Test body"));
354        assert!(entry.body.contains("New body"));
355    }
356
357    #[test]
358    fn category_change_survives_reopening_store() {
359        let dir = tempfile::tempdir().unwrap();
360        let mut store = MemoryStore::new(dir.path().to_path_buf());
361
362        store.write(test_entry("same", "Old", MemoryCategory::Fact)).unwrap();
363        let mut replacement = test_entry("same", "New", MemoryCategory::Command);
364        replacement.body = "New body".into();
365        store.upsert(replacement).unwrap();
366
367        let reopened = MemoryStore::new(dir.path().to_path_buf());
368        let entry = reopened.read("same").unwrap();
369        assert_eq!(entry.category, MemoryCategory::Command);
370        assert_eq!(entry.description, "New");
371        assert!(entry.body.contains("New body"));
372    }
373
374    #[test]
375    fn query_filters_status_and_tags() {
376        let dir = tempfile::tempdir().unwrap();
377        let mut store = MemoryStore::new(dir.path().to_path_buf());
378
379        let mut fix = test_entry("fix", "Fix", MemoryCategory::Fact);
380        fix.status = MemoryStatus::NeedsFix;
381        fix.tags = vec!["error".into()];
382        store.write(fix).unwrap();
383        store.write(test_entry("ok", "Ok", MemoryCategory::Fact)).unwrap();
384
385        let results = store.query(MemoryQuery {
386            status: Some(MemoryStatus::NeedsFix),
387            tags: vec!["error".into()],
388            include_body: false,
389            ..MemoryQuery::default()
390        });
391
392        assert_eq!(results.len(), 1);
393        assert_eq!(results[0].name, "fix");
394        assert!(results[0].body.is_empty());
395    }
396
397    #[test]
398    fn record_use_increments_times_used() {
399        let dir = tempfile::tempdir().unwrap();
400        let mut store = MemoryStore::new(dir.path().to_path_buf());
401
402        store.write(test_entry("used", "Used", MemoryCategory::Fact)).unwrap();
403        store.record_use("used").unwrap();
404
405        let entry = store.read("used").unwrap();
406        assert_eq!(entry.times_used, 2);
407    }
408
409    #[test]
410    fn maintenance_report_archives_deprecated_and_excess_auto_learned_commands() {
411        let dir = tempfile::tempdir().unwrap();
412        let mut store = MemoryStore::new(dir.path().to_path_buf());
413
414        let mut old_auto = test_entry("old-auto", "Old auto command", MemoryCategory::Command);
415        old_auto.tags = vec!["auto-learned".into(), "bash".into()];
416        old_auto.updated = "100".into();
417        old_auto.times_used = 0;
418
419        let mut new_auto = test_entry("new-auto", "New auto command", MemoryCategory::Command);
420        new_auto.tags = vec!["auto-learned".into(), "bash".into()];
421        new_auto.updated = "200".into();
422        new_auto.times_used = 0;
423
424        let mut used_auto = test_entry("used-auto", "Used auto command", MemoryCategory::Command);
425        used_auto.tags = vec!["auto-learned".into(), "bash".into()];
426        used_auto.updated = "50".into();
427        used_auto.times_used = 3;
428
429        let mut deprecated = test_entry("deprecated", "Deprecated fact", MemoryCategory::Fact);
430        deprecated.status = MemoryStatus::Deprecated;
431
432        store.write(old_auto).unwrap();
433        store.write(new_auto).unwrap();
434        store.write(used_auto).unwrap();
435        store.write(deprecated).unwrap();
436        store.write(test_entry("manual", "Manual fact", MemoryCategory::Fact)).unwrap();
437
438        let report = store.maintenance_report(&MemoryMaintenancePolicy {
439            max_auto_learned_commands: Some(2),
440            archive_deprecated: true,
441            ..MemoryMaintenancePolicy::default()
442        });
443
444        let names: Vec<&str> = report.actions.iter().map(|action| action.name.as_str()).collect();
445        assert_eq!(names, vec!["deprecated", "old-auto"]);
446        assert!(report.actions.iter().any(|action| action.reason.contains("deprecated")));
447        assert!(report.actions.iter().any(|action| action.reason.contains("auto-learned")));
448    }
449
450    #[test]
451    fn maintenance_apply_moves_candidates_to_archive() {
452        let dir = tempfile::tempdir().unwrap();
453        let mut store = MemoryStore::new(dir.path().to_path_buf());
454
455        let mut old_auto = test_entry("old-auto", "Old auto command", MemoryCategory::Command);
456        old_auto.tags = vec!["auto-learned".into(), "bash".into()];
457        old_auto.updated = "100".into();
458        old_auto.times_used = 0;
459
460        let mut new_auto = test_entry("new-auto", "New auto command", MemoryCategory::Command);
461        new_auto.tags = vec!["auto-learned".into(), "bash".into()];
462        new_auto.updated = "200".into();
463        new_auto.times_used = 0;
464
465        store.write(old_auto).unwrap();
466        store.write(new_auto).unwrap();
467
468        let report = store
469            .apply_maintenance(&MemoryMaintenancePolicy {
470                max_auto_learned_commands: Some(1),
471                archive_deprecated: false,
472                ..MemoryMaintenancePolicy::default()
473            })
474            .unwrap();
475
476        assert!(report.applied);
477        assert_eq!(report.archived_count, 1);
478        assert!(store.read("old-auto").is_none());
479        assert!(store.read("new-auto").is_some());
480        assert!(dir.path().join("_archived").join("old-auto.md").exists());
481    }
482}
telos_agent/knowledge/memory/index.rs

telos_agent/knowledge/memory/
index.rs
