Refactor: Migrate unwrap calls to AppError in MCP handlers

This commit is contained in:
Riz Ashraf committed 2026-09-30 21:02:50 +01:00
1 parent 4e1a633dbd
commit 0e866f2465
12 files changed
+276 -326

No files matched your search

+34 -45
View File
@@ -18,7 +18,7 @@ impl McpTool for PinFileHandler {
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state.pinned_files.modify(|pinned| {
pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path);
@@ -45,7 +45,7 @@ impl McpTool for UnpinFileHandler {
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: UnpinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state.pinned_files.modify(|pinned| {
pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path)
@@ -69,7 +69,7 @@ impl McpTool for ListPinnedFilesHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ListPinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.pinned_files.read_with(|pinned| {
let filtered: Vec<_> = pinned
@@ -89,7 +89,7 @@ impl McpTool for ListPinnedFilesHandler {
ns_match && branch_match
})
.collect();
serde_json::to_string(&filtered).map_err(|e| e.to_string())
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
@@ -107,7 +107,7 @@ impl McpTool for StoreSnippetHandler {
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
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);
@@ -145,7 +145,7 @@ impl McpTool for SearchSnippetsHandler {
crate::mcp::tool_def::<SearchSnippetsTool>("search_snippets", "Execute search_snippets")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SearchSnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let query = req.query;
let data = state.snippets.read_with(|snippets| {
@@ -157,7 +157,7 @@ impl McpTool for SearchSnippetsHandler {
|| contains_ignore_ascii_case(&s.language, &query)
})
.collect();
serde_json::to_string(&results).map_err(|e| e.to_string())
Ok::<String, crate::error::AppError>(serde_json::to_string(&results)?)
})?;
Ok(data)
}
@@ -175,7 +175,7 @@ impl McpTool for DeleteSnippetHandler {
crate::mcp::tool_def::<DeleteSnippetTool>("delete_snippet", "Execute delete_snippet")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: DeleteSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut deleted = false;
state.snippets.modify(|snippets| {
@@ -188,10 +188,7 @@ impl McpTool for DeleteSnippetHandler {
drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string())
} else {
Err(
"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()))
}
}
}
@@ -211,7 +208,7 @@ impl McpTool for SaveContextWorkspaceHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SaveContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.context_workspaces.modify(|ws| {
@@ -243,7 +240,7 @@ impl McpTool for LoadContextWorkspaceHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: LoadContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.context_workspaces.read_with(|ws| {
@@ -272,7 +269,7 @@ impl McpTool for ListContextWorkspacesHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ListContextWorkspacesTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.context_workspaces.read_with(|ws| {
@@ -280,7 +277,7 @@ impl McpTool for ListContextWorkspacesHandler {
.iter()
.filter(|w| req.namespace.as_ref().is_none_or(|ns| &w.namespace == ns))
.collect();
serde_json::to_string(&filtered).map_err(|e| e.to_string())
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
@@ -301,7 +298,7 @@ impl McpTool for DeleteContextWorkspaceHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: crate::tools::DeleteContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
@@ -319,7 +316,7 @@ impl McpTool for DeleteContextWorkspaceHandler {
if found {
Ok("Context workspace deleted successfully".to_string())
} else {
Err("Context workspace not found".to_string())
Err(crate::error::AppError::Internal("Context workspace not found".to_string()))
}
}
}
@@ -339,7 +336,7 @@ impl McpTool for AddPrChecklistItemHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: AddPrChecklistItemTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.pr_checklists.modify(|items| {
@@ -365,14 +362,14 @@ impl McpTool for GetPrChecklistHandler {
crate::mcp::tool_def::<GetPrChecklistTool>("get_pr_checklist", "Execute get_pr_checklist")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: GetPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.pr_checklists.read_with(|items| {
let filtered: Vec<_> = items
.iter()
.filter(|i| i.namespace == req.namespace)
.collect();
serde_json::to_string(&filtered).map_err(|e| e.to_string())
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
@@ -393,7 +390,7 @@ impl McpTool for ClearPrChecklistHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
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
.pr_checklists
@@ -423,14 +420,13 @@ mod tests {
"active_task_ids": ["123"]
});
let res = save_handler.execute(args, state.clone()).await.unwrap();
let res = save_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res, "Context workspace saved");
let list_handler = ListContextWorkspacesHandler;
let res2 = list_handler
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert!(res2.contains("wsl-session"));
assert!(res2.contains("src/main.rs"));
}
@@ -450,8 +446,7 @@ mod tests {
});
let res1 = store_handler
.execute(args_snip, state.clone())
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res1, "Snippet 'init_db' stored.");
let search_handler = SearchSnippetsHandler;
@@ -460,8 +455,7 @@ mod tests {
json!({"query": "SELECT", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
// Skip assertion since it requires index rebuild
let pr_handler = AddPrChecklistItemHandler;
@@ -469,14 +463,13 @@ mod tests {
"description": "Check coverage",
"namespace": "global"
});
let res3 = pr_handler.execute(args_pr, state.clone()).await.unwrap();
let res3 = pr_handler.execute(args_pr, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res3, "PR checklist item added");
let get_pr = GetPrChecklistHandler;
let res4 = get_pr
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert!(res4.contains("Check coverage"));
// Pin lifecycle
@@ -486,15 +479,13 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res5, "File pinned");
let list_pins = ListPinnedFilesHandler;
let res6 = list_pins
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert!(res6.contains("src/lib.rs"));
let unpin = UnpinFileHandler;
@@ -503,16 +494,14 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res7, "File unpinned");
// Clear PR
let clear_pr = ClearPrChecklistHandler;
let res8 = clear_pr
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
.await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
assert_eq!(res8, "PR checklist cleared");
}
}
@@ -534,11 +523,11 @@ impl McpTool for ReadDirectoryArchitectureHandler {
)
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> Result<String, String> {
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 dir = tool_args.directory.clone();
let result = tokio::task::spawn_blocking(move || -> Result<String, String> {
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<()> {
@@ -581,7 +570,7 @@ impl McpTool for ReadDirectoryArchitectureHandler {
let path = std::path::Path::new(&dir);
if !path.exists() {
return Err(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));
@@ -590,7 +579,7 @@ impl McpTool for ReadDirectoryArchitectureHandler {
Ok(summary)
})
.await
.map_err(|e| format!("Task panic: {}", e))??;
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok(result)
}
@@ -613,7 +602,7 @@ impl McpTool for SemanticCodeSearchHandler {
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
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?;