1use std::sync::Arc;
7
8use pe_core::PeError;
9use serde_json::{Value, json};
10
11use crate::dependency::{DependencyType, TaskDependency};
12use crate::registry::{TaskFilter, TaskRegistry};
13use crate::task::{Task, TaskPriority, TaskStatus, TaskType};
14
15pub async fn task_create(
20 registry: &Arc<dyn TaskRegistry>,
21 agent_id: &str,
22 args: Value,
23) -> Result<Value, PeError> {
24 let title = args["title"]
25 .as_str()
26 .ok_or_else(|| PeError::InvalidUpdate {
27 details: "task_create requires 'title' field".into(),
28 })?;
29
30 let mut task = Task::new(title);
31 task.task_type = TaskType::Agent;
32 task.agent_id = Some(agent_id.to_string());
33 task.created_by = agent_id.to_string();
34
35 if let Some(desc) = args["description"].as_str() {
36 task.description = desc.to_string();
37 }
38 if let Some(p) = args["priority"].as_str() {
39 task.priority = match p {
40 "urgent" => TaskPriority::Urgent,
41 "high" => TaskPriority::High,
42 "low" => TaskPriority::Low,
43 _ => TaskPriority::Medium,
44 };
45 }
46 if let Some(tags) = args["tags"].as_array() {
47 task.tags = tags
48 .iter()
49 .filter_map(|t| t.as_str().map(String::from))
50 .collect();
51 }
52 if let Some(parent) = args["parent_task_id"].as_str() {
53 task.parent_task_id = Some(parent.to_string());
54 }
55 if let Some(assignee) = args["assignee"].as_str() {
56 task.assignee = assignee.to_string();
57 }
58
59 let created = registry.create(task).await?;
60 Ok(json!({
61 "task_id": created.id,
62 "title": created.title,
63 "status": format!("{:?}", created.status),
64 }))
65}
66
67pub async fn task_list(registry: &Arc<dyn TaskRegistry>, args: Value) -> Result<Value, PeError> {
72 let mut filter = TaskFilter::default();
73 if let Some(s) = args["status"].as_str() {
74 filter.status = Some(match s {
75 "pending" => TaskStatus::Pending,
76 "in_progress" => TaskStatus::InProgress,
77 "completed" => TaskStatus::Completed,
78 "failed" => TaskStatus::Failed,
79 "blocked" => TaskStatus::Blocked,
80 "cancelled" => TaskStatus::Cancelled,
81 other => {
82 return Err(PeError::InvalidUpdate {
83 details: format!(
84 "Unknown status '{other}'. Valid: pending, in_progress, completed, failed, blocked, cancelled"
85 ),
86 });
87 }
88 });
89 }
90 if let Some(a) = args["agent_id"].as_str() {
91 filter.agent_id = Some(a.to_string());
92 }
93 if let Some(a) = args["assignee"].as_str() {
94 filter.assignee = Some(a.to_string());
95 }
96 if let Some(p) = args["parent_id"].as_str() {
97 filter.parent_id = Some(p.to_string());
98 }
99 if let Some(t) = args["tag"].as_str() {
100 filter.tag = Some(t.to_string());
101 }
102 if let Some(l) = args["limit"].as_u64() {
103 filter.limit = l as usize;
104 }
105
106 let tasks = registry.list(&filter).await?;
107 let mut items: Vec<Value> = Vec::with_capacity(tasks.len());
108 for t in &tasks {
109 let dependencies = registry.get_dependencies(&t.id).await?;
110 let dependency_items = dependencies
111 .iter()
112 .map(|dep| {
113 json!({
114 "task_id": dep.task_id,
115 "depends_on_id": dep.depends_on_id,
116 "dependency_type": format!("{:?}", dep.dependency_type),
117 "created_at": dep.created_at.to_rfc3339(),
118 })
119 })
120 .collect::<Vec<_>>();
121 items.push(json!({
122 "id": t.id,
123 "title": t.title,
124 "status": format!("{:?}", t.status),
125 "priority": format!("{:?}", t.priority),
126 "assignee": t.assignee,
127 "parent_task_id": t.parent_task_id,
128 "dependencies": dependency_items,
129 }));
130 }
131
132 let count = items.len();
133 Ok(json!({ "tasks": items, "count": count }))
134}
135
136pub async fn task_complete(
141 registry: &Arc<dyn TaskRegistry>,
142 args: Value,
143) -> Result<Value, PeError> {
144 let task_id = args["task_id"]
145 .as_str()
146 .ok_or_else(|| PeError::InvalidUpdate {
147 details: "task_complete requires 'task_id' field".into(),
148 })?;
149
150 let task = registry
151 .get(&task_id.to_string())
152 .await?
153 .ok_or_else(|| PeError::NodeNotFound {
154 node: task_id.to_string(),
155 })?;
156
157 match &task.status {
159 TaskStatus::Pending => {
160 registry
162 .update_status(&task_id.to_string(), TaskStatus::InProgress, None, None)
163 .await?;
164 }
165 TaskStatus::InProgress => {} TaskStatus::Blocked | TaskStatus::Failed => {
167 return Err(PeError::InvalidUpdate {
168 details: format!("Task is {:?}. Resolve it before completing.", task.status),
169 });
170 }
171 TaskStatus::Completed | TaskStatus::Cancelled => {
172 return Err(PeError::InvalidUpdate {
173 details: format!("Task is already {:?}.", task.status),
174 });
175 }
176 }
177
178 let result = args.get("result").cloned();
179 let completed = registry
180 .update_status(&task_id.to_string(), TaskStatus::Completed, result, None)
181 .await?;
182
183 Ok(json!({
184 "task_id": completed.id,
185 "status": "Completed",
186 "completed_at": completed.completed_at.map(|t| t.to_rfc3339()),
187 }))
188}
189
190pub async fn task_decompose(
195 registry: &Arc<dyn TaskRegistry>,
196 agent_id: &str,
197 args: Value,
198) -> Result<Value, PeError> {
199 let parent_id = args["task_id"]
200 .as_str()
201 .ok_or_else(|| PeError::InvalidUpdate {
202 details: "task_decompose requires 'task_id' field".into(),
203 })?;
204
205 let subtasks_arr = args["subtasks"]
206 .as_array()
207 .ok_or_else(|| PeError::InvalidUpdate {
208 details: "task_decompose requires 'subtasks' array".into(),
209 })?;
210
211 let parent =
213 registry
214 .get(&parent_id.to_string())
215 .await?
216 .ok_or_else(|| PeError::NodeNotFound {
217 node: parent_id.to_string(),
218 })?;
219 if parent.status == TaskStatus::Pending {
220 registry
221 .update_status(&parent_id.to_string(), TaskStatus::InProgress, None, None)
222 .await?;
223 }
224
225 let mut created_ids: Vec<String> = Vec::new();
227 for sub_def in subtasks_arr {
228 let title = sub_def["title"].as_str().unwrap_or("Untitled subtask");
229 let mut sub = Task::agent_task(title, agent_id).with_parent(parent_id);
230 if let Some(desc) = sub_def["description"].as_str() {
231 sub.description = desc.to_string();
232 }
233 if let Some(p) = sub_def["priority"].as_str() {
234 sub.priority = match p {
235 "urgent" => TaskPriority::Urgent,
236 "high" => TaskPriority::High,
237 "low" => TaskPriority::Low,
238 _ => TaskPriority::Medium,
239 };
240 }
241 let created = registry.create(sub).await?;
242 created_ids.push(created.id);
243 }
244
245 let mut deps_created = 0;
247 for (i, sub_def) in subtasks_arr.iter().enumerate() {
248 if let Some(dep_indices) = sub_def["depends_on"].as_array() {
249 for dep_idx in dep_indices {
250 if let Some(idx) = dep_idx.as_u64() {
251 let idx = idx as usize;
252 if idx < created_ids.len() && idx != i {
253 let dep = TaskDependency::new(
254 &created_ids[i],
255 &created_ids[idx],
256 DependencyType::Blocks,
257 );
258 registry.add_dependency(dep).await?;
259 deps_created += 1;
260 }
261 }
262 }
263 }
264 }
265
266 Ok(json!({
267 "parent_id": parent_id,
268 "subtask_ids": created_ids,
269 "subtask_count": created_ids.len(),
270 "dependencies_created": deps_created,
271 }))
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277 use crate::InMemoryTaskRegistry;
278
279 fn reg() -> Arc<dyn TaskRegistry> {
280 Arc::new(InMemoryTaskRegistry::new())
281 }
282
283 #[tokio::test]
284 async fn test_task_create_tool() {
285 let r = reg();
286 let result = task_create(&r, "agent-1", json!({ "title": "Write tests" }))
287 .await
288 .unwrap();
289 assert_eq!(result["title"], "Write tests");
290 assert!(!result["task_id"].as_str().unwrap().is_empty());
291 }
292
293 #[tokio::test]
294 async fn test_task_list_tool() {
295 let r = reg();
296 task_create(&r, "a1", json!({ "title": "Task 1" }))
297 .await
298 .unwrap();
299 task_create(&r, "a1", json!({ "title": "Task 2" }))
300 .await
301 .unwrap();
302 let result = task_list(&r, json!({})).await.unwrap();
303 assert_eq!(result["count"], 2);
304 }
305
306 #[tokio::test]
307 async fn test_task_complete_tool() {
308 let r = reg();
309 let created = task_create(&r, "a1", json!({ "title": "Finish me" }))
310 .await
311 .unwrap();
312 let tid = created["task_id"].as_str().unwrap();
313 let result = task_complete(&r, json!({ "task_id": tid })).await.unwrap();
314 assert_eq!(result["status"], "Completed");
315 }
316
317 #[tokio::test]
318 async fn test_task_decompose_tool() {
319 let r = reg();
320 let parent = task_create(&r, "a1", json!({ "title": "Big project" }))
321 .await
322 .unwrap();
323 let pid = parent["task_id"].as_str().unwrap();
324
325 let result = task_decompose(
326 &r,
327 "a1",
328 json!({
329 "task_id": pid,
330 "subtasks": [
331 { "title": "Write code" },
332 { "title": "Write tests", "depends_on": [0] },
333 { "title": "Deploy", "depends_on": [0, 1] },
334 ]
335 }),
336 )
337 .await
338 .unwrap();
339
340 assert_eq!(result["subtask_count"], 3);
341 assert_eq!(result["dependencies_created"], 3);
342 }
343}