Skip to main content

pe_tasks/
tools.rs

1//! # Task tools — agent-facing API for task management.
2//!
3//! Four tools that agents use to create, query, complete, and decompose tasks.
4//! Each operates on a shared `Arc<dyn TaskRegistry>`.
5
6use 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
15/// Create a new task.
16///
17/// Args: `{ title, description?, priority?, tags?, parent_task_id?, assignee?, task_type? }`
18/// Returns: `{ task_id, title, status }`
19pub 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
67/// List tasks matching a filter.
68///
69/// Args: `{ status?, task_type?, agent_id?, assignee?, parent_id?, tag?, limit? }`
70/// Returns: `{ tasks: [{ id, title, status, priority, assignee }] }`
71pub 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
136/// Mark a task as completed.
137///
138/// Args: `{ task_id, result? }`
139/// Returns: `{ task_id, status, completed_at }`
140pub 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    // Handle current state before completing
158    match &task.status {
159        TaskStatus::Pending => {
160            // Auto-transition through InProgress
161            registry
162                .update_status(&task_id.to_string(), TaskStatus::InProgress, None, None)
163                .await?;
164        }
165        TaskStatus::InProgress => {} // ready to complete
166        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
190/// Decompose a task into subtasks with optional dependencies.
191///
192/// Args: `{ task_id, subtasks: [{ title, description?, priority?, depends_on?: [indices] }] }`
193/// Returns: `{ parent_id, subtask_ids: [...], dependencies_created }`
194pub 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    // Move parent to InProgress
212    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    // Create subtasks
226    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    // Wire dependencies (by index into subtasks array)
246    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}