walrus_daemon/hook/system/memory/

mod.rs

//! Built-in memory — file-per-entry storage at `{config_dir}/memory/`.
//!
//! [`Memory`] manages individual entry files under `entries/`, a curated
//! `MEMORY.md` overview, and session summaries under `sessions/`. Entry
//! recall uses BM25 ranking. All I/O goes through the [`Storage`] trait
//! for testability.

use crate::hook::system::MemoryConfig;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::RwLock,
};
use wcore::model::{Message, Role};

pub mod bm25;
pub mod entry;
pub mod storage;
pub(crate) mod tool;

use entry::MemoryEntry;
use storage::Storage;

const MEMORY_PROMPT: &str = include_str!("../../../../prompts/memory.md");

const DEFAULT_SOUL: &str = include_str!("../../../../prompts/walrus.md");

pub struct Memory {
    storage: Box<dyn Storage>,
    entries: RwLock<HashMap<String, MemoryEntry>>,
    index: RwLock<String>,
    soul: RwLock<String>,
    index_path: PathBuf,
    soul_path: PathBuf,
    entries_dir: PathBuf,
    sessions_dir: PathBuf,
    config: MemoryConfig,
}

impl Memory {
    /// Open (or create) memory storage at the given directory.
    ///
    /// `config_dir` is the parent config directory where `Walrus.md` lives.
    /// `dir` is the memory-specific subdirectory (`{config_dir}/memory/`).
    pub fn open(dir: PathBuf, config: MemoryConfig, storage: Box<dyn Storage>) -> Self {
        let entries_dir = dir.join("entries");
        let sessions_dir = dir.join("sessions");
        let index_path = dir.join("MEMORY.md");
        // Walrus.md lives in the parent config dir, not inside memory/
        let soul_path = dir
            .parent()
            .map(|p| p.join("Walrus.md"))
            .unwrap_or_else(|| dir.join("Walrus.md"));

        storage.create_dir_all(&entries_dir).ok();
        storage.create_dir_all(&sessions_dir).ok();

        // Seed Walrus.md if it doesn't exist
        if !storage.exists(&soul_path) {
            storage.write(&soul_path, DEFAULT_SOUL).ok();
        }

        let soul_content = storage
            .read(&soul_path)
            .unwrap_or_else(|_| DEFAULT_SOUL.to_owned());

        let mem = Self {
            storage,
            entries: RwLock::new(HashMap::new()),
            index: RwLock::new(String::new()),
            soul: RwLock::new(soul_content),
            index_path,
            soul_path,
            entries_dir,
            sessions_dir,
            config,
        };

        mem.migrate_legacy(&dir);
        mem.load_entries();
        mem.load_index();
        mem
    }

    /// Load all entry files from the entries directory.
    fn load_entries(&self) {
        let paths = match self.storage.list(&self.entries_dir) {
            Ok(p) => p,
            Err(_) => return,
        };

        let mut entries = self.entries.write().unwrap();
        for path in paths {
            if path.extension().and_then(|e| e.to_str()) != Some("md") {
                continue;
            }
            let raw = match self.storage.read(&path) {
                Ok(r) => r,
                Err(_) => continue,
            };
            match MemoryEntry::parse(path, &raw) {
                Ok(entry) => {
                    entries.insert(entry.name.clone(), entry);
                }
                Err(e) => {
                    tracing::warn!("failed to parse memory entry: {e}");
                }
            }
        }
    }

    /// Load MEMORY.md index content.
    fn load_index(&self) {
        if let Ok(content) = self.storage.read(&self.index_path) {
            *self.index.write().unwrap() = content;
        }
    }

    /// BM25-ranked recall over all entries.
    pub fn recall(&self, query: &str, limit: usize) -> String {
        let entries = self.entries.read().unwrap();
        if entries.is_empty() {
            return "no memories found".to_owned();
        }

        // Single vector for both scoring and result lookup — avoids HashMap
        // iteration order aliasing between separate `.values()` calls.
        let entry_vec: Vec<&MemoryEntry> = entries.values().collect();
        let docs: Vec<(usize, String)> = entry_vec
            .iter()
            .enumerate()
            .map(|(i, e)| (i, e.search_text()))
            .collect();
        let doc_refs: Vec<(usize, &str)> = docs.iter().map(|(i, s)| (*i, s.as_str())).collect();

        let results = bm25::score(&doc_refs, query, limit);
        if results.is_empty() {
            return "no memories found".to_owned();
        }

        results
            .iter()
            .map(|(idx, _score)| {
                let e = &entry_vec[*idx];
                format!("## {}\n{}\n\n{}", e.name, e.description, e.content)
            })
            .collect::<Vec<_>>()
            .join("\n---\n")
    }

    /// Create or update a memory entry.
    pub fn remember(&self, name: String, description: String, content: String) -> String {
        let entry = MemoryEntry::new(name.clone(), description, content, &self.entries_dir);
        if let Err(e) = entry.save(self.storage.as_ref()) {
            return format!("failed to save entry: {e}");
        }
        self.entries.write().unwrap().insert(name.clone(), entry);
        format!("remembered: {name}")
    }

    /// Delete a memory entry by name.
    pub fn forget(&self, name: &str) -> String {
        let mut entries = self.entries.write().unwrap();
        match entries.remove(name) {
            Some(entry) => {
                if let Err(e) = entry.delete(self.storage.as_ref()) {
                    tracing::warn!("failed to delete entry file: {e}");
                }
                format!("forgot: {name}")
            }
            None => format!("no entry named: {name}"),
        }
    }

    /// Overwrite MEMORY.md (the curated overview).
    pub fn write_index(&self, content: &str) -> String {
        if let Err(e) = self.storage.write(&self.index_path, content) {
            return format!("failed to write MEMORY.md: {e}");
        }
        *self.index.write().unwrap() = content.to_owned();
        "MEMORY.md updated".to_owned()
    }

    /// Overwrite Walrus.md (the soul/identity file). Gated by `soul_editable`.
    pub fn write_soul(&self, content: &str) -> String {
        if !self.config.soul_editable {
            return "soul editing is disabled in config".to_owned();
        }
        if let Err(e) = self.storage.write(&self.soul_path, content) {
            return format!("failed to write Walrus.md: {e}");
        }
        *self.soul.write().unwrap() = content.to_owned();
        "Walrus.md updated".to_owned()
    }

    /// Return the soul content for system prompt injection.
    pub fn build_soul(&self) -> String {
        self.soul.read().unwrap().clone()
    }

    /// Build system prompt block from MEMORY.md content.
    pub fn build_prompt(&self) -> String {
        let index = self.index.read().unwrap();
        if index.is_empty() {
            return format!("\n\n{MEMORY_PROMPT}");
        }
        format!("\n\n<memory>\n{}\n</memory>\n\n{MEMORY_PROMPT}", *index)
    }

    /// Auto-recall from last user message, injected before each turn.
    pub fn before_run(&self, history: &[Message]) -> Vec<Message> {
        let last_user = history
            .iter()
            .rev()
            .find(|m| m.role == Role::User && !m.content.is_empty());

        let Some(msg) = last_user else {
            return Vec::new();
        };

        let query: String = msg
            .content
            .split_whitespace()
            .take(8)
            .collect::<Vec<_>>()
            .join(" ");

        if query.is_empty() {
            return Vec::new();
        }

        let limit = self.config.recall_limit;
        let result = self.recall(&query, limit);
        if result == "no memories found" {
            return Vec::new();
        }

        vec![Message {
            role: Role::User,
            content: format!("<recall>\n{result}\n</recall>"),
            auto_injected: true,
            ..Default::default()
        }]
    }

    /// Save a session summary after compaction.
    pub fn after_compact(&self, agent: &str, summary: &str) {
        let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
        let filename = format!("{agent}_{timestamp}.md");
        let path = self.sessions_dir.join(filename);
        if let Err(e) = self.storage.write(&path, summary) {
            tracing::warn!("failed to save session summary: {e}");
        }
    }

    /// Migrate legacy files (memory.md, user.md, facts.toml) to entry format.
    fn migrate_legacy(&self, dir: &Path) {
        // Only migrate if entries dir is empty.
        let existing = self.storage.list(&self.entries_dir).unwrap_or_default();
        if !existing.is_empty() {
            return;
        }

        let memory_path = dir.join("memory.md");
        let user_path = dir.join("user.md");
        let facts_path = dir.join("facts.toml");

        let has_legacy = self.storage.exists(&memory_path)
            || self.storage.exists(&user_path)
            || self.storage.exists(&facts_path);

        if !has_legacy {
            return;
        }

        // memory.md → split by double-newline into entries + seed MEMORY.md
        if let Ok(content) = self.storage.read(&memory_path)
            && !content.trim().is_empty()
        {
            self.storage.write(&self.index_path, &content).ok();

            for (i, chunk) in content.split("\n\n").enumerate() {
                let chunk = chunk.trim();
                if chunk.is_empty() {
                    continue;
                }
                let name = format!("migrated-memory-{}", i + 1);
                let entry = MemoryEntry::new(
                    name,
                    "Migrated from memory.md".to_owned(),
                    chunk.to_owned(),
                    &self.entries_dir,
                );
                entry.save(self.storage.as_ref()).ok();
            }
            self.storage
                .rename(&memory_path, &dir.join("memory.md.bak"))
                .ok();
        }

        // user.md → single entry
        if let Ok(content) = self.storage.read(&user_path)
            && !content.trim().is_empty()
        {
            let entry = MemoryEntry::new(
                "user-profile".to_owned(),
                "User profile migrated from user.md".to_owned(),
                content,
                &self.entries_dir,
            );
            entry.save(self.storage.as_ref()).ok();
            self.storage
                .rename(&user_path, &dir.join("user.md.bak"))
                .ok();
        }

        // facts.toml → single entry
        if let Ok(content) = self.storage.read(&facts_path)
            && !content.trim().is_empty()
        {
            let entry = MemoryEntry::new(
                "known-facts".to_owned(),
                "Known facts migrated from facts.toml".to_owned(),
                content,
                &self.entries_dir,
            );
            entry.save(self.storage.as_ref()).ok();
            self.storage
                .rename(&facts_path, &dir.join("facts.toml.bak"))
                .ok();
        }
    }
}