mcp-memory/server/src/handlers_v2/tasks.rs

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 deps = req.dependencies.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: req.parent_id,
            dependencies: deps,
            acceptance_criteria: vec![],
        };
        let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
        let _ = 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;
        let mut actually_deleted = Vec::new();
        state.tasks.modify(|tasks| {
            let initial_len = tasks.len();
            
            // Build index-based children map
            let mut children_map: std::collections::HashMap<usize, Vec<usize>> = std::collections::HashMap::new();
            let mut id_to_index = std::collections::HashMap::new();
            for (idx, t) in tasks.iter().enumerate() {
                id_to_index.insert(t.id.as_str(), idx);
            }
            
            for (idx, t) in tasks.iter().enumerate() {
                if let Some(pid) = &t.parent_id {
                    if let Some(&parent_idx) = id_to_index.get(pid.as_str()) {
                        children_map.entry(parent_idx).or_default().push(idx);
                    }
                }
            }

            let mut to_delete_idx = std::collections::HashSet::new();
            if let Some(&start_idx) = id_to_index.get(req.id.as_str()) {
                let mut queue = std::collections::VecDeque::new();
                queue.push_back(start_idx);

                while let Some(curr) = queue.pop_front() {
                    if to_delete_idx.insert(curr) {
                        if let Some(children) = children_map.get(&curr) {
                            queue.extend(children.iter().copied());
                        }
                    }
                }
            }
            
            for &idx in &to_delete_idx {
                actually_deleted.push(tasks[idx].id.clone());
            }
            
            tasks.retain(|t| !actually_deleted.contains(&t.id));
            deleted_count = initial_len - tasks.len();
        });

        if deleted_count > 0 {
            let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
            for id in actually_deleted {
                let _ = idx.delete_document(&id);
            }
            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 target_idx = tasks.iter().position(|t| t.id == req.id || t.title == req.id);
            let target_idx = match target_idx {
                Some(idx) => idx,
                None => return,
            };
            
            found = true;
            let target_id = tasks[target_idx].id.clone();

            if target_status == "done" || target_status == "completed" {
                // 1. Check Acceptance Criteria
                if tasks[target_idx].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();
                    for dep_id in &tasks[target_idx].dependencies {
                        if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id) {
                            if dep_task.status != "completed" && dep_task.status != "done" {
                                uncompleted_deps.push(dep_task.title.as_str());
                            }
                        }
                    }
                    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.as_str());
                        }
                    }
                    if !uncompleted_children.is_empty() {
                        blocked = true;
                        blocker_details = format!("Blocked by child tasks: {}", uncompleted_children.join(", "));
                    }
                }
            }

            if !blocked {
                // Apply update
                tasks[target_idx].status = target_status.clone();
                tasks[target_idx].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<usize, Vec<usize>> =
                        std::collections::HashMap::new();
                    
                    // First pass: map string ID to index to build the adjacency list by index
                    let mut id_to_idx = std::collections::HashMap::new();
                    for (idx, t) in tasks.iter().enumerate() {
                        id_to_idx.insert(t.id.as_str(), idx);
                    }

                    for (idx, t) in tasks.iter().enumerate() {
                        if let Some(pid) = &t.parent_id
                            && let Some(&p_idx) = id_to_idx.get(pid.as_str()) {
                                children_map.entry(p_idx).or_default().push(idx);
                            }
                    }

                    if let Some(&start_idx) = id_to_idx.get(target_id.as_str()) {
                        let mut queue = std::collections::VecDeque::new();
                        queue.push_back(start_idx);

                        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(idx);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        });

        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 data = state.tasks.read_with(|tasks| {
            let filtered: Vec<_> = tasks.iter().filter(|t| {
                let status_match = t.status != "done";
                let branch_match = match &req.git_branch {
                    Some(branch) => t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str()),
                    None => true,
                };
                status_match && branch_match
            }).collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data)
    }
}

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("Acceptance criteria set successfully.".to_string())
        } 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(format!(
                "Acceptance criteria verified with proof: {}",
                req.proof
            ))
        } else if already_met {
            Ok("Acceptance criteria was already met.".to_string())
        } else {
            Ok("Acceptance criteria or task not found.".to_string())
        }
    }
}

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 data = state.milestones.read_with(|items| {
            let filtered: Vec<_> = items.iter().filter(|i| {
                if let Some(ns) = &req.namespace {
                    &i.namespace == ns
                } else {
                    true
                }
            }).collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data)
    }
}