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

refactor: apply 5-pass audit optimizations across mcp-memory codebase

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-06 06:05:38 +01:00
1 parent 924b6d09fa
commit 5bd8b1587a
43 files changed
+1866 -1658
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/state.rs
+78 -46
View File
@@ -4,7 +4,7 @@ use crate::store::Store;
use crate::vector_db::VectorDB;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::sync::Arc;
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct GenericEvent {
@@ -56,7 +56,7 @@ pub struct MemoryState {
pub ttl_notify: Arc<tokio::sync::Notify>,
pub condense_notify: Arc<tokio::sync::Notify>,
pub graph: Store<KnowledgeGraph>,
pub search_index: RwLock<MemoryIndex>,
pub search_index: tokio::sync::RwLock<MemoryIndex>,
pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,
pub project: ProjectStores,
@@ -70,14 +70,37 @@ pub struct MemoryState {
}
impl MemoryState {
pub fn new_in_memory() -> Self {
Self::new(":memory:")
}
pub fn new(base_dir_str: &str) -> Self {
let is_in_memory = base_dir_str == ":memory:";
let base = std::path::PathBuf::from(base_dir_str);
if let Err(e) = std::fs::create_dir_all(&base) {
tracing::error!("Failed to create store directory at {:?}: {}", base, e);
if !is_in_memory {
if let Err(e) = std::fs::create_dir_all(&base) {
tracing::error!("Failed to create store directory at {:?}: {}", base, e);
}
}
let db = crate::db::init_redb(&base);
let search_index = if is_in_memory {
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
} else {
match crate::search::MemoryIndex::new(&base) {
Ok(idx) => idx,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
}
}
};
let state = Self {
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
clipboard_watch_mode: tokio::sync::RwLock::new(false),
@@ -87,17 +110,7 @@ impl MemoryState {
condense_notify: Arc::new(tokio::sync::Notify::new()),
graph: Store::new("knowledge_graph_master", db.clone()),
base_dir: base.clone(),
search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) {
Ok(idx) => idx,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
}
}),
search_index: tokio::sync::RwLock::new(search_index),
vector_db: tokio::sync::RwLock::new(None),
project: ProjectStores {
@@ -164,11 +177,8 @@ impl MemoryState {
self.graph.modify(update_fn);
}
pub fn get_search_index(&self) -> MemoryIndex {
self.search_index
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
pub async fn get_search_index(&self) -> MemoryIndex {
self.search_index.read().await.clone()
}
pub fn search(self: &Arc<Self>) -> SearchService {
@@ -176,8 +186,18 @@ impl MemoryState {
}
pub async fn rebuild_index(self: &Arc<Self>) {
let idx = self.get_search_index();
idx.delete_all();
let is_in_memory = self.base_dir.to_str() == Some(":memory:");
let new_idx = if is_in_memory {
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
} else {
match crate::search::MemoryIndex::new(&self.base_dir) {
Ok(idx) => idx,
Err(e) => {
tracing::warn!("Failed to create disk MemoryIndex for rebuild ({}), falling back to RAM", e);
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
}
}
};
let entities: Vec<_> = self
.graph
@@ -192,9 +212,8 @@ impl MemoryState {
tasks.len()
);
let idx_clone = idx.clone();
let idx_clone = new_idx.clone();
tokio::task::spawn_blocking(move || {
// tracing::info!("spawn_blocking started in rebuild_index");
for e in entities {
idx_clone.add_entity_sync(&e);
}
@@ -213,10 +232,8 @@ impl MemoryState {
tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
});
let _ = idx.commit().await;
if let Ok(mut w) = self.search_index.write() {
*w = idx;
}
let _ = new_idx.commit().await;
*self.search_index.write().await = new_idx;
self.index_commit_notify.notify_waiters();
}
@@ -350,7 +367,7 @@ mod tests {
arc_state.rebuild_index().await;
// Check search index initialization
let idx = arc_state.search_index.read().unwrap();
let idx = arc_state.search_index.read().await;
// Force reload reader to ensure it sees the commit made by rebuild_index
idx.reader.reload().unwrap();
// tracing::info!(
@@ -464,31 +481,46 @@ impl SearchService {
}
if !vdb_search {
let (mut texts_to_embed, mut metadata) = self.state.code.snippets.read_with(|snips| {
let mut texts = Vec::with_capacity(snips.len().min(50));
let mut meta = Vec::with_capacity(snips.len().min(50));
let mut cached_items = Vec::new();
let mut uncached_texts = Vec::new();
let mut uncached_meta = Vec::new();
self.state.code.snippets.read_with(|snips| {
for snippet in snips.iter().take(50) {
texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code));
meta.push((snippet.name.clone(), "snippet".to_string(), snippet.description.clone()));
let title = snippet.name.clone();
let desc = snippet.description.clone();
if let Some(ref emb) = snippet.embedding {
cached_items.push((title, "snippet".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code));
uncached_meta.push((title, "snippet".to_string(), desc));
}
}
(texts, meta)
});
self.state.code.sticky.read_with(|sticky| {
for note in sticky.iter().take(50) {
texts_to_embed.push(note.content.clone());
metadata.push((
"StickyNote".to_string(),
"sticky".to_string(),
note.content.chars().take(200).collect::<String>(),
));
let content_preview = note.content.chars().take(200).collect::<String>();
uncached_texts.push(note.content.clone());
uncached_meta.push(("StickyNote".to_string(), "sticky".to_string(), content_preview));
}
});
if !texts_to_embed.is_empty()
&& let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await
for (title, doc_type, body, emb) in cached_items {
let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: title.clone(),
doc_type,
title,
body,
score: sim,
});
}
if !uncached_texts.is_empty()
&& let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
{
for (emb, meta) in embeddings.into_iter().zip(metadata) {
for (emb, meta) in embeddings.into_iter().zip(uncached_meta) {
let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: meta.0.clone(),
@@ -511,13 +543,13 @@ impl SearchService {
Ok(results)
}
pub fn keyword_search(
pub async fn keyword_search(
&self,
query: &str,
filter_namespace: Option<&str>,
limit: usize,
) -> crate::error::Result<Vec<UnifiedSearchResult>> {
let idx = self.state.get_search_index();
let idx = self.state.get_search_index().await;
let matches = idx
.search(query, filter_namespace)
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;