use std::sync::Arc;
use pe_core::PeError;
use serde_json::{Value, json};
use crate::dependency::{DependencyType, TaskDependency};
use crate::registry::{TaskFilter, TaskRegistry};
use crate::task::{Task, TaskPriority, TaskStatus, TaskType};
pub async fn task_create(
registry: &Arc<dyn TaskRegistry>,
agent_id: &str,
args: Value,
) -> Result<Value, PeError> {
let title = args["title"]
.as_str()
.ok_or_else(|| PeError::InvalidUpdate {
details: "task_create requires 'title' field".into(),
})?;
let mut task = Task::new(title);
task.task_type = TaskType::Agent;
task.agent_id = Some(agent_id.to_string());
task.created_by = agent_id.to_string();
if let Some(desc) = args["description"].as_str() {
task.description = desc.to_string();
}
if let Some(p) = args["priority"].as_str() {
task.priority = match p {
"urgent" => TaskPriority::Urgent,
"high" => TaskPriority::High,
"low" => TaskPriority::Low,
_ => TaskPriority::Medium,
};
}
if let Some(tags) = args["tags"].as_array() {
task.tags = tags
.iter()
.filter_map(|t| t.as_str().map(String::from))
.collect();
}
if let Some(parent) = args["parent_task_id"].as_str() {
task.parent_task_id = Some(parent.to_string());
}
if let Some(assignee) = args["assignee"].as_str() {
task.assignee = assignee.to_string();
}
let created = registry.create(task).await?;
Ok(json!({
"task_id": created.id,
"title": created.title,
"status": format!("{:?}", created.status),
}))
}
pub async fn task_list(registry: &Arc<dyn TaskRegistry>, args: Value) -> Result<Value, PeError> {
let mut filter = TaskFilter::default();
if let Some(s) = args["status"].as_str() {
filter.status = Some(match s {
"pending" => TaskStatus::Pending,
"in_progress" => TaskStatus::InProgress,
"completed" => TaskStatus::Completed,
"failed" => TaskStatus::Failed,
"blocked" => TaskStatus::Blocked,
"cancelled" => TaskStatus::Cancelled,
other => {
return Err(PeError::InvalidUpdate {
details: format!(
"Unknown status '{other}'. Valid: pending, in_progress, completed, failed, blocked, cancelled"
),
});
}
});
}
if let Some(a) = args["agent_id"].as_str() {
filter.agent_id = Some(a.to_string());
}
if let Some(a) = args["assignee"].as_str() {
filter.assignee = Some(a.to_string());
}
if let Some(p) = args["parent_id"].as_str() {
filter.parent_id = Some(p.to_string());
}
if let Some(t) = args["tag"].as_str() {
filter.tag = Some(t.to_string());
}
if let Some(l) = args["limit"].as_u64() {
filter.limit = l as usize;
}
let tasks = registry.list(&filter).await?;
let mut items: Vec<Value> = Vec::with_capacity(tasks.len());
for t in &tasks {
let dependencies = registry.get_dependencies(&t.id).await?;
let dependency_items = dependencies
.iter()
.map(|dep| {
json!({
"task_id": dep.task_id,
"depends_on_id": dep.depends_on_id,
"dependency_type": format!("{:?}", dep.dependency_type),
"created_at": dep.created_at.to_rfc3339(),
})
})
.collect::<Vec<_>>();
items.push(json!({
"id": t.id,
"title": t.title,
"status": format!("{:?}", t.status),
"priority": format!("{:?}", t.priority),
"assignee": t.assignee,
"parent_task_id": t.parent_task_id,
"dependencies": dependency_items,
}));
}
let count = items.len();
Ok(json!({ "tasks": items, "count": count }))
}
pub async fn task_complete(
registry: &Arc<dyn TaskRegistry>,
args: Value,
) -> Result<Value, PeError> {
let task_id = args["task_id"]
.as_str()
.ok_or_else(|| PeError::InvalidUpdate {
details: "task_complete requires 'task_id' field".into(),
})?;
let task = registry
.get(&task_id.to_string())
.await?
.ok_or_else(|| PeError::NodeNotFound {
node: task_id.to_string(),
})?;
match &task.status {
TaskStatus::Pending => {
registry
.update_status(&task_id.to_string(), TaskStatus::InProgress, None, None)
.await?;
}
TaskStatus::InProgress => {} TaskStatus::Blocked | TaskStatus::Failed => {
return Err(PeError::InvalidUpdate {
details: format!("Task is {:?}. Resolve it before completing.", task.status),
});
}
TaskStatus::Completed | TaskStatus::Cancelled => {
return Err(PeError::InvalidUpdate {
details: format!("Task is already {:?}.", task.status),
});
}
}
let result = args.get("result").cloned();
let completed = registry
.update_status(&task_id.to_string(), TaskStatus::Completed, result, None)
.await?;
Ok(json!({
"task_id": completed.id,
"status": "Completed",
"completed_at": completed.completed_at.map(|t| t.to_rfc3339()),
}))
}
pub async fn task_decompose(
registry: &Arc<dyn TaskRegistry>,
agent_id: &str,
args: Value,
) -> Result<Value, PeError> {
let parent_id = args["task_id"]
.as_str()
.ok_or_else(|| PeError::InvalidUpdate {
details: "task_decompose requires 'task_id' field".into(),
})?;
let subtasks_arr = args["subtasks"]
.as_array()
.ok_or_else(|| PeError::InvalidUpdate {
details: "task_decompose requires 'subtasks' array".into(),
})?;
let parent =
registry
.get(&parent_id.to_string())
.await?
.ok_or_else(|| PeError::NodeNotFound {
node: parent_id.to_string(),
})?;
if parent.status == TaskStatus::Pending {
registry
.update_status(&parent_id.to_string(), TaskStatus::InProgress, None, None)
.await?;
}
let mut created_ids: Vec<String> = Vec::new();
for sub_def in subtasks_arr {
let title = sub_def["title"].as_str().unwrap_or("Untitled subtask");
let mut sub = Task::agent_task(title, agent_id).with_parent(parent_id);
if let Some(desc) = sub_def["description"].as_str() {
sub.description = desc.to_string();
}
if let Some(p) = sub_def["priority"].as_str() {
sub.priority = match p {
"urgent" => TaskPriority::Urgent,
"high" => TaskPriority::High,
"low" => TaskPriority::Low,
_ => TaskPriority::Medium,
};
}
let created = registry.create(sub).await?;
created_ids.push(created.id);
}
let mut deps_created = 0;
for (i, sub_def) in subtasks_arr.iter().enumerate() {
if let Some(dep_indices) = sub_def["depends_on"].as_array() {
for dep_idx in dep_indices {
if let Some(idx) = dep_idx.as_u64() {
let idx = idx as usize;
if idx < created_ids.len() && idx != i {
let dep = TaskDependency::new(
&created_ids[i],
&created_ids[idx],
DependencyType::Blocks,
);
registry.add_dependency(dep).await?;
deps_created += 1;
}
}
}
}
}
Ok(json!({
"parent_id": parent_id,
"subtask_ids": created_ids,
"subtask_count": created_ids.len(),
"dependencies_created": deps_created,
}))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::InMemoryTaskRegistry;
fn reg() -> Arc<dyn TaskRegistry> {
Arc::new(InMemoryTaskRegistry::new())
}
#[tokio::test]
async fn test_task_create_tool() {
let r = reg();
let result = task_create(&r, "agent-1", json!({ "title": "Write tests" }))
.await
.unwrap();
assert_eq!(result["title"], "Write tests");
assert!(!result["task_id"].as_str().unwrap().is_empty());
}
#[tokio::test]
async fn test_task_list_tool() {
let r = reg();
task_create(&r, "a1", json!({ "title": "Task 1" }))
.await
.unwrap();
task_create(&r, "a1", json!({ "title": "Task 2" }))
.await
.unwrap();
let result = task_list(&r, json!({})).await.unwrap();
assert_eq!(result["count"], 2);
}
#[tokio::test]
async fn test_task_complete_tool() {
let r = reg();
let created = task_create(&r, "a1", json!({ "title": "Finish me" }))
.await
.unwrap();
let tid = created["task_id"].as_str().unwrap();
let result = task_complete(&r, json!({ "task_id": tid })).await.unwrap();
assert_eq!(result["status"], "Completed");
}
#[tokio::test]
async fn test_task_decompose_tool() {
let r = reg();
let parent = task_create(&r, "a1", json!({ "title": "Big project" }))
.await
.unwrap();
let pid = parent["task_id"].as_str().unwrap();
let result = task_decompose(
&r,
"a1",
json!({
"task_id": pid,
"subtasks": [
{ "title": "Write code" },
{ "title": "Write tests", "depends_on": [0] },
{ "title": "Deploy", "depends_on": [0, 1] },
]
}),
)
.await
.unwrap();
assert_eq!(result["subtask_count"], 3);
assert_eq!(result["dependencies_created"], 3);
}
}