use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
const MAX_MEMORY_FILE_SIZE: u64 = 64 * 1024;
const DEFAULT_CHAR_LIMIT: usize = 4000;
pub struct MemoryStore {
path: PathBuf,
}
impl MemoryStore {
pub fn new(path: PathBuf) -> Self {
Self { path }
}
pub fn global() -> Self {
let dir = super::Config::config_dir();
Self::new(dir.join("memory.md"))
}
pub fn project(project_root: &Path) -> Self {
Self::new(project_root.join(".atomcode").join("memory.md"))
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn load(&self) -> Vec<String> {
let content = match fs::metadata(&self.path) {
Ok(meta) => {
if meta.len() > MAX_MEMORY_FILE_SIZE {
let bytes = fs::read(&self.path).unwrap_or_default();
let start = bytes.len().saturating_sub(MAX_MEMORY_FILE_SIZE as usize);
let safe_start = bytes[start..].iter()
.position(|&b| b == b'\n')
.map(|pos| start + pos + 1)
.unwrap_or(start);
String::from_utf8_lossy(&bytes[safe_start..]).to_string()
} else {
fs::read_to_string(&self.path).unwrap_or_default()
}
}
Err(_) => return Vec::new(),
};
content
.lines()
.filter_map(|line| {
let trimmed = line.trim();
if trimmed.starts_with("- ") {
Some(trimmed[2..].to_string())
} else {
None
}
})
.collect()
}
pub fn append(&self, content: &str) -> io::Result<()> {
if let Some(parent) = self.path.parent() {
fs::create_dir_all(parent)?;
}
let existing = fs::read_to_string(&self.path).unwrap_or_default();
let needs_newline = !existing.is_empty() && !existing.ends_with('\n');
let mut file = fs::OpenOptions::new()
.create(true)
.append(true)
.open(&self.path)?;
if needs_newline {
writeln!(file)?;
}
writeln!(file, "- {}", content.trim())
}
pub fn remove_matching(&self, keyword: &str) -> io::Result<Vec<String>> {
let content = fs::read_to_string(&self.path).unwrap_or_default();
let keyword_lower = keyword.to_lowercase();
let mut removed = Vec::new();
let mut kept = Vec::new();
for line in content.lines() {
let trimmed = line.trim();
if trimmed.starts_with("- ") && trimmed.to_lowercase().contains(&keyword_lower) {
removed.push(trimmed[2..].to_string());
} else {
kept.push(line.to_string());
}
}
if !removed.is_empty() {
let mut out = kept.join("\n");
if !out.is_empty() && !out.ends_with('\n') {
out.push('\n');
}
fs::write(&self.path, out)?;
}
Ok(removed)
}
pub fn find_matching(&self, keyword: &str) -> Vec<String> {
let keyword_lower = keyword.to_lowercase();
self.load()
.into_iter()
.filter(|entry| entry.to_lowercase().contains(&keyword_lower))
.collect()
}
pub fn merged_for_prompt(global: &MemoryStore, project: &MemoryStore, project_name: &str) -> String {
let global_entries = global.load();
let project_entries = project.load();
if global_entries.is_empty() && project_entries.is_empty() {
return String::new();
}
let mut result = String::from("=== MEMORY ===\nThe user has asked you to remember these facts and preferences:\n");
if !global_entries.is_empty() {
result.push_str("\n[Global]\n");
for entry in &global_entries {
result.push_str(&format!("- {}\n", entry));
}
}
if !project_entries.is_empty() {
result.push_str(&format!("\n[Project: {}]\n", project_name));
for entry in &project_entries {
result.push_str(&format!("- {}\n", entry));
}
}
if result.chars().count() > DEFAULT_CHAR_LIMIT {
let truncated: String = result.chars().take(DEFAULT_CHAR_LIMIT).collect();
format!("{}\n[...truncated, run /memory to review]", truncated)
} else {
result
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_append_creates_file() {
let dir = tempfile::tempdir().unwrap();
let store = MemoryStore::new(dir.path().join("sub").join("memory.md"));
store.append("test entry").unwrap();
let content = fs::read_to_string(store.path()).unwrap();
assert_eq!(content, "- test entry\n");
}
#[test]
fn test_append_to_existing() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
fs::write(&path, "- first\n").unwrap();
let store = MemoryStore::new(path);
store.append("second").unwrap();
let entries = store.load();
assert_eq!(entries, vec!["first", "second"]);
}
#[test]
fn test_load_skips_non_entries() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
fs::write(&path, "# Header\n\n- real entry\nnot an entry\n- another\n").unwrap();
let store = MemoryStore::new(path);
assert_eq!(store.load(), vec!["real entry", "another"]);
}
#[test]
fn test_load_empty_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
fs::write(&path, "").unwrap();
let store = MemoryStore::new(path);
assert!(store.load().is_empty());
}
#[test]
fn test_load_nonexistent() {
let store = MemoryStore::new(PathBuf::from("/nonexistent/memory.md"));
assert!(store.load().is_empty());
}
#[test]
fn test_remove_matching_case_insensitive() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
fs::write(&path, "- Use tabs\n- use spaces\n- pnpm only\n").unwrap();
let store = MemoryStore::new(path);
let removed = store.remove_matching("use").unwrap();
assert_eq!(removed, vec!["Use tabs", "use spaces"]);
assert_eq!(store.load(), vec!["pnpm only"]);
}
#[test]
fn test_remove_matching_no_match() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
fs::write(&path, "- keep this\n").unwrap();
let store = MemoryStore::new(path.clone());
let removed = store.remove_matching("nonexistent").unwrap();
assert!(removed.is_empty());
assert_eq!(fs::read_to_string(&path).unwrap(), "- keep this\n");
}
#[test]
fn test_merged_for_prompt_truncation() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("memory.md");
let long_entry = "x".repeat(5000);
fs::write(&path, format!("- {}\n", long_entry)).unwrap();
let store = MemoryStore::new(path);
let empty = MemoryStore::new(PathBuf::from("/none"));
let result = MemoryStore::merged_for_prompt(&store, &empty, "p");
assert!(result.contains("[...truncated"));
assert!(result.chars().count() < 5000);
}
}
