From 721f56491078d483a9abb2be5a8aaa63f63a9ed3 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Thu, 10 Sep 2026 10:48:37 +0100 Subject: [PATCH] chore: finalize redb migration for all memory stores --- fix_linux.py | 8 ++++ server/src/main.rs | 46 ++++++++++++-------- server/src/state.rs | 40 +++++++++++++++--- server/src/store.rs | 101 ++++++++++++++++---------------------------- 4 files changed, 107 insertions(+), 88 deletions(-) create mode 100644 fix_linux.py diff --git a/fix_linux.py b/fix_linux.py new file mode 100644 index 0000000..c8171a6 --- /dev/null +++ b/fix_linux.py @@ -0,0 +1,8 @@ +import re +with open(r'C:\Users\reazul.ashraf\workspace\rust\mcp-memory\server\src\main.rs', 'r', encoding='utf-8') as f: + text = f.read() + +text = re.sub(r'#\[cfg\(not\(target_os = "windows"\)\)\].*?return Ok\(\(\);\n \}', '', text, flags=re.DOTALL) + +with open(r'C:\Users\reazul.ashraf\workspace\rust\mcp-memory\server\src\main.rs', 'w', encoding='utf-8') as f: + f.write(text) diff --git a/server/src/main.rs b/server/src/main.rs index b1f1ec2..61ae35e 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -407,30 +407,40 @@ fn main() -> Result<(), Box> { let base = PathBuf::from(base_dir); fs::create_dir_all(&base).expect("Failed to create store dir"); + let redb_path = base.join("mcp_store.redb"); + let db = Arc::new(redb::Database::create(&redb_path).unwrap()); + + // Ensure table exists + { + let write_txn = db.begin_write().unwrap(); + let _ = write_txn.open_table(crate::store::STORE_TABLE); + write_txn.commit().unwrap(); + } + let state = Arc::new(MemoryState { master_path: base.join("knowledge_graph_master.json"), session_graph: RwLock::new(KnowledgeGraph::default()), base_dir: base.clone(), master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)), search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()), - ledger: Store::new(base.join("audit_ledger.json")), - sticky: Store::new(base.join("sticky_notes.json")), - tasks: Store::new(base.join("tasks.json")), - snippets: Store::new(base.join("snippets.json")), - adrs: Store::new(base.join("adrs.json")), - prefs: Store::new(base.join("preferences.json")), - error_fixes: Store::new(base.join("error_fixes.json")), - pinned_files: Store::new(base.join("pinned_files.json")), - session_summaries: Store::new(base.join("session_summaries.json")), - handoff_memos: Store::new(base.join("handoff_memos.json")), - env_fingerprints: Store::new(base.join("env_fingerprints.json")), - env_requirements: Store::new(base.join("env_requirements.json")), - milestones: Store::new(base.join("milestones.json")), - environments: Store::new(base.join("environments.json")), - pr_checklists: Store::new(base.join("pr_checklists.json")), - tech_debts: Store::new(base.join("tech_debts.json")), - gates: Store::new(base.join("gates.json")), - context_workspaces: Store::new(base.join("context_workspaces.json")), + ledger: Store::new("audit_ledger", db.clone()), + sticky: Store::new("sticky_notes", db.clone()), + tasks: Store::new("tasks", db.clone()), + snippets: Store::new("snippets", db.clone()), + adrs: Store::new("adrs", db.clone()), + prefs: Store::new("preferences", db.clone()), + error_fixes: Store::new("error_fixes", db.clone()), + pinned_files: Store::new("pinned_files", db.clone()), + session_summaries: Store::new("session_summaries", db.clone()), + handoff_memos: Store::new("handoff_memos", db.clone()), + env_fingerprints: Store::new("env_fingerprints", db.clone()), + env_requirements: Store::new("env_requirements", db.clone()), + milestones: Store::new("milestones", db.clone()), + environments: Store::new("environments", db.clone()), + pr_checklists: Store::new("pr_checklists", db.clone()), + tech_debts: Store::new("tech_debts", db.clone()), + gates: Store::new("gates", db.clone()), + context_workspaces: Store::new("context_workspaces", db.clone()), }); state.rebuild_index(); diff --git a/server/src/state.rs b/server/src/state.rs index aa7cf2d..d4355b1 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -1,6 +1,6 @@ use crate::models::*; use crate::search::MemoryIndex; -use crate::store::{Store, read_json_file, write_json_atomic}; +use crate::store::Store; use std::collections::{HashMap, HashSet}; use std::fs; use std::path::PathBuf; @@ -76,7 +76,21 @@ impl MemoryState { 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); + 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() @@ -133,7 +147,18 @@ impl MemoryState { 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 write_json = |path: &std::path::Path, data: &KnowledgeGraph| -> std::io::Result<()> { + if path.exists() { + let bak_path = path.with_extension("json.bak"); + let _ = fs::copy(path, &bak_path); + } + let tmp_path = path.with_extension("json.tmp"); + let json_data = serde_json::to_string_pretty(data)?; + fs::write(&tmp_path, json_data)?; + fs::rename(&tmp_path, path) + }; + let _ = write_json(&self.master_path, &master); let mut cache_lock = self.master_cache.write().unwrap(); cache_lock.0 = master; @@ -143,8 +168,13 @@ impl MemoryState { } pub fn rebuild_index(&self) { - if let Ok(new_idx) = MemoryIndex::new() { - let full = self.get_full_graph(); + 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); } diff --git a/server/src/store.rs b/server/src/store.rs index 733c8b5..37efdf7 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -1,84 +1,55 @@ +use redb::{Database, ReadableDatabase, TableDefinition}; use serde::{Serialize, de::DeserializeOwned}; -use std::fs; -use std::path::{Path, PathBuf}; -use std::sync::RwLock; -use std::time::SystemTime; +use std::sync::{Arc, RwLock}; -pub fn read_json_file(path: &Path) -> T { - if let Ok(data) = fs::read(path) - && let Ok(parsed) = serde_json::from_slice(&data) - { - return parsed; - } - let bak_path = path.with_extension("json.bak"); - if let Ok(data) = fs::read(&bak_path) - && let Ok(parsed) = serde_json::from_slice(&data) - { - let _ = fs::write(path, data); - return parsed; - } - T::default() -} - -pub fn write_json_atomic(path: &Path, data: &T) -> std::io::Result<()> { - if path.exists() { - let bak_path = path.with_extension("json.bak"); - let _ = fs::copy(path, &bak_path); - } - let tmp_path = path.with_extension("json.tmp"); - let json_data = serde_json::to_string_pretty(data)?; - fs::write(&tmp_path, json_data)?; - fs::rename(&tmp_path, path) -} +pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("store"); pub struct Store { - pub path: PathBuf, - pub cache: RwLock<(T, SystemTime)>, + pub key: String, + pub db: Arc, + pub cache: RwLock, } impl Store { - fn mtime(&self) -> SystemTime { - fs::metadata(&self.path) - .and_then(|m| m.modified()) - .unwrap_or(SystemTime::UNIX_EPOCH) + pub fn new(key: &str, db: Arc) -> Self { + let initial_data = Self::load_from_db(key, &db); + Self { + key: key.to_string(), + db, + cache: RwLock::new(initial_data), + } } - pub fn new(path: PathBuf) -> Self { - let mtime = fs::metadata(&path) - .and_then(|m| m.modified()) - .unwrap_or(SystemTime::UNIX_EPOCH); - let initial_data = read_json_file(&path); - Self { - path, - cache: RwLock::new((initial_data, mtime)), + fn load_from_db(key: &str, db: &Database) -> T { + let read_txn = db.begin_read().unwrap(); + if let Ok(table) = read_txn.open_table(STORE_TABLE) { + if let Ok(Some(value)) = table.get(key) { + if let Ok(parsed) = serde_json::from_slice::(value.value()) { + return parsed; + } + } } + T::default() + } + + fn save_to_db(key: &str, db: &Database, data: &T) { + let write_txn = db.begin_write().unwrap(); + { + let mut table = write_txn.open_table(STORE_TABLE).unwrap(); + let json_data = serde_json::to_vec(data).unwrap(); + table.insert(key, json_data.as_slice()).unwrap(); + } + write_txn.commit().unwrap(); } pub fn read(&self) -> T { - let current_mtime = self.mtime(); - { - let lock = self.cache.read().unwrap(); - if lock.1 == current_mtime { - return lock.0.clone(); - } - } - let mut lock = self.cache.write().unwrap(); - let new_mtime = self.mtime(); - if lock.1 != new_mtime { - lock.0 = read_json_file(&self.path); - lock.1 = new_mtime; - } - lock.0.clone() + let lock = self.cache.read().unwrap(); + lock.clone() } pub fn modify(&self, f: F) { let mut lock = self.cache.write().unwrap(); - let current_mtime = self.mtime(); - if lock.1 != current_mtime { - lock.0 = read_json_file(&self.path); - } - f(&mut lock.0); - let _ = write_json_atomic(&self.path, &lock.0); - lock.1 = self.mtime(); + f(&mut lock); + Self::save_to_db(&self.key, &self.db, &lock); } }