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

refactor(mcp): strip sse fallback architecture in favor of pure websockets and fix nvim NDJSON bug

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-17 15:26:22 +01:00
1 parent 0e29b12ac8
commit 3716c3e698
33 files changed
+2082 -1756
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/state.rs
+18 -169
View File
@@ -2,16 +2,12 @@ use crate::models::*;
use crate::search::MemoryIndex;
use crate::store::Store;
use std::collections::HashMap;
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 graph: Store<KnowledgeGraph>,
pub search_index: RwLock<MemoryIndex>,
pub ledger: Store<Vec<CodeChange>>,
pub sticky: Store<Vec<StickyNote>>,
@@ -31,189 +27,42 @@ pub struct MemoryState {
pub tech_debts: Store<Vec<TechDebt>>,
pub gates: Store<Vec<GateRecord>>,
pub context_workspaces: Store<Vec<ContextWorkspace>>,
pub activity_tx: tokio::sync::broadcast::Sender<String>,
}
impl MemoryState {
pub fn unique_items<T: Eq + std::hash::Hash + Clone>(input: Vec<T>) -> Vec<T> {
let mut keys = std::collections::HashSet::new();
input.into_iter().filter(|entry| keys.insert(entry.clone())).collect()
}
fn master_mtime(&self) -> SystemTime {
fs::metadata(&self.master_path)
.and_then(|m| m.modified())
.unwrap_or(SystemTime::UNIX_EPOCH)
pub fn broadcast_activity(&self, message: &str) {
let payload = serde_json::json!({
"type": "activity",
"data": message
}).to_string();
let _ = self.activity_tx.send(payload);
}
pub fn recover_wal(&self) {
let wal_path = self.base_dir.join("wal.jsonl");
if let Ok(content) = std::fs::read_to_string(&wal_path) {
let mut session = self.session_graph.write().unwrap();
for line in content.lines() {
if let Ok(d) = serde_json::from_str::<KnowledgeGraph>(line) {
Self::merge_graphs(&mut session, &d);
}
}
}
}
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(|| crate::models::Entity {
name: src_ent.name.clone(),
entity_type: src_ent.entity_type.clone(),
observations: Vec::new(),
namespace: src_ent.namespace.clone(),
git_branch: src_ent.git_branch.clone(),
});
for obs in &src_ent.observations {
if !dest_ent.observations.contains(obs) {
dest_ent.observations.push(obs.clone());
}
}
}
for rel in &src.relations {
if !dest.relations.contains(rel) {
dest.relations.push(rel.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 session_graph = self.session_graph.read().unwrap();
Self::merge_graphs(&mut master, &session_graph);
master
self.graph.read()
}
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 fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
self.graph.modify(update_fn);
}
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 apply_sync_write<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
self.graph.modify(update_fn);
}
pub fn rebuild_index(&self) {
if let Ok(new_idx) = MemoryIndex::new(&self.base_dir) {
let session_clone = { self.session_graph.read().unwrap().clone() };
let cache_clone = { self.master_cache.read().unwrap().0.clone() };
// Index entities that are only in master, or merge if they are in both
for (name, e) in &cache_clone.entities {
if let Some(session_e) = session_clone.entities.get(name) {
let mut merged_e = e.clone();
for obs in &session_e.observations {
if !merged_e.observations.contains(obs) {
merged_e.observations.push(obs.clone());
}
}
let _ = new_idx.index_entity(&merged_e);
} else {
let _ = new_idx.index_entity(e);
}
}
// Index entities that are only in session
for (name, session_e) in &session_clone.entities {
if !cache_clone.entities.contains_key(name) {
let _ = new_idx.index_entity(session_e);
}
let graph = self.graph.read();
for (_, e) in &graph.entities {
let _ = new_idx.index_entity(e);
}
let tasks = self.tasks.read();