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

refactor: Implement unified search abstraction, MemoryState refactoring, error handling, and unit test expansion

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-01 08:37:52 +01:00
1 parent a34554b7ff
commit 462f65f66d
21 files changed
+425 -535
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/workspaces.rs
+13 -56
View File
@@ -420,13 +420,13 @@ mod tests {
"active_task_ids": ["123"]
});
let res = save_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
let res = save_handler.execute(args, state.clone()).await.unwrap();
assert_eq!(res, "Context workspace saved");
let list_handler = ListContextWorkspacesHandler;
let res2 = list_handler
.execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert!(res2.contains("wsl-session"));
assert!(res2.contains("src/main.rs"));
}
@@ -446,7 +446,7 @@ mod tests {
});
let res1 = store_handler
.execute(args_snip, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert_eq!(res1, "Snippet 'init_db' stored.");
let search_handler = SearchSnippetsHandler;
@@ -455,7 +455,7 @@ mod tests {
json!({"query": "SELECT", "namespace": "global"}),
state.clone(),
)
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
// Skip assertion since it requires index rebuild
let pr_handler = AddPrChecklistItemHandler;
@@ -463,13 +463,13 @@ mod tests {
"description": "Check coverage",
"namespace": "global"
});
let res3 = pr_handler.execute(args_pr, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
let res3 = pr_handler.execute(args_pr, state.clone()).await.unwrap();
assert_eq!(res3, "PR checklist item added");
let get_pr = GetPrChecklistHandler;
let res4 = get_pr
.execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert!(res4.contains("Check coverage"));
// Pin lifecycle
@@ -479,13 +479,13 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert_eq!(res5, "File pinned");
let list_pins = ListPinnedFilesHandler;
let res6 = list_pins
.execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert!(res6.contains("src/lib.rs"));
let unpin = UnpinFileHandler;
@@ -494,14 +494,14 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert_eq!(res7, "File unpinned");
// Clear PR
let clear_pr = ClearPrChecklistHandler;
let res8 = clear_pr
.execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
.await.unwrap();
assert_eq!(res8, "PR checklist cleared");
}
}
@@ -585,7 +585,6 @@ impl McpTool for ReadDirectoryArchitectureHandler {
}
}
use crate::tools::SemanticCodeSearchTool;
use crate::embedding::{generate_embedding_async, generate_embeddings_async, cosine_similarity};
pub struct SemanticCodeSearchHandler;
@@ -605,57 +604,15 @@ impl McpTool for SemanticCodeSearchHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: SemanticCodeSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let query_emb = generate_embedding_async(tool_args.query.clone()).await?;
let mut results = Vec::new();
// Search using VectorDB if available
let mut vdb_search = false;
if let Some(vdb) = &*state.vector_db.read().await {
vdb_search = true;
if let Ok(search_results) = vdb.search(query_emb.clone(), 5).await {
for res in search_results {
results.push((res.score, res.id, res.text));
}
}
}
// Fallback to manual loop if VectorDB is not initialized
if !vdb_search {
let mut texts_to_embed = Vec::new();
let mut metadata = Vec::new();
let snippets = state.code.snippets.read_with(|snips| snips.clone());
for snippet in snippets {
let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code);
texts_to_embed.push(combined);
metadata.push((snippet.name, snippet.description));
}
let sticky = state.code.sticky.read_with(|s| s.clone());
for note in sticky {
texts_to_embed.push(note.content.clone());
metadata.push(("StickyNote".to_string(), note.content.chars().take(200).collect::<String>()));
}
if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await {
for (emb, meta) in embeddings.into_iter().zip(metadata) {
let sim = cosine_similarity(&query_emb, &emb);
results.push((sim, meta.0, meta.1));
}
}
results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
results.truncate(5);
}
let results = state.search().semantic_search(&tool_args.query, None, 5).await?;
if results.is_empty() {
return Ok(format!("No semantic matches found for query: {}", tool_args.query));
}
let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query);
for (score, title, desc) in results {
out.push_str(&format!("- [{:.2}] {}: {}\n", score, title, desc));
for res in results {
out.push_str(&format!("- [{:.2}] {}: {}\n", res.score, res.title, res.body));
}
Ok(out)