Skip to main content

agentzero_tools/
task_plan.rs

1use agentzero_core::{Tool, ToolContext, ToolResult};
2use anyhow::{anyhow, Context};
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use std::path::PathBuf;
6use std::sync::Mutex;
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
9#[serde(rename_all = "snake_case")]
10pub enum TaskStatus {
11    Pending,
12    InProgress,
13    Completed,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17struct TaskItem {
18    id: usize,
19    title: String,
20    status: TaskStatus,
21}
22
23#[derive(Debug, Deserialize)]
24#[serde(tag = "action")]
25#[serde(rename_all = "snake_case")]
26enum TaskAction {
27    Create { tasks: Vec<TaskCreate> },
28    Add { title: String },
29    Update { id: usize, status: TaskStatus },
30    List,
31    Delete,
32}
33
34#[derive(Debug, Deserialize)]
35struct TaskCreate {
36    title: String,
37    #[serde(default = "default_pending")]
38    status: TaskStatus,
39}
40
41fn default_pending() -> TaskStatus {
42    TaskStatus::Pending
43}
44
45pub struct TaskPlanTool {
46    tasks: Mutex<Vec<TaskItem>>,
47}
48
49impl Default for TaskPlanTool {
50    fn default() -> Self {
51        Self {
52            tasks: Mutex::new(Vec::new()),
53        }
54    }
55}
56
57impl TaskPlanTool {
58    fn persist(&self, workspace_root: &str) -> anyhow::Result<()> {
59        let tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
60        let dir = PathBuf::from(workspace_root).join(".agentzero");
61        std::fs::create_dir_all(&dir).ok();
62        let path = dir.join("task-plan.json");
63        let json = serde_json::to_string_pretty(&*tasks)?;
64        std::fs::write(path, json)?;
65        Ok(())
66    }
67
68    fn load(&self, workspace_root: &str) {
69        let path = PathBuf::from(workspace_root)
70            .join(".agentzero")
71            .join("task-plan.json");
72        if let Ok(content) = std::fs::read_to_string(path) {
73            if let Ok(items) = serde_json::from_str::<Vec<TaskItem>>(&content) {
74                if let Ok(mut tasks) = self.tasks.lock() {
75                    if tasks.is_empty() {
76                        *tasks = items;
77                    }
78                }
79            }
80        }
81    }
82
83    fn format_tasks(tasks: &[TaskItem]) -> String {
84        if tasks.is_empty() {
85            return "no tasks".to_string();
86        }
87        tasks
88            .iter()
89            .map(|t| {
90                let marker = match t.status {
91                    TaskStatus::Pending => "[ ]",
92                    TaskStatus::InProgress => "[~]",
93                    TaskStatus::Completed => "[x]",
94                };
95                format!("{} {}. {}", marker, t.id, t.title)
96            })
97            .collect::<Vec<_>>()
98            .join("\n")
99    }
100}
101
102#[async_trait]
103impl Tool for TaskPlanTool {
104    fn name(&self) -> &'static str {
105        "task_plan"
106    }
107
108    fn description(&self) -> &'static str {
109        "Manage a structured task plan: create, list, update status, or clear tasks for tracking multi-step work."
110    }
111
112    fn input_schema(&self) -> Option<serde_json::Value> {
113        Some(serde_json::json!({
114            "type": "object",
115            "properties": {
116                "action": { "type": "string", "enum": ["create", "add", "update", "list", "delete"], "description": "The task plan action to perform" },
117                "tasks": { "type": "array", "items": { "type": "object", "properties": { "title": { "type": "string" }, "status": { "type": "string" } }, "required": ["title"] }, "description": "Tasks to create (for create action)" },
118                "title": { "type": "string", "description": "Task title (for add action)" },
119                "id": { "type": "integer", "description": "Task ID (for update action)" },
120                "status": { "type": "string", "enum": ["pending", "in_progress", "completed"], "description": "New status (for update action)" }
121            },
122            "required": ["action"],
123            "additionalProperties": false
124        }))
125    }
126
127    async fn execute(&self, input: &str, ctx: &ToolContext) -> anyhow::Result<ToolResult> {
128        let action: TaskAction =
129            serde_json::from_str(input).context("task_plan expects JSON with \"action\" field")?;
130
131        self.load(&ctx.workspace_root);
132
133        match action {
134            TaskAction::Create { tasks: new_tasks } => {
135                let mut tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
136                tasks.clear();
137                for (i, tc) in new_tasks.into_iter().enumerate() {
138                    tasks.push(TaskItem {
139                        id: i + 1,
140                        title: tc.title,
141                        status: tc.status,
142                    });
143                }
144                let output = Self::format_tasks(&tasks);
145                drop(tasks);
146                self.persist(&ctx.workspace_root)?;
147                Ok(ToolResult {
148                    output: format!("plan created:\n{output}"),
149                })
150            }
151
152            TaskAction::Add { title } => {
153                if title.trim().is_empty() {
154                    return Err(anyhow!("title must not be empty"));
155                }
156                let mut tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
157                let id = tasks.iter().map(|t| t.id).max().unwrap_or(0) + 1;
158                tasks.push(TaskItem {
159                    id,
160                    title: title.clone(),
161                    status: TaskStatus::Pending,
162                });
163                drop(tasks);
164                self.persist(&ctx.workspace_root)?;
165                Ok(ToolResult {
166                    output: format!("added task {id}: {title}"),
167                })
168            }
169
170            TaskAction::Update { id, status } => {
171                let mut tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
172                let task = tasks
173                    .iter_mut()
174                    .find(|t| t.id == id)
175                    .ok_or_else(|| anyhow!("task {id} not found"))?;
176                task.status = status;
177                let output = format!("updated task {}: {}", task.id, task.title);
178                drop(tasks);
179                self.persist(&ctx.workspace_root)?;
180                Ok(ToolResult { output })
181            }
182
183            TaskAction::List => {
184                let tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
185                Ok(ToolResult {
186                    output: Self::format_tasks(&tasks),
187                })
188            }
189
190            TaskAction::Delete => {
191                let mut tasks = self.tasks.lock().map_err(|_| anyhow!("lock poisoned"))?;
192                tasks.clear();
193                drop(tasks);
194                self.persist(&ctx.workspace_root)?;
195                Ok(ToolResult {
196                    output: "plan deleted".to_string(),
197                })
198            }
199        }
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use std::fs;
207    use std::sync::atomic::{AtomicU64, Ordering};
208    use std::time::{SystemTime, UNIX_EPOCH};
209
210    static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
211
212    fn temp_dir() -> PathBuf {
213        let nanos = SystemTime::now()
214            .duration_since(UNIX_EPOCH)
215            .expect("clock")
216            .as_nanos();
217        let seq = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
218        let dir = std::env::temp_dir().join(format!(
219            "agentzero-task-plan-{}-{nanos}-{seq}",
220            std::process::id()
221        ));
222        fs::create_dir_all(&dir).expect("temp dir should be created");
223        dir
224    }
225
226    #[tokio::test]
227    async fn task_plan_create_list_update() {
228        let dir = temp_dir();
229        let tool = TaskPlanTool::default();
230        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
231
232        let result = tool
233            .execute(
234                r#"{"action": "create", "tasks": [{"title": "Step 1"}, {"title": "Step 2"}]}"#,
235                &ctx,
236            )
237            .await
238            .expect("create should succeed");
239        assert!(result.output.contains("Step 1"));
240        assert!(result.output.contains("Step 2"));
241
242        let result = tool
243            .execute(r#"{"action": "list"}"#, &ctx)
244            .await
245            .expect("list should succeed");
246        assert!(result.output.contains("[ ] 1. Step 1"));
247
248        let result = tool
249            .execute(
250                r#"{"action": "update", "id": 1, "status": "completed"}"#,
251                &ctx,
252            )
253            .await
254            .expect("update should succeed");
255        assert!(result.output.contains("updated task 1"));
256
257        let result = tool
258            .execute(r#"{"action": "list"}"#, &ctx)
259            .await
260            .expect("list should succeed");
261        assert!(result.output.contains("[x] 1. Step 1"));
262        fs::remove_dir_all(dir).ok();
263    }
264
265    #[tokio::test]
266    async fn task_plan_add_and_delete() {
267        let dir = temp_dir();
268        let tool = TaskPlanTool::default();
269        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
270
271        tool.execute(r#"{"action": "add", "title": "New task"}"#, &ctx)
272            .await
273            .expect("add should succeed");
274
275        let result = tool
276            .execute(r#"{"action": "list"}"#, &ctx)
277            .await
278            .expect("list should succeed");
279        assert!(result.output.contains("New task"));
280
281        tool.execute(r#"{"action": "delete"}"#, &ctx)
282            .await
283            .expect("delete should succeed");
284
285        let result = tool
286            .execute(r#"{"action": "list"}"#, &ctx)
287            .await
288            .expect("list should succeed");
289        assert!(result.output.contains("no tasks"));
290        fs::remove_dir_all(dir).ok();
291    }
292
293    #[tokio::test]
294    async fn task_plan_update_nonexistent_fails() {
295        let dir = temp_dir();
296        let tool = TaskPlanTool::default();
297        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
298
299        let err = tool
300            .execute(
301                r#"{"action": "update", "id": 99, "status": "completed"}"#,
302                &ctx,
303            )
304            .await
305            .expect_err("update nonexistent should fail");
306        assert!(err.to_string().contains("not found"));
307        fs::remove_dir_all(dir).ok();
308    }
309
310    #[tokio::test]
311    async fn task_plan_empty_create_succeeds() {
312        let dir = temp_dir();
313        let tool = TaskPlanTool::default();
314        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
315
316        let result = tool
317            .execute(r#"{"action": "create", "tasks": []}"#, &ctx)
318            .await
319            .expect("create with empty tasks should succeed");
320        // Should acknowledge creation even if no tasks.
321        assert!(!result.output.is_empty());
322        fs::remove_dir_all(dir).ok();
323    }
324
325    #[tokio::test]
326    async fn task_plan_status_transitions() {
327        let dir = temp_dir();
328        let tool = TaskPlanTool::default();
329        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
330
331        tool.execute(
332            r#"{"action": "create", "tasks": [{"title": "My task"}]}"#,
333            &ctx,
334        )
335        .await
336        .expect("create");
337
338        // pending → in_progress
339        tool.execute(
340            r#"{"action": "update", "id": 1, "status": "in_progress"}"#,
341            &ctx,
342        )
343        .await
344        .expect("pending to in_progress");
345        let list = tool
346            .execute(r#"{"action": "list"}"#, &ctx)
347            .await
348            .expect("list");
349        assert!(list.output.contains("[~] 1. My task"));
350
351        // in_progress → completed
352        tool.execute(
353            r#"{"action": "update", "id": 1, "status": "completed"}"#,
354            &ctx,
355        )
356        .await
357        .expect("in_progress to completed");
358        let list = tool
359            .execute(r#"{"action": "list"}"#, &ctx)
360            .await
361            .expect("list");
362        assert!(list.output.contains("[x] 1. My task"));
363        fs::remove_dir_all(dir).ok();
364    }
365
366    #[tokio::test]
367    async fn task_plan_invalid_input_returns_error() {
368        let dir = temp_dir();
369        let tool = TaskPlanTool::default();
370        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
371
372        let err = tool
373            .execute("not json at all", &ctx)
374            .await
375            .expect_err("invalid JSON should fail");
376        assert!(!err.to_string().is_empty());
377        fs::remove_dir_all(dir).ok();
378    }
379}