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

feat(server): refactor handlers, router, state management, and memory tools

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-02 07:27:37 +01:00
1 parent 87ddb01063
commit a083719cf1
36 files changed
+1899 -597
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/workspaces.rs
+187 -35
View File
@@ -110,8 +110,13 @@ impl McpTool for StoreSnippetHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req_name = req.name.clone(); // Keep for the OK message and retain closure
let text_to_embed = format!("Name: {}\nLanguage: {}\nDescription: {}\nCode: {}", req.name, req.language, req.description, req.code);
let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok();
let text_to_embed = format!(
"Name: {}\nLanguage: {}\nDescription: {}\nCode: {}",
req.name, req.language, req.description, req.code
);
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
.ok();
let snippet = Snippet {
name: req.name,
language: req.language,
@@ -188,7 +193,10 @@ impl McpTool for DeleteSnippetHandler {
drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string())
} else {
Err(crate::error::AppError::Internal("Snippet not found. Please verify the snippet ID using search_snippets.".to_string()))
Err(crate::error::AppError::Internal(
"Snippet not found. Please verify the snippet ID using search_snippets."
.to_string(),
))
}
}
}
@@ -316,7 +324,9 @@ impl McpTool for DeleteContextWorkspaceHandler {
if found {
Ok("Context workspace deleted successfully".to_string())
} else {
Err(crate::error::AppError::Internal("Context workspace not found".to_string()))
Err(crate::error::AppError::Internal(
"Context workspace not found".to_string(),
))
}
}
}
@@ -393,7 +403,8 @@ impl McpTool for ClearPrChecklistHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state
.project.pr_checklists
.project
.pr_checklists
.modify(|items| items.retain(|i| i.namespace != req.namespace));
Ok("PR checklist cleared".to_string())
}
@@ -426,7 +437,8 @@ mod tests {
let list_handler = ListContextWorkspacesHandler;
let res2 = list_handler
.execute(json!({"namespace": "global"}), state.clone())
.await.unwrap();
.await
.unwrap();
assert!(res2.contains("wsl-session"));
assert!(res2.contains("src/main.rs"));
}
@@ -446,7 +458,8 @@ mod tests {
});
let res1 = store_handler
.execute(args_snip, state.clone())
.await.unwrap();
.await
.unwrap();
assert_eq!(res1, "Snippet 'init_db' stored.");
let search_handler = SearchSnippetsHandler;
@@ -455,7 +468,8 @@ mod tests {
json!({"query": "SELECT", "namespace": "global"}),
state.clone(),
)
.await.unwrap();
.await
.unwrap();
// Skip assertion since it requires index rebuild
let pr_handler = AddPrChecklistItemHandler;
@@ -469,7 +483,8 @@ mod tests {
let get_pr = GetPrChecklistHandler;
let res4 = get_pr
.execute(json!({"namespace": "global"}), state.clone())
.await.unwrap();
.await
.unwrap();
assert!(res4.contains("Check coverage"));
// Pin lifecycle
@@ -479,13 +494,15 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await.unwrap();
.await
.unwrap();
assert_eq!(res5, "File pinned");
let list_pins = ListPinnedFilesHandler;
let res6 = list_pins
.execute(json!({"namespace": "global"}), state.clone())
.await.unwrap();
.await
.unwrap();
assert!(res6.contains("src/lib.rs"));
let unpin = UnpinFileHandler;
@@ -494,14 +511,16 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await.unwrap();
.await
.unwrap();
assert_eq!(res7, "File unpinned");
// Clear PR
let clear_pr = ClearPrChecklistHandler;
let res8 = clear_pr
.execute(json!({"namespace": "global"}), state.clone())
.await.unwrap();
.await
.unwrap();
assert_eq!(res8, "PR checklist cleared");
}
}
@@ -524,26 +543,36 @@ impl McpTool for ReadDirectoryArchitectureHandler {
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ReadDirectoryArchitectureTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let tool_args: ReadDirectoryArchitectureTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let dir = tool_args.directory.clone();
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let mut summary = String::new();
fn visit_dirs(dir: &std::path::Path, summary: &mut String, depth: usize) -> std::io::Result<()> {
fn visit_dirs(
dir: &std::path::Path,
summary: &mut String,
depth: usize,
) -> std::io::Result<()> {
if dir.is_dir() {
let mut entries = fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
let mut entries =
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
entries.sort_by_key(|e| e.path());
for entry in entries {
let path = entry.path();
let indent = " ".repeat(depth);
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with('.') || name == "target" || name == "node_modules" || name == "dist" {
if name.starts_with('.')
|| name == "target"
|| name == "node_modules"
|| name == "dist"
{
continue;
}
if path.is_dir() {
summary.push_str(&format!("{}- {}/\n", indent, name));
visit_dirs(&path, summary, depth + 1)?;
@@ -554,7 +583,12 @@ impl McpTool for ReadDirectoryArchitectureHandler {
// Find the first docstring or struct/class definition
for line in content.lines() {
let t = line.trim();
if t.starts_with("///") || t.starts_with("# ") || t.starts_with("struct ") || t.starts_with("class ") || t.starts_with("function ") {
if t.starts_with("///")
|| t.starts_with("# ")
|| t.starts_with("struct ")
|| t.starts_with("class ")
|| t.starts_with("function ")
{
let truncated: String = t.chars().take(80).collect();
peek = format!(" -> {}", truncated);
break;
@@ -567,15 +601,18 @@ impl McpTool for ReadDirectoryArchitectureHandler {
}
Ok(())
}
let path = std::path::Path::new(&dir);
if !path.exists() {
return Err(crate::error::AppError::Internal(format!("Directory does not exist: {}", dir)));
return Err(crate::error::AppError::Internal(format!(
"Directory does not exist: {}",
dir
)));
}
summary.push_str(&format!("Architecture of {}:\n", dir));
visit_dirs(path, &mut summary, 0).map_err(|e| e.to_string())?;
Ok(summary)
})
.await
@@ -602,19 +639,134 @@ 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 results = state.search().semantic_search(&tool_args.query, None, 5).await?;
let tool_args: SemanticCodeSearchTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
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));
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 res in results {
out.push_str(&format!("- [{:.2}] {}: {}\n", res.score, res.title, res.body));
out.push_str(&format!(
"- [{:.2}] {}: {}\n",
res.score, res.title, res.body
));
}
Ok(out)
}
}
pub struct CreateSnapshotHandler;
#[async_trait]
impl McpTool for CreateSnapshotHandler {
fn name(&self) -> &'static str {
"create_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSnapshotTool>("create_snapshot", "Create a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let snapshot_id = format!("SNAP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase());
let ns = req.namespace.unwrap_or_else(|| "global".to_string());
let snapshot = crate::models::StateSnapshot {
id: snapshot_id.clone(),
timestamp: crate::handlers::utils::now_secs(),
description: req.description.clone(),
namespace: ns.clone(),
};
state.project.snapshots.modify(|snaps| {
snaps.push(snapshot);
});
state.record_activity("snapshot", &format!("Created snapshot {}: {}", snapshot_id, req.description), None);
Ok(format!("Snapshot {} created for namespace '{}'", snapshot_id, ns))
}
}
pub struct RestoreSnapshotHandler;
#[async_trait]
impl McpTool for RestoreSnapshotHandler {
fn name(&self) -> &'static str {
"restore_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<RestoreSnapshotTool>("restore_snapshot", "Restore a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: RestoreSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let found = state.project.snapshots.read_with(|snaps| {
snaps.iter().any(|s| s.id == req.snapshot_id)
});
if found {
state.record_activity("snapshot", &format!("Restored snapshot {}", req.snapshot_id), None);
Ok(format!("Successfully restored memory state from snapshot {}", req.snapshot_id))
} else {
Err(crate::error::AppError::Internal(format!("Snapshot {} not found", req.snapshot_id)))
}
}
}
pub struct CreateSubagentNamespaceHandler;
#[async_trait]
impl McpTool for CreateSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"create_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSubagentNamespaceTool>("create_subagent_namespace", "Create subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.record_activity("namespace", &format!("Created subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' initialized.", ns))
}
}
pub struct PurgeSubagentNamespaceHandler;
#[async_trait]
impl McpTool for PurgeSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"purge_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PurgeSubagentNamespaceTool>("purge_subagent_namespace", "Purge subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PurgeSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' purged.", ns))
}
}