roma_memory/

lib.rs

//! Layered memory system (L0-L4).
//!
//! Provides [`FileMemoryStore`] as the default file-backed implementation.
//! The [`MemoryStore`] trait, [`MemoryError`], [`MemoryLevel`], and
//! [`NullMemoryStore`] are defined in `roma-core` and re-exported here
//! for backward compatibility.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use async_trait::async_trait;
use roma_core::safe_resolve;
use tokio::fs;
use tracing::{debug, warn};

pub use roma_core::{MemoryError, MemoryLevel, MemoryStore, NullMemoryStore, PatchError};

// ---------------------------------------------------------------------------
// Validators
// ---------------------------------------------------------------------------

/// Structural validator for a memory level.
pub trait MemoryValidator: Send + Sync {
    fn validate(&self, content: &str) -> Result<(), MemoryError>;
}

/// L1 validator: enforce a maximum line count.
pub struct L1MaxLines {
    pub max: usize,
}

impl MemoryValidator for L1MaxLines {
    fn validate(&self, content: &str) -> Result<(), MemoryError> {
        let lines = content.lines().count();
        if lines > self.max {
            return Err(MemoryError::ValidationFailed {
                level: MemoryLevel::L1,
                reason: format!("{lines} lines exceed maximum of {}", self.max),
            });
        }
        Ok(())
    }
}

/// L2 validator: enforce required markdown sections exist.
pub struct L2SectionGuard {
    pub required: Vec<String>,
}

impl MemoryValidator for L2SectionGuard {
    fn validate(&self, content: &str) -> Result<(), MemoryError> {
        for section in &self.required {
            let header = format!("# {section}");
            let header2 = format!("## {section}");
            if !content.contains(&header) && !content.contains(&header2) {
                return Err(MemoryError::ValidationFailed {
                    level: MemoryLevel::L2,
                    reason: format!("missing required section: {section}"),
                });
            }
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// FileMemoryStore
// ---------------------------------------------------------------------------

/// File-backed memory store with per-level validators.
///
/// Directory layout:
/// ```text
/// base_dir/L0/...
/// base_dir/L1/...
/// ...
/// ```
///
/// Paths passed to read/write/patch/delete are relative to a level directory.
/// The format is `<level_prefix>/<name>`, e.g. `"L1/key_info.md"`. The level
/// prefix is automatically stripped to resolve the filesystem path.
pub struct FileMemoryStore {
    base_dir: PathBuf,
    validators: HashMap<MemoryLevel, Box<dyn MemoryValidator>>,
}

impl FileMemoryStore {
    /// Create a new file-backed store rooted at `base_dir`.
    /// Level subdirectories are created lazily on first write.
    pub fn new(base_dir: impl Into<PathBuf>) -> Self {
        Self {
            base_dir: base_dir.into(),
            validators: HashMap::new(),
        }
    }

    /// Attach a validator for a specific memory level.
    pub fn with_validator(
        mut self,
        level: MemoryLevel,
        validator: Box<dyn MemoryValidator>,
    ) -> Self {
        self.validators.insert(level, validator);
        self
    }

    /// Parse a logical path like `"L1/key_info.md"` into `(MemoryLevel, relative_path)`.
    ///
    /// Performs only lexical validation: rejects `..` components and
    /// absolute sub-paths. Full filesystem-boundary resolution (including
    /// symlink following) happens in [`Self::safe_fs_path`].
    fn parse_path(&self, path: &str) -> Result<(MemoryLevel, PathBuf), MemoryError> {
        let (level, rest) = path
            .split_once('/')
            .ok_or_else(|| MemoryError::NotFound(format!("invalid memory path: {path}")))?;
        let level = match level {
            "L0" => MemoryLevel::L0,
            "L1" => MemoryLevel::L1,
            "L2" => MemoryLevel::L2,
            "L3" => MemoryLevel::L3,
            "L4" => MemoryLevel::L4,
            _ => return Err(MemoryError::NotFound(format!("unknown level: {level}"))),
        };
        let relative = PathBuf::from(rest);
        // Lexical rejection of traversal and absolute paths. Symlink
        // escape is caught later by `safe_fs_path`.
        for comp in relative.components() {
            match comp {
                std::path::Component::ParentDir => {
                    return Err(MemoryError::PathDenied(format!(
                        "path traversal rejected: {path}"
                    )));
                }
                std::path::Component::RootDir | std::path::Component::Prefix(_) => {
                    return Err(MemoryError::PathDenied(format!(
                        "absolute sub-path rejected: {path}"
                    )));
                }
                _ => {}
            }
        }
        Ok((level, relative))
    }

    /// Ensure the level directory exists, then resolve `relative` to an
    /// absolute path that is guaranteed to stay inside
    /// `base_dir/<level>/`, following symlinks.
    async fn safe_fs_path(
        &self,
        level: MemoryLevel,
        relative: &Path,
    ) -> Result<PathBuf, MemoryError> {
        let level_dir = self.base_dir.join(level.dir_name());
        if !level_dir.exists() {
            fs::create_dir_all(&level_dir).await?;
        }
        let resolved = safe_resolve(&level_dir, relative)?;
        Ok(resolved)
    }
}

#[async_trait]
impl MemoryStore for FileMemoryStore {
    async fn read(&self, path: &str) -> Result<String, MemoryError> {
        let (level, relative) = self.parse_path(path)?;
        let level_dir = self.base_dir.join(level.dir_name());
        // If the level directory doesn't exist yet, nothing can be read
        // there — short-circuit to NotFound without creating the dir.
        if !level_dir.exists() {
            return Err(MemoryError::NotFound(path.to_string()));
        }
        let full = self.safe_fs_path(level, &relative).await?;
        fs::read_to_string(&full).await.map_err(|e| {
            if e.kind() == std::io::ErrorKind::NotFound {
                MemoryError::NotFound(path.to_string())
            } else {
                MemoryError::Io(e)
            }
        })
    }

    async fn write(&self, path: &str, content: &str) -> Result<(), MemoryError> {
        let (level, relative) = self.parse_path(path)?;
        self.validate(level, content)?;
        let full = self.safe_fs_path(level, &relative).await?;
        if let Some(parent) = full.parent() {
            fs::create_dir_all(parent).await?;
        }
        debug!(path = %full.display(), "writing memory file");
        fs::write(&full, content).await?;
        Ok(())
    }

    async fn patch(&self, path: &str, old: &str, new: &str) -> Result<(), MemoryError> {
        let content = self.read(path).await?;
        let count = content.matches(old).count();
        if count == 0 {
            return Err(MemoryError::Patch(PatchError::NotFound));
        }
        if count > 1 {
            return Err(MemoryError::Patch(PatchError::NotUnique { count }));
        }
        let patched = content.replacen(old, new, 1);
        self.write(path, &patched).await
    }

    async fn delete(&self, path: &str) -> Result<(), MemoryError> {
        let (level, relative) = self.parse_path(path)?;
        let level_dir = self.base_dir.join(level.dir_name());
        if !level_dir.exists() {
            // Nothing to delete; stay idempotent.
            return Ok(());
        }
        let full = self.safe_fs_path(level, &relative).await?;
        match fs::remove_file(&full).await {
            Ok(()) => {
                debug!(path = %full.display(), "deleted memory file");
                Ok(())
            }
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                warn!(path = %full.display(), "delete called on non-existent file");
                Ok(())
            }
            Err(e) => Err(MemoryError::Io(e)),
        }
    }

    async fn list(&self, level: MemoryLevel) -> Result<Vec<String>, MemoryError> {
        let dir = self.base_dir.join(level.dir_name());
        let mut result = Vec::new();
        // Walk directory tree recursively using an explicit stack.
        let mut stack = vec![(dir, String::new())];
        while let Some((current_dir, prefix)) = stack.pop() {
            let mut entries = match fs::read_dir(&current_dir).await {
                Ok(rd) => rd,
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue,
                Err(e) => return Err(MemoryError::Io(e)),
            };
            while let Ok(Some(entry)) = entries.next_entry().await {
                let name = entry.file_name();
                let name_str = match name.to_str() {
                    Some(s) => s,
                    None => continue,
                };
                let path = entry.path();
                if path.is_dir() {
                    let sub_prefix = if prefix.is_empty() {
                        format!("{}/", name_str)
                    } else {
                        format!("{prefix}{name_str}/")
                    };
                    stack.push((path, sub_prefix));
                } else {
                    let logical = if prefix.is_empty() {
                        format!("{}/{}", level.dir_name(), name_str)
                    } else {
                        format!("{}/{}{}", level.dir_name(), prefix, name_str)
                    };
                    result.push(logical);
                }
            }
        }
        result.sort();
        Ok(result)
    }

    fn validate(&self, level: MemoryLevel, content: &str) -> Result<(), MemoryError> {
        if let Some(v) = self.validators.get(&level) {
            v.validate(content)?;
        }
        Ok(())
    }

    async fn read_root(&self, name: &str) -> Result<String, MemoryError> {
        // Reject traversal attempts.
        if name.contains("..") || name.contains('/') || name.contains('\\') {
            return Err(MemoryError::PathDenied(format!(
                "root file name must be plain: {name}"
            )));
        }
        // Resolve symlinks and verify the canonical path stays within base_dir.
        let resolved = safe_resolve(&self.base_dir, Path::new(name))?;
        fs::read_to_string(&resolved).await.map_err(|e| {
            if e.kind() == std::io::ErrorKind::NotFound {
                MemoryError::NotFound(name.to_string())
            } else {
                MemoryError::Io(e)
            }
        })
    }
}

// ---------------------------------------------------------------------------
// Forgetter trait
// ---------------------------------------------------------------------------

/// Strategy for pruning/consolidating long-term memory when it grows too large.
///
/// Implementations decide *which* files to merge, *how* to merge them, and
/// what to discard. The trait is intentionally minimal — the "how" of merging
/// (LLM summarization, simple concatenation, etc.) is left to the implementer.
#[async_trait]
pub trait Forgetter: Send + Sync {
    /// Run consolidation if the L3 file count exceeds `threshold`.
    /// Returns the number of files removed.
    async fn maybe_consolidate(
        &self,
        store: &Arc<dyn MemoryStore>,
        threshold: usize,
    ) -> Result<usize, MemoryError>;
}

/// A no-op forgetter that never consolidates. Useful as a default.
pub struct NullForgetter;

#[async_trait]
impl Forgetter for NullForgetter {
    async fn maybe_consolidate(
        &self,
        _store: &Arc<dyn MemoryStore>,
        _threshold: usize,
    ) -> Result<usize, MemoryError> {
        Ok(0)
    }
}

/// Simple forgetter that merges the oldest L3 files when the count exceeds
/// the threshold. Files are concatenated into a single `_consolidated_<n>.md`
/// and the originals are deleted.
pub struct SimpleForgetter;

#[async_trait]
impl Forgetter for SimpleForgetter {
    async fn maybe_consolidate(
        &self,
        store: &Arc<dyn MemoryStore>,
        threshold: usize,
    ) -> Result<usize, MemoryError> {
        let files = store.list(MemoryLevel::L3).await?;
        if files.len() <= threshold {
            return Ok(0);
        }

        let excess = files.len() - threshold;
        let to_merge = &files[..excess];
        let mut merged = String::from("# Consolidated Memory\n\n");
        let mut read_ok: Vec<&str> = Vec::new();
        for path in to_merge {
            match store.read(path).await {
                Ok(content) => {
                    let name = path.rsplit('/').next().unwrap_or(path);
                    merged.push_str(&format!("## {name}\n\n{content}\n\n"));
                    read_ok.push(path);
                }
                Err(e) => {
                    warn!(path, error = %e, "failed to read file during consolidation, skipping");
                }
            }
        }

        let n = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default();
        let consolidated_path = format!("L3/_consolidated_{}{}.md", n.as_secs(), n.subsec_millis());
        store.write(&consolidated_path, &merged).await?;

        let mut removed = 0;
        for path in read_ok {
            match store.delete(path).await {
                Ok(()) | Err(MemoryError::NotFound(_)) => removed += 1,
                Err(e) => {
                    warn!(path, error = %e, "failed to delete consolidated file");
                }
            }
        }

        debug!(
            consolidated = consolidated_path,
            removed, "L3 consolidation complete"
        );
        Ok(removed)
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn make_store(dir: &TempDir) -> FileMemoryStore {
        FileMemoryStore::new(dir.path())
            .with_validator(MemoryLevel::L1, Box::new(L1MaxLines { max: 30 }))
            .with_validator(
                MemoryLevel::L2,
                Box::new(L2SectionGuard {
                    required: vec!["Findings".into()],
                }),
            )
    }

    #[tokio::test]
    async fn write_and_read_roundtrip() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.write("L0/rules.md", "be helpful").await.unwrap();
        let content = store.read("L0/rules.md").await.unwrap();
        assert_eq!(content, "be helpful");
    }

    #[tokio::test]
    async fn read_missing_returns_not_found() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let err = store.read("L0/absent.md").await.unwrap_err();
        assert!(matches!(err, MemoryError::NotFound(_)));
    }

    #[tokio::test]
    async fn patch_unique_replaces_content() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.write("L0/rules.md", "old value here").await.unwrap();
        store.patch("L0/rules.md", "old", "new").await.unwrap();
        assert_eq!(store.read("L0/rules.md").await.unwrap(), "new value here");
    }

    #[tokio::test]
    async fn patch_not_unique_returns_error() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store
            .write("L0/rules.md", "foo and foo again")
            .await
            .unwrap();
        let err = store.patch("L0/rules.md", "foo", "bar").await.unwrap_err();
        assert!(matches!(
            err,
            MemoryError::Patch(PatchError::NotUnique { count: 2 })
        ));
    }

    #[tokio::test]
    async fn patch_not_found_returns_error() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.write("L0/rules.md", "hello").await.unwrap();
        let err = store
            .patch("L0/rules.md", "absent", "new")
            .await
            .unwrap_err();
        assert!(matches!(err, MemoryError::Patch(PatchError::NotFound)));
    }

    #[tokio::test]
    async fn delete_removes_file() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.write("L0/rules.md", "temp").await.unwrap();
        store.delete("L0/rules.md").await.unwrap();
        assert!(matches!(
            store.read("L0/rules.md").await.unwrap_err(),
            MemoryError::NotFound(_)
        ));
    }

    #[tokio::test]
    async fn delete_nonexistent_is_ok() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.delete("L0/ghost.md").await.unwrap();
    }

    #[tokio::test]
    async fn list_returns_files_in_level() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store.write("L0/rules.md", "r").await.unwrap();
        store.write("L0/sop.md", "s").await.unwrap();
        store.write("L1/key.md", "k").await.unwrap();
        let l0 = store.list(MemoryLevel::L0).await.unwrap();
        assert_eq!(l0, vec!["L0/rules.md", "L0/sop.md"]);
        let l1 = store.list(MemoryLevel::L1).await.unwrap();
        assert_eq!(l1, vec!["L1/key.md"]);
        let l2 = store.list(MemoryLevel::L2).await.unwrap();
        assert!(l2.is_empty());
    }

    #[tokio::test]
    async fn l1_validator_rejects_excess_lines() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let long = (0..31)
            .map(|i| format!("line {i}"))
            .collect::<Vec<_>>()
            .join("\n");
        let err = store.write("L1/key.md", &long).await.unwrap_err();
        assert!(matches!(
            err,
            MemoryError::ValidationFailed {
                level: MemoryLevel::L1,
                ..
            }
        ));
    }

    #[tokio::test]
    async fn l1_validator_accepts_at_limit() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let exact = (0..30)
            .map(|i| format!("line {i}"))
            .collect::<Vec<_>>()
            .join("\n");
        store.write("L1/key.md", &exact).await.unwrap();
    }

    #[tokio::test]
    async fn l2_validator_rejects_missing_section() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let err = store
            .write("L2/session.md", "# No findings here")
            .await
            .unwrap_err();
        assert!(matches!(
            err,
            MemoryError::ValidationFailed {
                level: MemoryLevel::L2,
                ..
            }
        ));
    }

    #[tokio::test]
    async fn l2_validator_accepts_with_section() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        store
            .write("L2/session.md", "# Findings\nsome insight")
            .await
            .unwrap();
    }

    #[test]
    fn l1_max_lines_validator_direct() {
        let v = L1MaxLines { max: 2 };
        assert!(v.validate("a\nb").is_ok());
        assert!(v.validate("a\nb\nc").is_err());
    }

    #[test]
    fn l2_section_guard_validator_direct() {
        let v = L2SectionGuard {
            required: vec!["Summary".into()],
        };
        assert!(v.validate("# Summary\nhello").is_ok());
        assert!(v.validate("## Summary\nhello").is_ok());
        assert!(v.validate("# No summary here").is_err());
    }

    #[tokio::test]
    async fn path_traversal_is_rejected() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let err = store.read("L1/../../etc/passwd").await.unwrap_err();
        assert!(matches!(err, MemoryError::PathDenied(_)));
        let err = store.write("L0/../secret", "data").await.unwrap_err();
        assert!(matches!(err, MemoryError::PathDenied(_)));
    }

    #[tokio::test]
    async fn absolute_subpath_is_rejected() {
        let dir = TempDir::new().unwrap();
        let store = make_store(&dir);
        let err = store.read("L0//etc/passwd").await.unwrap_err();
        // "L0//etc/passwd" → level "L0", rest "/etc/passwd" → absolute → PathDenied
        assert!(matches!(err, MemoryError::PathDenied(_)));
        let err = store.write("L0//tmp/evil", "x").await.unwrap_err();
        assert!(matches!(err, MemoryError::PathDenied(_)));
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn symlink_escape_is_rejected() {
        // Create a memory base with an L0 directory that contains a
        // symlink pointing outside the base. Reads via that symlink
        // must be rejected.
        let outside = TempDir::new().unwrap();
        std::fs::write(outside.path().join("target.md"), "secret").unwrap();

        let dir = TempDir::new().unwrap();
        let l0 = dir.path().join("L0");
        std::fs::create_dir_all(&l0).unwrap();
        std::os::unix::fs::symlink(outside.path().join("target.md"), l0.join("link.md")).unwrap();

        let store = make_store(&dir);
        let err = store.read("L0/link.md").await.unwrap_err();
        assert!(
            matches!(err, MemoryError::PathDenied(_)),
            "expected PathDenied, got {err:?}"
        );
    }

    #[tokio::test]
    async fn l3_crud_roundtrip() {
        let dir = TempDir::new().unwrap();
        let store = FileMemoryStore::new(dir.path());

        store
            .write("L3/debug_rust_build.md", "# SOP\n1. cargo check")
            .await
            .unwrap();
        let content = store.read("L3/debug_rust_build.md").await.unwrap();
        assert!(content.contains("cargo check"));

        store
            .patch("L3/debug_rust_build.md", "cargo check", "cargo clippy")
            .await
            .unwrap();
        let patched = store.read("L3/debug_rust_build.md").await.unwrap();
        assert!(patched.contains("cargo clippy"));

        store.delete("L3/debug_rust_build.md").await.unwrap();
        assert!(matches!(
            store.read("L3/debug_rust_build.md").await.unwrap_err(),
            MemoryError::NotFound(_)
        ));
    }

    #[tokio::test]
    async fn list_isolation_between_levels() {
        let dir = TempDir::new().unwrap();
        let store = FileMemoryStore::new(dir.path());

        store.write("L0/rules.md", "r").await.unwrap();
        store.write("L3/sop1.md", "s1").await.unwrap();
        store.write("L3/sop2.md", "s2").await.unwrap();

        let l0 = store.list(MemoryLevel::L0).await.unwrap();
        let l3 = store.list(MemoryLevel::L3).await.unwrap();

        assert_eq!(l0, vec!["L0/rules.md"]);
        assert_eq!(l3, vec!["L3/sop1.md", "L3/sop2.md"]);
    }

    #[tokio::test]
    async fn write_creates_subdirectories() {
        let dir = TempDir::new().unwrap();
        let store = FileMemoryStore::new(dir.path());

        store
            .write("L3/debugging/rust_build.md", "steps")
            .await
            .unwrap();

        let files = store.list(MemoryLevel::L3).await.unwrap();
        assert_eq!(files, vec!["L3/debugging/rust_build.md"]);

        let content = store.read("L3/debugging/rust_build.md").await.unwrap();
        assert_eq!(content, "steps");
    }

    #[tokio::test]
    async fn read_root_rejects_traversal() {
        let dir = TempDir::new().unwrap();
        let store = FileMemoryStore::new(dir.path());

        let err = store.read_root("../etc/passwd").await.unwrap_err();
        assert!(matches!(err, MemoryError::PathDenied(_)));

        let err = store.read_root("sub/file").await.unwrap_err();
        assert!(matches!(err, MemoryError::PathDenied(_)));
    }

    #[tokio::test]
    async fn read_root_roundtrip() {
        let dir = TempDir::new().unwrap();
        std::fs::write(dir.path().join("README.md"), "project info").unwrap();

        let store = FileMemoryStore::new(dir.path());
        let content = store.read_root("README.md").await.unwrap();
        assert_eq!(content, "project info");
    }
}