riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: finalize redb migration for all memory stores

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-10 10:48:37 +01:00
1 parent 9a8b6e52e9
commit 721f564910
4 files changed
+107 -88
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/main.rs
+28 -18
View File
@@ -407,30 +407,40 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
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();
server/src/state.rs
+35 -5
View File
@@ -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);
}
server/src/store.rs
+36 -65
View File
@@ -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<T: DeserializeOwned + Default>(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<T: Serialize>(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<T> {
pub path: PathBuf,
pub cache: RwLock<(T, SystemTime)>,
pub key: String,
pub db: Arc<Database>,
pub cache: RwLock<T>,
}
impl<T: DeserializeOwned + Default + Serialize + Clone> Store<T> {
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<Database>) -> 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::<T>(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<F: FnOnce(&mut T)>(&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);
}
}