mcp-memory/server/src/state.rs

use crate::models::*;
use crate::store::{Store, read_json_file, write_json_atomic};
use crate::search::MemoryIndex;
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::PathBuf;
use std::sync::RwLock;
use std::time::{Duration, SystemTime};

pub struct MemoryState {
    pub base_dir: PathBuf,
    pub master_path: PathBuf,
    pub session_graph: RwLock<KnowledgeGraph>,
    pub master_cache: RwLock<(KnowledgeGraph, SystemTime)>,
    pub search_index: RwLock<MemoryIndex>,
    pub ledger: Store<Vec<CodeChange>>,
    pub sticky: Store<Vec<StickyNote>>,
    pub tasks: Store<Vec<Task>>,
    pub snippets: Store<Vec<Snippet>>,
    pub adrs: Store<Vec<Adr>>,
    pub prefs: Store<HashMap<String, Preference>>,
    pub error_fixes: Store<Vec<ErrorFix>>,
    pub pinned_files: Store<Vec<PinnedFile>>,
    pub session_summaries: Store<Vec<SessionSummary>>,
    pub handoff_memos: Store<Vec<HandoffMemo>>,
    pub env_fingerprints: Store<HashMap<String, EnvFingerprint>>,
    pub env_requirements: Store<Vec<EnvRequirement>>,
    pub milestones: Store<Vec<Milestone>>,
    pub environments: Store<Vec<EnvironmentDetail>>,
    pub pr_checklists: Store<Vec<PrChecklistItem>>,
    pub tech_debts: Store<Vec<TechDebt>>,
    pub gates: Store<Vec<GateRecord>>,
    pub context_workspaces: Store<Vec<ContextWorkspace>>,
}

impl MemoryState {
    fn master_mtime(&self) -> SystemTime {
        fs::metadata(&self.master_path)
            .and_then(|m| m.modified())
            .unwrap_or(SystemTime::UNIX_EPOCH)
    }

    pub fn unique_items<T: Eq + std::hash::Hash + Clone>(input: Vec<T>) -> Vec<T> {
        let mut keys = HashSet::new();
        let mut list = Vec::new();
        for entry in input {
            if keys.insert(entry.clone()) {
                list.push(entry);
            }
        }
        list
    }
    pub fn merge_graphs(dest: &mut KnowledgeGraph, src: &KnowledgeGraph) {
        for (name, src_ent) in &src.entities {
            let dest_ent = dest
                .entities
                .entry(name.clone())
                .or_insert_with(|| src_ent.clone());
            if dest_ent.name == src_ent.name {
                dest_ent.observations.extend(src_ent.observations.clone());
                dest_ent.observations = Self::unique_items(dest_ent.observations.clone());
            }
        }
        dest.relations.extend(src.relations.clone());
        dest.relations = Self::unique_items(dest.relations.clone());
    }

    pub fn read_master_cached(&self) -> KnowledgeGraph {
        let current_mtime = self.master_mtime();
        {
            let lock = self.master_cache.read().unwrap();
            if lock.1 == current_mtime {
                return lock.0.clone();
            }
        }
        let mut lock = self.master_cache.write().unwrap();
        let new_mtime = self.master_mtime();
        if lock.1 != new_mtime {
            lock.0 = read_json_file(&self.master_path);
            lock.1 = new_mtime;
        }
        lock.0.clone()
    }

    pub fn get_full_graph(&self) -> KnowledgeGraph {
        let mut master = self.read_master_cached();
        let wal_path = self.base_dir.join("wal.jsonl");
        if let Ok(content) = std::fs::read_to_string(&wal_path) {
            for line in content.lines() {
                if let Ok(d) = serde_json::from_str::<KnowledgeGraph>(line) {
                    Self::merge_graphs(&mut master, &d);
                }
            }
        }
        let session_graph = self.session_graph.read().unwrap();
        Self::merge_graphs(&mut master, &session_graph);
        master
    }
    pub fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
        let mut session_graph = self.session_graph.write().unwrap();
        update_fn(&mut session_graph);
        let wal_path = self.base_dir.join("wal.jsonl");
        if let Ok(payload) = serde_json::to_string(&*session_graph) {
            if let Ok(mut file) = std::fs::OpenOptions::new().create(true).append(true).open(&wal_path) {
                use std::io::Write;
                let _ = writeln!(file, "{}", payload);
            }
        }
    }
    pub fn apply_sync_write<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
        let lock_path = self.base_dir.join("master.lock");
        let mut attempts = 0;
        loop {
            if fs::OpenOptions::new()
                .create_new(true)
                .write(true)
                .open(&lock_path)
                .is_ok()
            {
                break;
            }
            if attempts > 100 {
                let _ = fs::remove_file(&lock_path);
            }
            attempts += 1;
            std::thread::sleep(Duration::from_millis(50));
        }
        let mut master = self.get_full_graph();
        let wal_path = self.base_dir.join("wal.jsonl");
        let _ = fs::remove_file(&wal_path);
        *self.session_graph.write().unwrap() = KnowledgeGraph::default();
        update_fn(&mut master);
        let _ = write_json_atomic(&self.master_path, &master);

        let mut cache_lock = self.master_cache.write().unwrap();
        cache_lock.0 = master;
        cache_lock.1 = self.master_mtime();

        let _ = fs::remove_file(&lock_path);
    }

    pub fn rebuild_index(&self) {
        if let Ok(new_idx) = MemoryIndex::new() {
            let full = self.get_full_graph();
            for (_, e) in &full.entities {
                let _ = new_idx.index_entity(e);
            }
            for t in self.tasks.read() {
                let _ = new_idx.index_task(&t);
            }
            for s in self.snippets.read() {
                let _ = new_idx.index_snippet(&s);
            }
            for a in self.adrs.read() {
                let _ = new_idx.index_adr(&a);
            }
            if let Ok(mut w) = self.search_index.write() {
                *w = new_idx;
            }
        }
    }
}