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 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 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 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}