Files
mcp-memory/server/src/handlers/tasks.rs
T

740 lines
32 KiB
Rust

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;
pub struct TasksHandler;
#[async_trait]
impl McpTool for TasksHandler {
fn name(&self) -> &'static str {
"tasks"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<TasksTool>(
"tasks",
"Consolidated task management (add, update, delete, list, set_criteria, verify)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: TasksTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
TaskAction::Add => {
let title = req.title.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'title' for action 'add'. Next step: Provide non-empty 'title' string in request and retry.".to_string())
})?;
let description = req.description.unwrap_or_default();
let now = crate::handlers::utils::now_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,
status: "pending".to_string(),
description,
created_at: now,
updated_at: now,
git_branch: req.git_branch,
parent_id: req.parent_id,
dependencies: deps,
acceptance_criteria: vec![],
expires_at: None,
namespace: req.namespace.unwrap_or_else(|| "default".to_string()),
repo_name: req.repo_name,
priority: req.priority,
assigned_agent: req.assigned_agent,
verification_command: req.verification_command,
..Default::default()
};
let idx = state.get_search_index().await;
drop(idx.index_task(&task));
state.project.tasks.modify(|tasks| {
tasks.push(task.clone());
});
state.record_activity(
"task_create",
&format!("Created task: {}", task.title),
Some(&task.description),
);
state.broadcast_task_event(TaskEvent {
task_id: task_id.clone(),
status: "created".to_string(),
action: Some("add".to_string()),
result: Some(
serde_json::json!({ "title": task.title, "git_branch": task.git_branch }),
),
error: None,
timestamp: now,
session_id: None,
..Default::default()
});
Ok(format!("Task added with ID: {}", task_id))
}
TaskAction::Update => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'update'. Next step: Provide task 'id' string in request and retry.".to_string())
})?;
let status = req.status.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'status' for action 'update'. Next step: Provide valid 'status' ('pending', 'active', 'in_progress', 'blocked', 'completed', or 'cancelled') in request and retry.".to_string())
})?;
let target_status = status.to_lowercase();
let mut found = false;
let mut blocked = false;
let mut blocker_details = String::new();
let mut updated_task: Option<Task> = None;
state.project.tasks.modify(|tasks| {
let target_idx = tasks.iter().position(|t| t.id == id || t.title == id);
let target_idx = match target_idx {
Some(i) => i,
None => return,
};
found = true;
if target_status == "done" || target_status == "completed" {
if tasks[target_idx]
.acceptance_criteria
.iter()
.any(|c| !c.is_met)
{
blocked = true;
blocker_details = "Unmet acceptance criteria exist.".to_string();
}
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)
&& dep_task.status != "completed"
&& dep_task.status != "done"
&& dep_task.status != "cancelled"
{
uncompleted_deps.push(dep_task.title.as_str());
}
}
if !uncompleted_deps.is_empty() {
blocked = true;
blocker_details = format!(
"Blocked by dependencies: {}",
uncompleted_deps.join(", ")
);
}
}
if !blocked {
let target_id_ref = tasks[target_idx].id.as_str();
let mut uncompleted_children = Vec::new();
for child in tasks
.iter()
.filter(|t| t.parent_id.as_deref() == Some(target_id_ref))
{
if child.status != "completed"
&& child.status != "done"
&& child.status != "cancelled"
{
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 {
tasks[target_idx].status = target_status.clone();
tasks[target_idx].updated_at = crate::handlers::utils::now_secs();
updated_task = Some(tasks[target_idx].clone());
}
});
if blocked {
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: "blocked".to_string(),
action: Some("update".to_string()),
result: None,
error: Some(blocker_details.clone()),
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Err(crate::error::AppError::Internal(format!(
"Error: Cannot transition task. {}",
blocker_details
)))
} else if found {
if let Some(ref ut) = updated_task {
let idx = state.get_search_index().await;
drop(idx.index_task(ut));
}
state.record_activity(
"task_update",
&format!("Task {} status -> {}", id, status),
None,
);
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: target_status.clone(),
action: Some("update".to_string()),
result: Some(serde_json::json!({ "status": target_status })),
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
if target_status == "completed" {
crate::handlers::reconciliation::cascade_unblock_tasks(&state);
crate::handlers::reconciliation::recalculate_milestones(&state);
}
Ok("Task status updated.".to_string())
} else {
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::Delete => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'delete'. Next step: Provide task 'id' string in request and retry.".to_string())
})?;
let mut deleted_count = 0;
let mut actually_deleted = Vec::new();
state.project.tasks.modify(|tasks| {
let initial_len = tasks.len();
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);
}
let mut children_map: std::collections::HashMap<usize, Vec<usize>> =
std::collections::HashMap::new();
for (idx, t) in tasks.iter().enumerate() {
if let Some(pid) = &t.parent_id
&& let Some(&p_idx) = id_to_index.get(pid.as_str())
{
children_map.entry(p_idx).or_default().push(idx);
}
}
let mut to_delete_idx = std::collections::HashSet::new();
if let Some(&start_idx) = id_to_index.get(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)
&& 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.get_search_index().await;
for deleted_id in actually_deleted {
drop(idx.delete_document(&deleted_id));
}
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: "deleted".to_string(),
action: Some("delete".to_string()),
result: Some(serde_json::json!({ "deleted_count": deleted_count })),
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok(format!(
"Deleted task and its children ({} total).",
deleted_count
))
} else {
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::List => {
let level = req.summary_level.as_deref().unwrap_or("detailed");
let offset = req.offset.unwrap_or(0);
let data = state.project.tasks.read_with(|tasks| {
let matching: Vec<_> = tasks
.iter()
.filter(|t| {
let status_match = match req.status.as_deref() {
Some("all") => true,
Some(s) => t.status.eq_ignore_ascii_case(s),
None => t.is_active(),
};
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();
let mut filtered: Vec<_> = matching
.into_iter()
.skip(offset)
.map(|t| match level {
"compact" => serde_json::json!({ "id": t.id, "title": t.title, "status": t.status }),
"full" => serde_json::to_value(t).unwrap_or_default(),
_ => serde_json::json!({ "id": t.id, "title": t.title, "status": t.status, "description": t.description, "git_branch": t.git_branch }),
})
.collect();
let mut json_str = serde_json::to_string(&filtered)?;
if let Some(max_t) = req.max_tokens {
let char_limit = max_t * 4;
while json_str.len() > char_limit && !filtered.is_empty() {
filtered.pop();
json_str = serde_json::to_string(&filtered)?;
}
}
Ok::<String, crate::error::AppError>(json_str)
})?;
Ok(data)
}
TaskAction::SetCriteria => {
let id = req.id.or(req.title.clone()).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' or 'title' for action 'set_criteria'. Next step: Provide task 'id' string in request and retry.".to_string())
})?;
let criteria_list = req.criteria.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'criteria' for action 'set_criteria'. Next step: Provide array of acceptance criteria strings in request and retry.".to_string())
})?;
let mut success = false;
let mut updated_task: Option<Task> = None;
state.project.tasks.modify(|tasks| {
if let Some(task) = tasks.iter_mut().find(|t| t.id == id || t.title == id) {
task.acceptance_criteria = criteria_list
.into_iter()
.map(|desc| crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(),
description: desc,
is_met: false,
..Default::default()
})
.collect();
task.updated_at = crate::handlers::utils::now_secs();
updated_task = Some(task.clone());
success = true;
}
});
if success {
if let Some(ref ut) = updated_task {
let idx = state.get_search_index().await;
drop(idx.index_task(ut));
}
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: "criteria_set".to_string(),
action: Some("set_criteria".to_string()),
result: Some(serde_json::json!({ "id": id })),
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok("Acceptance criteria set successfully.".to_string())
} else {
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::Verify => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'verify'. Next step: Provide task 'id' string in request and retry.".to_string())
})?;
let proof_str = req.proof.unwrap_or_else(|| "Verified".to_string());
let mut success = false;
let mut already_met = false;
let mut updated_task: Option<Task> = None;
state.project.tasks.modify(|tasks| {
if let Some(task) = tasks.iter_mut().find(|t| t.id == id || t.title == id) {
if let Some(ac) = task.acceptance_criteria.first_mut() {
if ac.is_met {
already_met = true;
} else {
ac.is_met = true;
success = true;
task.updated_at = crate::handlers::utils::now_secs();
updated_task = Some(task.clone());
}
} else {
task.acceptance_criteria
.push(crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(),
description: proof_str.clone(),
is_met: true,
..Default::default()
});
task.updated_at = crate::handlers::utils::now_secs();
updated_task = Some(task.clone());
success = true;
}
}
});
if success {
if let Some(ref ut) = updated_task {
let idx = state.get_search_index().await;
drop(idx.index_task(ut));
}
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: "verified".to_string(),
action: Some("verify".to_string()),
result: Some(serde_json::json!({ "proof": proof_str })),
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok(format!(
"Acceptance criteria verified with proof: {}",
proof_str
))
} else if already_met {
Ok("Acceptance criteria was already met.".to_string())
} else {
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
}
}
}
pub struct MilestonesHandler;
#[async_trait]
impl McpTool for MilestonesHandler {
fn name(&self) -> &'static str {
"milestones"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<MilestonesTool>(
"milestones",
"Consolidated milestone management (add, update, list)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: MilestonesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
MilestoneAction::Add => {
let title = req.title.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'title' for action 'add'. Next step: Provide non-empty 'title' string in request and retry.".to_string())
})?;
let ns = req
.namespace
.unwrap_or_else(crate::models::default_namespace);
state.project.milestones.modify(|ms| {
ms.push(crate::models::Milestone {
id: uuid::Uuid::new_v4().to_string(),
title,
status: "pending".to_string(),
namespace: ns,
target_date: req
.target_date
.as_deref()
.and_then(|s| s.parse::<u64>().ok()),
description: req.description,
task_ids: req.task_ids.unwrap_or_default(),
deliverables: req.deliverables.unwrap_or_default(),
repo_name: req.repo_name,
..Default::default()
})
});
crate::handlers::reconciliation::recalculate_milestones(&state);
Ok("Milestone added".to_string())
}
MilestoneAction::Update => {
let id = req.id.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' for action 'update'. Next step: Provide milestone 'id' string in request and retry.".to_string())
})?;
let status = req.status.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'status' for action 'update'. Next step: Provide milestone 'status' in request and retry.".to_string())
})?;
let mut found = false;
state.project.milestones.modify(|ms| {
for m in ms.iter_mut() {
if m.id == id {
m.status = status.clone();
if (status.eq_ignore_ascii_case("completed")
|| status.eq_ignore_ascii_case("done"))
&& m.completed_at.is_none() {
m.completed_at = Some(crate::handlers::utils::now_secs());
}
if let Some(task_ids) = &req.task_ids {
m.task_ids = task_ids.clone();
if let Some(deliverables) = &req.deliverables {
m.deliverables = deliverables.clone();
}
}
found = true;
break;
}
}
});
if found {
crate::handlers::reconciliation::recalculate_milestones(&state);
Ok("Milestone updated".to_string())
} else {
Err(crate::error::AppError::Internal(
"Milestone not found. Please verify the milestone ID.".to_string(),
))
}
}
MilestoneAction::List => {
let data = state.project.milestones.read_with(|items| {
let filtered: Vec<_> = items
.iter()
.filter(|i| {
if let Some(ns) = &req.namespace {
&i.namespace == ns
} else {
true
}
})
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use tempfile::tempdir;
#[tokio::test]
async fn test_add_task_and_list() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = TasksHandler;
let args = json!({
"action": "add",
"title": "Fix the hyperdrive",
"description": "It's making a strange noise",
});
let res = handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("Task added with ID:"));
let res2 = handler
.execute(json!({"action": "list"}), state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Fix the hyperdrive"));
}
#[tokio::test]
async fn test_update_task_status() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = TasksHandler;
let res = handler
.execute(
json!({"action": "add", "title": "Test", "description": "test"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let id_start = res.find("ID: ").unwrap() + 4;
let task_id = res[id_start..].trim();
let args = json!({
"action": "update",
"id": task_id,
"status": "done"
});
let res3 = handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res3, "Task status updated.");
let res4 = handler
.execute(json!({"action": "list"}), state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res4.contains(task_id));
}
#[tokio::test]
async fn test_milestones_and_criteria() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler_ms = MilestonesHandler;
let args_ms = json!({
"action": "add",
"title": "Release 1.0",
"namespace": "global"
});
let res1 = handler_ms
.execute(args_ms, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res1.contains("Milestone added"));
let ms_id = state.project.milestones.read_with(|ms| ms[0].id.clone());
let args_ums = json!({
"action": "update",
"id": ms_id,
"status": "completed"
});
let res2 = handler_ms
.execute(args_ums, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res2, "Milestone updated");
let res3 = handler_ms
.execute(
json!({"action": "list", "namespace": "global"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res3.contains("completed"));
assert!(res3.contains("Release 1.0"));
let handler_t = TasksHandler;
let res_task = handler_t
.execute(
json!({"action": "add", "title": "Test", "description": "desc"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let task_id = res_task[res_task.find("ID: ").unwrap() + 4..].trim();
let args_ac = json!({
"action": "set_criteria",
"id": task_id,
"criteria": ["Do X", "Do Y"]
});
let res4 = handler_t
.execute(args_ac, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res4, "Acceptance criteria set successfully.");
let args_vac = json!({
"action": "verify",
"id": task_id,
"proof": "I did X"
});
let res5 = handler_t
.execute(args_vac, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res5.contains("Acceptance criteria verified"));
}
#[tokio::test]
async fn test_delete_task_with_children() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = TasksHandler;
let parent = handler
.execute(
json!({"action": "add", "title": "Parent", "description": "p"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let parent_id = parent[parent.find("ID: ").unwrap() + 4..]
.trim()
.to_string();
let child = handler
.execute(
json!({"action": "add", "title": "Child", "description": "c", "parent_id": parent_id}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let _child_id = child[child.find("ID: ").unwrap() + 4..].trim().to_string();
let res_del = handler
.execute(json!({"action": "delete", "id": parent_id}), state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res_del.contains("Deleted task and its children (2 total)."));
}
#[tokio::test]
async fn test_list_milestones_with_namespace() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = MilestonesHandler;
let args_ms = serde_json::json!({
"action": "add",
"title": "Release 1.0",
"namespace": "global"
});
let res1 = handler
.execute(args_ms, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res1.contains("Milestone added"));
let res2 = handler
.execute(
serde_json::json!({"action": "list", "namespace": "global"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Release 1.0"));
}
}