Skip to main content

agent_code_lib/memory/
writer.rs

1//! Memory write discipline.
2//!
3//! Enforces the two-step write pattern:
4//! 1. Write the memory file with proper frontmatter
5//! 2. Update MEMORY.md index with a one-line pointer
6//!
7//! Prevents entropy by never dumping content into the index.
8
9use std::path::{Path, PathBuf};
10
11use super::types::{MemoryMeta, MemoryType};
12
13/// Maximum index line length.
14const MAX_INDEX_LINE_CHARS: usize = 150;
15
16/// Maximum index lines before truncation.
17const MAX_INDEX_LINES: usize = 200;
18
19/// Write a memory file and update the index atomically.
20///
21/// Returns the path of the written memory file.
22pub fn write_memory(
23    memory_dir: &Path,
24    filename: &str,
25    meta: &MemoryMeta,
26    content: &str,
27) -> Result<PathBuf, String> {
28    let _ = std::fs::create_dir_all(memory_dir);
29
30    // Step 1: Write the memory file with frontmatter.
31    let type_str = match &meta.memory_type {
32        Some(MemoryType::User) => "user",
33        Some(MemoryType::Feedback) => "feedback",
34        Some(MemoryType::Project) => "project",
35        Some(MemoryType::Reference) => "reference",
36        None => "user",
37    };
38
39    let file_content = format!(
40        "---\nname: {}\ndescription: {}\ntype: {}\n---\n\n{}",
41        meta.name, meta.description, type_str, content
42    );
43
44    let file_path = memory_dir.join(filename);
45    std::fs::write(&file_path, &file_content)
46        .map_err(|e| format!("Failed to write memory file: {e}"))?;
47
48    // Step 2: Update MEMORY.md index.
49    update_index(memory_dir, filename, &meta.name, &meta.description)?;
50
51    Ok(file_path)
52}
53
54/// Update the MEMORY.md index with a pointer to a memory file.
55/// If an entry for this filename already exists, replace it.
56fn update_index(
57    memory_dir: &Path,
58    filename: &str,
59    name: &str,
60    description: &str,
61) -> Result<(), String> {
62    let index_path = memory_dir.join("MEMORY.md");
63
64    let existing = std::fs::read_to_string(&index_path).unwrap_or_default();
65
66    // Build the new index line (under 150 chars).
67    let mut line = format!("- [{}]({}) — {}", name, filename, description);
68    if line.len() > MAX_INDEX_LINE_CHARS {
69        line.truncate(MAX_INDEX_LINE_CHARS - 3);
70        line.push_str("...");
71    }
72
73    // Replace existing entry for this filename, or append.
74    let mut lines: Vec<String> = existing
75        .lines()
76        .filter(|l| !l.contains(&format!("({})", filename)))
77        .map(|l| l.to_string())
78        .collect();
79
80    lines.push(line);
81
82    // Enforce max lines.
83    if lines.len() > MAX_INDEX_LINES {
84        lines.truncate(MAX_INDEX_LINES);
85    }
86
87    let new_index = lines.join("\n") + "\n";
88    std::fs::write(&index_path, new_index).map_err(|e| format!("Failed to update index: {e}"))?;
89
90    Ok(())
91}
92
93/// Remove a memory file and its index entry.
94pub fn delete_memory(memory_dir: &Path, filename: &str) -> Result<(), String> {
95    let file_path = memory_dir.join(filename);
96    if file_path.exists() {
97        std::fs::remove_file(&file_path).map_err(|e| format!("Failed to delete: {e}"))?;
98    }
99
100    // Remove from index.
101    let index_path = memory_dir.join("MEMORY.md");
102    if let Ok(existing) = std::fs::read_to_string(&index_path) {
103        let filtered: Vec<&str> = existing
104            .lines()
105            .filter(|l| !l.contains(&format!("({})", filename)))
106            .collect();
107        let _ = std::fs::write(&index_path, filtered.join("\n") + "\n");
108    }
109
110    Ok(())
111}
112
113/// Rebuild MEMORY.md from the actual files in the memory directory.
114/// Scans all .md files (except MEMORY.md itself), reads their frontmatter,
115/// and regenerates the index.
116pub fn rebuild_index(memory_dir: &Path) -> Result<(), String> {
117    let headers = super::scanner::scan_memory_files(memory_dir);
118    let index_path = memory_dir.join("MEMORY.md");
119
120    let mut lines = Vec::new();
121    for h in &headers {
122        let name = h
123            .meta
124            .as_ref()
125            .map(|m| m.name.as_str())
126            .unwrap_or(&h.filename);
127        let desc = h
128            .meta
129            .as_ref()
130            .map(|m| m.description.as_str())
131            .unwrap_or("");
132
133        let mut line = format!("- [{}]({}) — {}", name, h.filename, desc);
134        if line.len() > MAX_INDEX_LINE_CHARS {
135            line.truncate(MAX_INDEX_LINE_CHARS - 3);
136            line.push_str("...");
137        }
138        lines.push(line);
139    }
140
141    if lines.len() > MAX_INDEX_LINES {
142        lines.truncate(MAX_INDEX_LINES);
143    }
144
145    let content = lines.join("\n") + "\n";
146    std::fs::write(&index_path, content).map_err(|e| format!("Failed to write index: {e}"))?;
147
148    Ok(())
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    fn test_meta() -> MemoryMeta {
156        MemoryMeta {
157            name: "Test Memory".to_string(),
158            description: "A test memory file".to_string(),
159            memory_type: Some(MemoryType::User),
160        }
161    }
162
163    #[test]
164    fn test_write_memory_creates_file_and_index() {
165        let dir = tempfile::tempdir().unwrap();
166        let meta = test_meta();
167        let path = write_memory(dir.path(), "test.md", &meta, "Hello world").unwrap();
168
169        assert!(path.exists());
170        let content = std::fs::read_to_string(&path).unwrap();
171        assert!(content.contains("name: Test Memory"));
172        assert!(content.contains("type: user"));
173        assert!(content.contains("Hello world"));
174
175        // Index should exist and contain a pointer.
176        let index = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap();
177        assert!(index.contains("[Test Memory](test.md)"));
178    }
179
180    #[test]
181    fn test_write_memory_updates_existing_index_entry() {
182        let dir = tempfile::tempdir().unwrap();
183        let meta = test_meta();
184        write_memory(dir.path(), "test.md", &meta, "version 1").unwrap();
185
186        let meta2 = MemoryMeta {
187            name: "Updated".to_string(),
188            description: "Updated description".to_string(),
189            memory_type: Some(MemoryType::Feedback),
190        };
191        write_memory(dir.path(), "test.md", &meta2, "version 2").unwrap();
192
193        let index = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap();
194        // Should have only one entry for test.md (replaced, not duplicated).
195        assert_eq!(index.matches("test.md").count(), 1);
196        assert!(index.contains("[Updated](test.md)"));
197    }
198
199    #[test]
200    fn test_delete_memory() {
201        let dir = tempfile::tempdir().unwrap();
202        let meta = test_meta();
203        write_memory(dir.path(), "test.md", &meta, "content").unwrap();
204
205        delete_memory(dir.path(), "test.md").unwrap();
206
207        assert!(!dir.path().join("test.md").exists());
208        let index = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap();
209        assert!(!index.contains("test.md"));
210    }
211
212    #[test]
213    fn test_delete_nonexistent_memory() {
214        let dir = tempfile::tempdir().unwrap();
215        // Should not error even if file doesn't exist.
216        assert!(delete_memory(dir.path(), "nope.md").is_ok());
217    }
218
219    #[test]
220    fn test_rebuild_index() {
221        let dir = tempfile::tempdir().unwrap();
222        let meta = test_meta();
223        write_memory(dir.path(), "one.md", &meta, "first").unwrap();
224
225        let meta2 = MemoryMeta {
226            name: "Second".to_string(),
227            description: "Second file".to_string(),
228            memory_type: Some(MemoryType::Project),
229        };
230        write_memory(dir.path(), "two.md", &meta2, "second").unwrap();
231
232        // Corrupt the index.
233        std::fs::write(dir.path().join("MEMORY.md"), "garbage").unwrap();
234
235        // Rebuild should restore it.
236        rebuild_index(dir.path()).unwrap();
237        let index = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap();
238        assert!(index.contains("one.md"));
239        assert!(index.contains("two.md"));
240    }
241
242    #[test]
243    fn test_index_line_length_cap() {
244        let dir = tempfile::tempdir().unwrap();
245        let meta = MemoryMeta {
246            name: "A".repeat(200),
247            description: "B".repeat(200),
248            memory_type: Some(MemoryType::User),
249        };
250        write_memory(dir.path(), "long.md", &meta, "content").unwrap();
251
252        let index = std::fs::read_to_string(dir.path().join("MEMORY.md")).unwrap();
253        for line in index.lines() {
254            assert!(line.len() <= MAX_INDEX_LINE_CHARS + 3); // +3 for "..."
255        }
256    }
257}