mcp-memory/server/src/state.rs

use crate::models::*;
use crate::search::MemoryIndex;
use crate::store::Store;
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 {
            if let Ok(data) = fs::read(&self.master_path)
                && let Ok(parsed) = serde_json::from_slice(&data)
            {
                lock.0 = parsed;
            } else {
                let bak_path = self.master_path.with_extension("json.bak");
                if let Ok(data) = fs::read(&bak_path)
                    && let Ok(parsed) = serde_json::from_slice(&data)
                {
                    let _ = fs::write(&self.master_path, data);
                    lock.0 = parsed;
                } else {
                    lock.0 = KnowledgeGraph::default();
                }
            }
            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 async fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
        let payload = {
            let mut session_graph = self.session_graph.write().unwrap();
            update_fn(&mut session_graph);
            serde_json::to_string(&*session_graph).ok()
        };
        if let Some(payload) = payload {
            let wal_path = self.base_dir.join("wal.jsonl");
            if let Ok(mut file) = tokio::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(&wal_path)
                .await
            {
                use tokio::io::AsyncWriteExt;
                let _ = file.write_all(payload.as_bytes()).await;
                let _ = file.write_all(b"\n").await;
            }
        }
    }
    pub async 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 tokio::fs::OpenOptions::new()
                .create_new(true)
                .write(true)
                .open(&lock_path)
                .await
                .is_ok()
            {
                break;
            }
            if attempts > 100 {
                let _ = tokio::fs::remove_file(&lock_path).await;
            }
            attempts += 1;
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
        let mut master = self.get_full_graph();
        let wal_path = self.base_dir.join("wal.jsonl");
        let _ = tokio::fs::remove_file(&wal_path).await;
        *self.session_graph.write().unwrap() = KnowledgeGraph::default();
        update_fn(&mut master);
        
        let master_path = self.master_path.clone();
        let master_clone = master.clone();
        let _ = tokio::task::spawn_blocking(move || {
            let write_json = |path: &std::path::Path, data: &KnowledgeGraph| -> std::io::Result<()> {
                if path.exists() {
                    let bak_path = path.with_extension("json.bak");
                    let _ = std::fs::copy(path, &bak_path);
                }
                let tmp_path = path.with_extension("json.tmp");
                let json_data = serde_json::to_string_pretty(data)?;
                std::fs::write(&tmp_path, json_data)?;
                std::fs::rename(&tmp_path, path)
            };
            let _ = write_json(&master_path, &master_clone);
        }).await;

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

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

    pub fn rebuild_index(&self) {
        if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
            let session = self.session_graph.read().unwrap();
            let mut full = {
                let cache = self.master_cache.read().unwrap();
                cache.0.clone()
            };
            Self::merge_graphs(&mut full, &session);
            for e in full.entities.values() {
                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;
            }
        }
    }
}