refactor: apply zero-unwrap policy and optimize locks in store.rs and handlers

server/src/handlers_v2/tasks.rs

@@ -0,0 +1,456 @@
+use crate::models::*;
+use crate::router::McpTool;
+use crate::state::MemoryState;
+use crate::tools::*;
+use async_trait::async_trait;
+use serde_json::Value;
+use std::sync::Arc;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+pub struct AddTaskHandler;
+
+#[async_trait]
+impl McpTool for AddTaskHandler {
+    fn name(&self) -> &'static str {
+        "add_task"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<AddTaskTool>("add_task", "Execute add_task")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: AddTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let now = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+        let task_id = uuid::Uuid::new_v4().to_string();
+
+        let parent_id = req.parent_id.clone();
+        let deps = req.dependencies.clone().unwrap_or_default();
+
+        let task = Task {
+            id: task_id.clone(),
+            title: req.title,
+            status: "pending".to_string(),
+            description: req.description,
+            created_at: now,
+            updated_at: now,
+            git_branch: req.git_branch,
+            parent_id,
+            dependencies: deps,
+            acceptance_criteria: vec![],
+        };
+        if let Ok(idx) = state.search_index.read() {
+            drop(idx.index_task(&task));
+        }
+        state.tasks.modify(|tasks| {
+            tasks.push(task);
+        });
+        Ok(format!("Task added with ID: {}", task_id).to_string())
+    }
+}
+
+pub struct DeleteTaskHandler;
+
+#[async_trait]
+impl McpTool for DeleteTaskHandler {
+    fn name(&self) -> &'static str {
+        "delete_task"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<DeleteTaskTool>("delete_task", "Execute delete_task")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: DeleteTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut deleted_count = 0;
+        state.tasks.modify(|tasks| {
+            let initial_len = tasks.len();
+            // Collect IDs of tasks to delete (this task + all its recursive children)
+            let mut to_delete = std::collections::HashSet::new();
+            to_delete.insert(req.id.clone());
+
+            let mut children_map: std::collections::HashMap<String, Vec<String>> =
+                std::collections::HashMap::new();
+            for t in tasks.iter() {
+                if let Some(pid) = &t.parent_id {
+                    children_map
+                        .entry(pid.clone())
+                        .or_default()
+                        .push(t.id.clone());
+                }
+            }
+
+            let mut queue = std::collections::VecDeque::new();
+            queue.push_back(req.id.clone());
+
+            while let Some(curr) = queue.pop_front() {
+                if to_delete.insert(curr.clone())
+                    && let Some(children) = children_map.get(&curr)
+                {
+                    queue.extend(children.iter().cloned());
+                }
+            }
+
+            tasks.retain(|t| !to_delete.contains(&t.id));
+            deleted_count = initial_len - tasks.len();
+        });
+
+        if deleted_count > 0 {
+            Ok(vec![
+                format!("Deleted task and its children ({} total).", deleted_count).to_string(),
+            ][0]
+            .clone())
+        } else {
+            Ok("Task not found.".to_string())
+        }
+    }
+}
+
+pub struct UpdateTaskStatusHandler;
+
+#[async_trait]
+impl McpTool for UpdateTaskStatusHandler {
+    fn name(&self) -> &'static str {
+        "update_task_status"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<UpdateTaskStatusTool>(
+            "update_task_status",
+            "Execute update_task_status",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: UpdateTaskStatusTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut found = false;
+        let mut blocked = false;
+        let mut blocker_details = String::new();
+        let target_status = req.status.to_lowercase();
+
+        state.tasks.modify(|tasks| {
+            // Find target task
+            let mut target_id = String::new();
+            if let Some(t) = tasks.iter().find(|t| t.id == req.id || t.title == req.id) {
+                target_id = t.id.clone();
+            }
+
+            if target_id.is_empty() {
+                return;
+            }
+            found = true;
+
+            if target_status == "done" || target_status == "completed" {
+                // 1. Check Acceptance Criteria
+                if let Some(t) = tasks.iter().find(|t| t.id == target_id)
+                    && t.acceptance_criteria.iter().any(|c| !c.is_met)
+                {
+                    blocked = true;
+                    blocker_details = "Unmet acceptance criteria exist.".to_string();
+                }
+
+                // 2. Check dependencies
+                if !blocked {
+                    let mut uncompleted_deps = Vec::new();
+                    if let Some(t) = tasks.iter().find(|t| t.id == target_id) {
+                        for dep_id in &t.dependencies {
+                            if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id)
+                                && dep_task.status != "completed"
+                                && dep_task.status != "done"
+                            {
+                                uncompleted_deps.push(dep_task.title.clone());
+                            }
+                        }
+                    }
+                    if !uncompleted_deps.is_empty() {
+                        blocked = true;
+                        blocker_details =
+                            format!("Blocked by dependencies: {}", uncompleted_deps.join(", "));
+                    }
+                }
+
+                // 3. Check child tasks
+                if !blocked {
+                    let mut uncompleted_children = Vec::new();
+                    for child in tasks
+                        .iter()
+                        .filter(|t| t.parent_id.as_ref() == Some(&target_id))
+                    {
+                        if child.status != "completed" && child.status != "done" {
+                            uncompleted_children.push(child.title.clone());
+                        }
+                    }
+                    if !uncompleted_children.is_empty() {
+                        blocked = true;
+                        blocker_details = format!(
+                            "Blocked by child tasks: {}",
+                            uncompleted_children.join(", ")
+                        );
+                    }
+                }
+            }
+
+            if !blocked {
+                // Apply update
+                if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) {
+                    t.status = target_status.clone();
+                    t.updated_at = SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .unwrap_or_default()
+                        .as_secs();
+                }
+
+                // Cascade cancellation to children
+                if target_status == "cancelled" || target_status == "abandoned" {
+                    let mut children_map: std::collections::HashMap<String, Vec<usize>> =
+                        std::collections::HashMap::new();
+                    for (idx, t) in tasks.iter().enumerate() {
+                        if let Some(pid) = &t.parent_id {
+                            children_map.entry(pid.clone()).or_default().push(idx);
+                        }
+                    }
+
+                    let mut queue = std::collections::VecDeque::new();
+                    queue.push_back(target_id.clone());
+
+                    while let Some(curr) = queue.pop_front() {
+                        if let Some(child_indices) = children_map.get(&curr) {
+                            for &idx in child_indices {
+                                if tasks[idx].status != "completed"
+                                    && tasks[idx].status != target_status
+                                {
+                                    tasks[idx].status = target_status.clone();
+                                    queue.push_back(tasks[idx].id.clone());
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        });
+
+        if blocked {
+            Ok(format!(
+                "Error: Cannot transition task. {}",
+                blocker_details
+            ))
+        } else if found {
+            Ok("Task status updated.".to_string())
+        } else {
+            Ok("Task not found.".to_string())
+        }
+    }
+}
+
+pub struct ListActiveTasksHandler;
+
+#[async_trait]
+impl McpTool for ListActiveTasksHandler {
+    fn name(&self) -> &'static str {
+        "list_active_tasks"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<ListActiveTasksTool>(
+            "list_active_tasks",
+            "Execute list_active_tasks",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: ListActiveTasksTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut tasks = state.tasks.read();
+        tasks.retain(|t| t.status != "done");
+        if let Some(branch) = req.git_branch {
+            tasks.retain(|t| {
+                t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str())
+            });
+        }
+        let data = serde_json::to_string(&tasks).unwrap_or_default();
+        Ok(data.to_string())
+    }
+}
+
+pub struct SetAcceptanceCriteriaHandler;
+
+#[async_trait]
+impl McpTool for SetAcceptanceCriteriaHandler {
+    fn name(&self) -> &'static str {
+        "set_acceptance_criteria"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<SetAcceptanceCriteriaTool>(
+            "set_acceptance_criteria",
+            "Execute set_acceptance_criteria",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: SetAcceptanceCriteriaTool =
+            serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut success = false;
+        state.tasks.modify(|tasks| {
+            if let Some(task) = tasks.iter_mut().rev().find(|t| t.title == req.task_title) {
+                task.acceptance_criteria = req
+                    .criteria
+                    .into_iter()
+                    .map(|desc| crate::models::AcceptanceCriteria {
+                        id: uuid::Uuid::new_v4().to_string(),
+                        description: desc,
+                        is_met: false,
+                    })
+                    .collect();
+                task.updated_at = SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap_or_default()
+                    .as_secs();
+                success = true;
+            }
+        });
+        if success {
+            Ok(vec!["Acceptance criteria set successfully.".to_string()][0].clone())
+        } else {
+            Ok("Task not found.".to_string())
+        }
+    }
+}
+
+pub struct VerifyAcceptanceCriteriaHandler;
+
+#[async_trait]
+impl McpTool for VerifyAcceptanceCriteriaHandler {
+    fn name(&self) -> &'static str {
+        "verify_acceptance_criteria"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<VerifyAcceptanceCriteriaTool>(
+            "verify_acceptance_criteria",
+            "Execute verify_acceptance_criteria",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: VerifyAcceptanceCriteriaTool =
+            serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut success = false;
+        let mut already_met = false;
+        state.tasks.modify(|tasks| {
+            if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id)
+                && let Some(ac) = task
+                    .acceptance_criteria
+                    .iter_mut()
+                    .find(|c| c.id == req.criteria || c.description == req.criteria)
+            {
+                if ac.is_met {
+                    already_met = true;
+                } else {
+                    ac.is_met = true;
+                    success = true;
+                    task.updated_at = SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .unwrap_or_default()
+                        .as_secs();
+                }
+            }
+        });
+        if success {
+            Ok(vec![format!(
+                "Acceptance criteria verified with proof: {}",
+                req.proof
+            )][0]
+                .clone())
+        } else if already_met {
+            Ok("Acceptance criteria was already met.".to_string())
+        } else {
+            Ok(vec!["Acceptance criteria or task not found.".to_string()][0].clone())
+        }
+    }
+}
+
+pub struct AddMilestoneHandler;
+
+#[async_trait]
+impl McpTool for AddMilestoneHandler {
+    fn name(&self) -> &'static str {
+        "add_milestone"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<AddMilestoneTool>("add_milestone", "Execute add_milestone")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: AddMilestoneTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.milestones.modify(|ms| {
+            ms.push(crate::models::Milestone {
+                id: uuid::Uuid::new_v4().to_string(),
+                title: req.title,
+                status: "pending".to_string(),
+                namespace: req.namespace,
+                target_date: None,
+            })
+        });
+        Ok("Milestone added".to_string())
+    }
+}
+
+pub struct UpdateMilestoneHandler;
+
+#[async_trait]
+impl McpTool for UpdateMilestoneHandler {
+    fn name(&self) -> &'static str {
+        "update_milestone"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<UpdateMilestoneTool>("update_milestone", "Execute update_milestone")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: UpdateMilestoneTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut found = false;
+        state.milestones.modify(|ms| {
+            for m in ms.iter_mut() {
+                if m.id == req.id {
+                    m.status = req.status.clone();
+                    found = true;
+                    break;
+                }
+            }
+        });
+        if found {
+            Ok("Milestone updated".to_string())
+        } else {
+            Ok("Milestone not found".to_string())
+        }
+    }
+}
+
+pub struct ListMilestonesHandler;
+
+#[async_trait]
+impl McpTool for ListMilestonesHandler {
+    fn name(&self) -> &'static str {
+        "list_milestones"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<ListMilestonesTool>("list_milestones", "Execute list_milestones")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: ListMilestonesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut items = state.milestones.read();
+        if let Some(ns) = req.namespace {
+            items.retain(|i| i.namespace == ns);
+        }
+        let data = serde_json::to_string(&items).unwrap_or_default();
+        Ok(data.to_string())
+    }
+}