Refactor: Migrate unwrap calls to AppError in MCP handlers

server/src/handlers/graph.rs

@@ -24,7 +24,7 @@ impl McpTool for QueryGraphPathHandler {
         crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path")
     }
 
-    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::QueryGraphPathTool =
             serde_json::from_value(args).map_err(|e| e.to_string())?;
         state.read_graph(|graph| {
@@ -115,7 +115,7 @@ impl McpTool for CreateEntitiesHandler {
         crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities")
     }
 
-    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: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut inserted = Vec::new();
         state.modify_graph(|g| {
@@ -146,7 +146,7 @@ impl McpTool for CreateRelationsHandler {
         crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations")
     }
 
-    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: CreateRelationsTool = match serde_json::from_value(args.clone()) {
             Ok(r) => r,
             Err(e) => {
@@ -155,12 +155,12 @@ impl McpTool for CreateRelationsHandler {
                     || err_msg.contains("missing field `to`")
                     || err_msg.contains("missing field `relation_type`")
                 {
-                    return Err(format!(
+                    return Err(crate::error::AppError::Internal(format!(
                         "Schema error: {}. Note that the relation schema strictly uses 'from', 'to', and 'relation_type' (not 'source', 'target', or 'relationType'). Please correct your tool call arguments.",
                         err_msg
-                    ));
+                    )));
                 }
-                return Err(err_msg);
+                return Err(err_msg.into());
             }
         };
         let mut missing_nodes = std::collections::HashSet::new();
@@ -184,10 +184,9 @@ impl McpTool for CreateRelationsHandler {
         });
         if !missing_nodes.is_empty() {
             let missing: Vec<_> = missing_nodes.into_iter().collect();
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Relations dropped due to missing entities: {}",
-                missing.join(", ")
-            ));
+                missing.join(", "))));
         }
         Ok("Relations created".to_string())
     }
@@ -205,7 +204,7 @@ impl McpTool for AddObservationsHandler {
         crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations")
     }
 
-    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: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut missing_entities = Vec::new();
         state.modify_graph(|g| {
@@ -218,10 +217,9 @@ impl McpTool for AddObservationsHandler {
             }
         });
         if !missing_entities.is_empty() {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Observations dropped for missing entities: {}",
-                missing_entities.join(", ")
-            ));
+                missing_entities.join(", "))));
         }
         Ok("Observations added".to_string())
     }
@@ -239,7 +237,7 @@ impl McpTool for DeleteEntitiesHandler {
         crate::mcp::tool_def::<DeleteEntitiesTool>("delete_entities", "Execute delete_entities")
     }
 
-    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: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
         let mut missing = Vec::new();
@@ -255,10 +253,9 @@ impl McpTool for DeleteEntitiesHandler {
         });
 
         if !missing.is_empty() {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
-                missing.join(", ")
-            ));
+                missing.join(", "))));
         }
 
         let idx = state.get_search_index();
@@ -284,7 +281,7 @@ impl McpTool for DeleteObservationsHandler {
         )
     }
 
-    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: DeleteObservationsTool =
             serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut missing = Vec::new();
@@ -299,10 +296,9 @@ impl McpTool for DeleteObservationsHandler {
             }
         });
         if !missing.is_empty() {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
-                missing.join(", ")
-            ));
+                missing.join(", "))));
         }
         Ok("Observations deleted".to_string())
     }
@@ -320,7 +316,7 @@ impl McpTool for DeleteRelationsHandler {
         crate::mcp::tool_def::<DeleteRelationsTool>("delete_relations", "Execute delete_relations")
     }
 
-    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: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut missing_count = 0;
         state.modify_graph(|master| {
@@ -330,10 +326,10 @@ impl McpTool for DeleteRelationsHandler {
             missing_count = to_rem.len() - (initial_len - master.relations.len());
         });
         if missing_count > 0 {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
                 missing_count
-            ));
+            )));
         }
         Ok("Relations deleted".to_string())
     }
@@ -351,9 +347,9 @@ impl McpTool for ReadGraphHandler {
         crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph")
     }
 
-    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: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
-        let data = state.read_graph(|full| -> Result<String, String> {
+        let data = state.read_graph(|full| -> crate::error::Result<String> {
             if let Some(ns) = req.namespace {
                 let mut filtered = BorrowedGraph::default();
                 for (k, v) in &full.entities {
@@ -366,9 +362,9 @@ impl McpTool for ReadGraphHandler {
                         filtered.relations.push(r);
                     }
                 }
-                serde_json::to_string(&filtered).map_err(|e| e.to_string())
+                Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
             } else {
-                serde_json::to_string(full).map_err(|e| e.to_string())
+                Ok::<String, crate::error::AppError>(serde_json::to_string(full)?)
             }
         })?;
         Ok(data)
@@ -387,7 +383,7 @@ impl McpTool for SearchNodesHandler {
         crate::mcp::tool_def::<SearchNodesTool>("search_nodes", "Execute search_nodes")
     }
 
-    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: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let matches = if let Ok(idx) = state.search_index.read() {
             idx.search(&req.query, req.namespace.as_deref())
@@ -396,7 +392,7 @@ impl McpTool for SearchNodesHandler {
             vec![]
         };
 
-        let data = state.read_graph(|full| -> Result<String, String> {
+        let data = state.read_graph(|full| -> crate::error::Result<String> {
             let mut result = BorrowedGraph::default();
             for (id, doc_type, _, _, _) in &matches {
                 if doc_type == "entity"
@@ -405,7 +401,7 @@ impl McpTool for SearchNodesHandler {
                     result.entities.insert(id, e);
                 }
             }
-            serde_json::to_string(&result).map_err(|e| e.to_string())
+            Ok::<String, crate::error::AppError>(serde_json::to_string(&result)?)
         })?;
         Ok(data)
     }
@@ -423,9 +419,9 @@ impl McpTool for OpenNodesHandler {
         crate::mcp::tool_def::<OpenNodesTool>("open_nodes", "Execute open_nodes")
     }
 
-    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: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
-        let data = state.read_graph(|full| -> Result<String, String> {
+        let data = state.read_graph(|full| -> crate::error::Result<String> {
             let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
             let mut result = BorrowedGraph::default();
             let mut connected: HashSet<&str> = HashSet::new();
@@ -443,7 +439,7 @@ impl McpTool for OpenNodesHandler {
                     result.entities.insert(name, e);
                 }
             }
-            serde_json::to_string(&result).map_err(|e| e.to_string())
+            Ok::<String, crate::error::AppError>(serde_json::to_string(&result)?)
         })?;
         Ok(data)
     }
@@ -461,7 +457,7 @@ impl McpTool for VisualizeGraphHandler {
         crate::mcp::tool_def::<VisualizeGraphTool>("visualize_graph", "Execute visualize_graph")
     }
 
-    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: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let query = req.query.unwrap_or_default();
         let mut output = state.read_graph(|full| {
@@ -550,7 +546,7 @@ impl McpTool for CondenseEntityHandler {
         crate::mcp::tool_def::<CondenseEntityTool>("condense_entity", "Execute condense_entity")
     }
 
-    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: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut missing = false;
         state.modify_graph(|master| {
@@ -561,10 +557,10 @@ impl McpTool for CondenseEntityHandler {
             }
         });
         if missing {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
                 req.entity_name
-            ));
+            )));
         }
         Ok("Entity condensed".to_string())
     }
@@ -582,7 +578,7 @@ impl McpTool for MergeEntitiesHandler {
         crate::mcp::tool_def::<MergeEntitiesTool>("merge_entities", "Execute merge_entities")
     }
 
-    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: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         let mut missing = false;
         state.modify_graph(|master| {
@@ -616,10 +612,10 @@ impl McpTool for MergeEntitiesHandler {
             });
         });
         if missing {
-            return Err(format!(
+            return Err(crate::error::AppError::Internal(format!(
                 "Error: Source entity '{}' not found. Please verify the exact entity name using search_nodes.",
                 req.source_entity
-            ));
+            )));
         }
         Ok("Entities merged".to_string())
     }
@@ -637,7 +633,7 @@ impl McpTool for FindOrphansHandler {
         crate::mcp::tool_def::<FindOrphansTool>("find_orphans", "Execute find_orphans")
     }
 
-    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 orphans = state.read_graph(|full| {
             let mut connected = std::collections::HashSet::new();
             for r in &full.relations {
@@ -650,7 +646,7 @@ impl McpTool for FindOrphansHandler {
                 .cloned()
                 .collect::<Vec<String>>()
         });
-        let data = serde_json::to_string(&orphans).map_err(|e| e.to_string())?;
+        let data = Ok::<String, crate::error::AppError>(serde_json::to_string(&orphans)?)?;
         Ok(data)
     }
 }
@@ -675,7 +671,7 @@ mod tests {
             ]
         });
 
-        let res = create_handler.execute(args, state.clone()).await.unwrap();
+        let res = create_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res, "Entities created");
 
         // Ensure graph contains the entity
@@ -720,7 +716,7 @@ mod tests {
                 {"from": "A", "to": "B", "relation_type": "knows"}
             ]
         });
-        let res = handler.execute(args, state.clone()).await.unwrap();
+        let res = handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res, "Relations created");
 
         // Test semantic LLM schema feedback (User request)
@@ -759,28 +755,25 @@ mod tests {
                 {"entity_name": "A", "contents": ["Obs 1", "Obs 2"]}
             ]
         });
-        let res1 = add_obs.execute(args_obs, state.clone()).await.unwrap();
+        let res1 = add_obs.execute(args_obs, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res1, "Observations added");
 
         let read_graph = ReadGraphHandler;
         let res2 = read_graph
             .execute(json!({"namespace": "global"}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(res2.contains("Obs 1"));
         assert!(res2.contains("Obs 2"));
 
         let del_entity = DeleteEntitiesHandler;
         let res4 = del_entity
             .execute(json!({"entity_names": ["A"]}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res4, "Entities deleted");
 
         let res5 = read_graph
             .execute(json!({"namespace": "global"}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(!res5.contains("A"));
     }
 
@@ -798,8 +791,7 @@ mod tests {
         });
         create_handler
             .execute(args_ent, state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
 
         let rel_handler = CreateRelationsHandler;
         let args_rel = json!({
@@ -807,28 +799,25 @@ mod tests {
                 {"from": "X", "to": "Y", "relation_type": "depends_on", "namespace": "global"}
             ]
         });
-        rel_handler.execute(args_rel, state.clone()).await.unwrap();
+        rel_handler.execute(args_rel, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
 
         let read_handler = ReadGraphHandler;
         let res_read = read_handler
             .execute(json!({"namespace": "global"}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(res_read.contains("X"));
         assert!(res_read.contains("depends_on"));
 
         let open_handler = OpenNodesHandler;
         let res_open = open_handler
             .execute(json!({"names": ["X"]}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(res_open.contains("Y"));
 
         let viz_handler = VisualizeGraphHandler;
         let res_viz = viz_handler
             .execute(json!({"query": "X"}), state.clone())
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(!res_viz.is_empty());
 
         let condense = CondenseEntityHandler;
@@ -837,8 +826,7 @@ mod tests {
                 json!({"entity_name": "X", "summarized_observations": ["X condensed"]}),
                 state.clone(),
             )
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res_cond, "Entity condensed");
 
         let merge = MergeEntitiesHandler;
@@ -847,12 +835,11 @@ mod tests {
                 json!({"source_entity": "X", "target_entity": "Y"}),
                 state.clone(),
             )
-            .await
-            .unwrap();
+            .await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert_eq!(res_merge, "Entities merged");
 
         let orphans = FindOrphansHandler;
-        let res_orphans = orphans.execute(json!({}), state.clone()).await.unwrap();
+        let res_orphans = orphans.execute(json!({}), state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
         assert!(!res_orphans.contains("Y"));
     }
 }