agent_air_runtime/controller/tools/
markdown_plan.rs1use std::collections::HashMap;
11use std::future::Future;
12use std::pin::Pin;
13use std::sync::Arc;
14
15use chrono::Utc;
16use tokio::fs;
17
18use super::plan_store::PlanStore;
19use super::types::{
20 DisplayConfig, DisplayResult, Executable, ResultContentType, ToolContext, ToolType,
21};
22
23pub const MARKDOWN_PLAN_TOOL_NAME: &str = "markdown_plan";
25
26pub const MARKDOWN_PLAN_TOOL_DESCRIPTION: &str = r#"Creates or updates a durable markdown plan file in the workspace.
28
29Usage:
30- Omit plan_id to create a new plan (an ID will be generated automatically)
31- Provide plan_id to overwrite an existing plan
32- Each step starts as pending ([ ])
33- The plan file is written to .agent-air/plans/ in the workspace root
34
35Returns:
36- The plan ID, file path, and rendered markdown content"#;
37
38pub const MARKDOWN_PLAN_TOOL_SCHEMA: &str = r#"{
40 "type": "object",
41 "properties": {
42 "plan_id": {
43 "type": "string",
44 "description": "Plan ID to update. Omit to create a new plan."
45 },
46 "title": {
47 "type": "string",
48 "description": "Plan title"
49 },
50 "steps": {
51 "type": "array",
52 "items": {
53 "type": "object",
54 "properties": {
55 "description": {
56 "type": "string"
57 },
58 "notes": {
59 "type": "string"
60 }
61 },
62 "required": ["description"]
63 }
64 },
65 "status": {
66 "type": "string",
67 "enum": ["draft", "active", "completed", "abandoned"]
68 }
69 },
70 "required": ["title", "steps"]
71}"#;
72
73pub struct MarkdownPlanTool {
78 plan_store: Arc<PlanStore>,
80}
81
82impl MarkdownPlanTool {
83 pub fn new(plan_store: Arc<PlanStore>) -> Self {
85 Self { plan_store }
86 }
87
88 fn generate_markdown(title: &str, plan_id: &str, status: &str, steps: &[PlanStep]) -> String {
90 let date = Utc::now().format("%Y-%m-%d");
91 let mut md = format!(
92 "# Plan: {}\n\n**ID**: {}\n**Status**: {}\n**Created**: {}\n\n## Steps\n\n",
93 title, plan_id, status, date
94 );
95
96 for (i, step) in steps.iter().enumerate() {
97 md.push_str(&format!("{}. [ ] {}\n", i + 1, step.description));
98 if let Some(ref notes) = step.notes {
99 md.push_str(&format!(" Notes: {}\n", notes));
100 }
101 }
102
103 md
104 }
105}
106
107struct PlanStep {
109 description: String,
110 notes: Option<String>,
111}
112
113impl Executable for MarkdownPlanTool {
114 fn name(&self) -> &str {
115 MARKDOWN_PLAN_TOOL_NAME
116 }
117
118 fn description(&self) -> &str {
119 MARKDOWN_PLAN_TOOL_DESCRIPTION
120 }
121
122 fn input_schema(&self) -> &str {
123 MARKDOWN_PLAN_TOOL_SCHEMA
124 }
125
126 fn tool_type(&self) -> ToolType {
127 ToolType::Custom
128 }
129
130 fn execute(
131 &self,
132 _context: ToolContext,
133 input: HashMap<String, serde_json::Value>,
134 ) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> {
135 let plan_store = self.plan_store.clone();
136
137 Box::pin(async move {
138 let title = input
140 .get("title")
141 .and_then(|v| v.as_str())
142 .ok_or_else(|| "Missing required 'title' parameter".to_string())?;
143
144 let steps_value = input
145 .get("steps")
146 .and_then(|v| v.as_array())
147 .ok_or_else(|| "Missing required 'steps' parameter".to_string())?;
148
149 if steps_value.is_empty() {
150 return Err("'steps' array must not be empty".to_string());
151 }
152
153 let status = input
154 .get("status")
155 .and_then(|v| v.as_str())
156 .unwrap_or("draft");
157
158 let mut steps = Vec::with_capacity(steps_value.len());
160 for (i, step_val) in steps_value.iter().enumerate() {
161 let description = step_val
162 .get("description")
163 .and_then(|v| v.as_str())
164 .ok_or_else(|| {
165 format!("Step {} is missing required 'description' field", i + 1)
166 })?;
167 let notes = step_val
168 .get("notes")
169 .and_then(|v| v.as_str())
170 .map(String::from);
171 steps.push(PlanStep {
172 description: description.to_string(),
173 notes,
174 });
175 }
176
177 let plan_id = match input.get("plan_id").and_then(|v| v.as_str()) {
179 Some(id) => id.to_string(),
180 None => plan_store.get_next_plan_id().await?,
181 };
182
183 let plans_dir = plan_store.plans_dir();
185 let file_name = format!("{}.md", plan_id);
186 let file_path = plans_dir.join(&file_name);
187 let file_path_str = file_path.to_string_lossy().to_string();
188
189 let lock = plan_store.acquire_lock(&file_path).await;
191 let _guard = lock.lock().await;
192
193 fs::create_dir_all(plans_dir)
195 .await
196 .map_err(|e| format!("Failed to create plans directory: {}", e))?;
197
198 let markdown = Self::generate_markdown(title, &plan_id, status, &steps);
200 fs::write(&file_path, &markdown)
201 .await
202 .map_err(|e| format!("Failed to write plan file: {}", e))?;
203
204 Ok(format!(
205 "Plan '{}' saved to {}\n\n{}",
206 plan_id, file_path_str, markdown
207 ))
208 })
209 }
210
211 fn handles_own_permissions(&self) -> bool {
212 true
213 }
214
215 fn display_config(&self) -> DisplayConfig {
216 DisplayConfig {
217 display_name: "Markdown Plan".to_string(),
218 display_title: Box::new(|input| {
219 input
220 .get("plan_id")
221 .and_then(|v| v.as_str())
222 .unwrap_or("new plan")
223 .to_string()
224 }),
225 display_content: Box::new(|_input, result| {
226 let lines: Vec<&str> = result.lines().take(15).collect();
227 let total_lines = result.lines().count();
228 let truncated = total_lines > 15;
229 let content = if truncated {
230 format!("{}...\n[truncated]", lines.join("\n"))
231 } else {
232 lines.join("\n")
233 };
234
235 DisplayResult {
236 content,
237 content_type: ResultContentType::Markdown,
238 is_truncated: truncated,
239 full_length: total_lines,
240 }
241 }),
242 }
243 }
244
245 fn compact_summary(&self, input: &HashMap<String, serde_json::Value>, _result: &str) -> String {
246 let plan_id = input
247 .get("plan_id")
248 .and_then(|v| v.as_str())
249 .unwrap_or("new");
250
251 let step_count = input
252 .get("steps")
253 .and_then(|v| v.as_array())
254 .map(|a| a.len())
255 .unwrap_or(0);
256
257 let status = input
258 .get("status")
259 .and_then(|v| v.as_str())
260 .unwrap_or("draft");
261
262 format!(
263 "[MarkdownPlan: {} ({} steps, {})]",
264 plan_id, step_count, status
265 )
266 }
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272 use tempfile::TempDir;
273
274 fn make_steps(descriptions: &[&str]) -> serde_json::Value {
275 let steps: Vec<serde_json::Value> = descriptions
276 .iter()
277 .map(|d| serde_json::json!({ "description": d }))
278 .collect();
279 serde_json::Value::Array(steps)
280 }
281
282 fn make_steps_with_notes(items: &[(&str, Option<&str>)]) -> serde_json::Value {
283 let steps: Vec<serde_json::Value> = items
284 .iter()
285 .map(|(desc, notes)| {
286 let mut step = serde_json::json!({ "description": desc });
287 if let Some(n) = notes {
288 step.as_object_mut()
289 .unwrap()
290 .insert("notes".to_string(), serde_json::json!(n));
291 }
292 step
293 })
294 .collect();
295 serde_json::Value::Array(steps)
296 }
297
298 fn make_context(tool_use_id: &str) -> ToolContext {
299 ToolContext {
300 session_id: 1,
301 tool_use_id: tool_use_id.to_string(),
302 turn_id: None,
303 permissions_pre_approved: false,
304 }
305 }
306
307 #[tokio::test]
308 async fn test_create_new_plan() {
309 let temp_dir = TempDir::new().unwrap();
310 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
311 let tool = MarkdownPlanTool::new(plan_store);
312
313 let mut input = HashMap::new();
314 input.insert(
315 "title".to_string(),
316 serde_json::Value::String("My Test Plan".to_string()),
317 );
318 input.insert("steps".to_string(), make_steps(&["Step one", "Step two"]));
319
320 let result = tool.execute(make_context("test-create"), input).await;
321 assert!(result.is_ok());
322
323 let output = result.unwrap();
324 assert!(output.contains("plan-001"));
325 assert!(output.contains("My Test Plan"));
326 assert!(output.contains("1. [ ] Step one"));
327 assert!(output.contains("2. [ ] Step two"));
328
329 let plan_file = temp_dir.path().join(".agent-air/plans/plan-001.md");
331 assert!(plan_file.exists());
332 }
333
334 #[tokio::test]
335 async fn test_upsert_existing_plan() {
336 let temp_dir = TempDir::new().unwrap();
337 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
338
339 let plans_dir = temp_dir.path().join(".agent-air/plans");
341 fs::create_dir_all(&plans_dir).await.unwrap();
342 fs::write(plans_dir.join("plan-001.md"), "# Old plan")
343 .await
344 .unwrap();
345
346 let tool = MarkdownPlanTool::new(plan_store);
347
348 let mut input = HashMap::new();
349 input.insert(
350 "plan_id".to_string(),
351 serde_json::Value::String("plan-001".to_string()),
352 );
353 input.insert(
354 "title".to_string(),
355 serde_json::Value::String("Updated Plan".to_string()),
356 );
357 input.insert("steps".to_string(), make_steps(&["New step"]));
358 input.insert(
359 "status".to_string(),
360 serde_json::Value::String("active".to_string()),
361 );
362
363 let result = tool.execute(make_context("test-upsert"), input).await;
364 assert!(result.is_ok());
365
366 let output = result.unwrap();
367 assert!(output.contains("plan-001"));
368 assert!(output.contains("Updated Plan"));
369 assert!(output.contains("active"));
370 assert!(output.contains("1. [ ] New step"));
371 }
372
373 #[tokio::test]
374 async fn test_sequential_id_generation() {
375 let temp_dir = TempDir::new().unwrap();
376 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
377
378 let plans_dir = temp_dir.path().join(".agent-air/plans");
380 fs::create_dir_all(&plans_dir).await.unwrap();
381 fs::write(plans_dir.join("plan-001.md"), "").await.unwrap();
382 fs::write(plans_dir.join("plan-002.md"), "").await.unwrap();
383
384 let next_id = plan_store.get_next_plan_id().await.unwrap();
385 assert_eq!(next_id, "plan-003");
386 }
387
388 #[tokio::test]
389 async fn test_missing_required_fields() {
390 let temp_dir = TempDir::new().unwrap();
391 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
392 let tool = MarkdownPlanTool::new(plan_store);
393
394 let mut input = HashMap::new();
396 input.insert("steps".to_string(), make_steps(&["A step"]));
397
398 let result = tool.execute(make_context("test-no-title"), input).await;
399 assert!(result.is_err());
400 assert!(result.unwrap_err().contains("Missing required 'title'"));
401
402 let mut input = HashMap::new();
404 input.insert(
405 "title".to_string(),
406 serde_json::Value::String("A Title".to_string()),
407 );
408
409 let result = tool.execute(make_context("test-no-steps"), input).await;
410 assert!(result.is_err());
411 assert!(result.unwrap_err().contains("Missing required 'steps'"));
412 }
413
414 #[tokio::test]
415 async fn test_default_status_is_draft() {
416 let temp_dir = TempDir::new().unwrap();
417 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
418 let tool = MarkdownPlanTool::new(plan_store);
419
420 let mut input = HashMap::new();
421 input.insert(
422 "title".to_string(),
423 serde_json::Value::String("Draft Plan".to_string()),
424 );
425 input.insert("steps".to_string(), make_steps(&["Step A"]));
426
427 let result = tool.execute(make_context("test-draft"), input).await;
428 assert!(result.is_ok());
429 assert!(result.unwrap().contains("**Status**: draft"));
430 }
431
432 #[test]
433 fn test_generate_markdown_format() {
434 let steps = vec![
435 PlanStep {
436 description: "First step".to_string(),
437 notes: None,
438 },
439 PlanStep {
440 description: "Second step".to_string(),
441 notes: None,
442 },
443 ];
444
445 let md = MarkdownPlanTool::generate_markdown("Test Plan", "plan-001", "draft", &steps);
446
447 assert!(md.starts_with("# Plan: Test Plan\n"));
448 assert!(md.contains("**ID**: plan-001"));
449 assert!(md.contains("**Status**: draft"));
450 assert!(md.contains("**Created**:"));
451 assert!(md.contains("## Steps"));
452 assert!(md.contains("1. [ ] First step"));
453 assert!(md.contains("2. [ ] Second step"));
454 }
455
456 #[test]
457 fn test_steps_with_notes() {
458 let steps = vec![
459 PlanStep {
460 description: "Step with notes".to_string(),
461 notes: Some("Important context here".to_string()),
462 },
463 PlanStep {
464 description: "Step without notes".to_string(),
465 notes: None,
466 },
467 ];
468
469 let md = MarkdownPlanTool::generate_markdown("Noted Plan", "plan-005", "active", &steps);
470
471 assert!(md.contains("1. [ ] Step with notes\n Notes: Important context here"));
472 assert!(md.contains("2. [ ] Step without notes\n"));
473 assert!(!md.contains("2. [ ] Step without notes\n Notes:"));
475 }
476
477 #[test]
478 fn test_compact_summary() {
479 let plan_store = Arc::new(PlanStore::new(std::path::PathBuf::from("/tmp")));
480 let tool = MarkdownPlanTool::new(plan_store);
481
482 let mut input = HashMap::new();
483 input.insert(
484 "plan_id".to_string(),
485 serde_json::Value::String("plan-001".to_string()),
486 );
487 input.insert("steps".to_string(), make_steps(&["A", "B", "C"]));
488 input.insert(
489 "status".to_string(),
490 serde_json::Value::String("active".to_string()),
491 );
492
493 let summary = tool.compact_summary(&input, "");
494 assert_eq!(summary, "[MarkdownPlan: plan-001 (3 steps, active)]");
495 }
496
497 #[test]
498 fn test_compact_summary_defaults() {
499 let plan_store = Arc::new(PlanStore::new(std::path::PathBuf::from("/tmp")));
500 let tool = MarkdownPlanTool::new(plan_store);
501
502 let input = HashMap::new();
503 let summary = tool.compact_summary(&input, "");
504 assert_eq!(summary, "[MarkdownPlan: new (0 steps, draft)]");
505 }
506
507 #[tokio::test]
508 async fn test_empty_steps_rejected() {
509 let temp_dir = TempDir::new().unwrap();
510 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
511 let tool = MarkdownPlanTool::new(plan_store);
512
513 let mut input = HashMap::new();
514 input.insert(
515 "title".to_string(),
516 serde_json::Value::String("Empty".to_string()),
517 );
518 input.insert("steps".to_string(), serde_json::Value::Array(vec![]));
519
520 let result = tool.execute(make_context("test-empty"), input).await;
521 assert!(result.is_err());
522 assert!(result.unwrap_err().contains("must not be empty"));
523 }
524
525 #[tokio::test]
526 async fn test_steps_with_notes_rendered() {
527 let temp_dir = TempDir::new().unwrap();
528 let plan_store = Arc::new(PlanStore::new(temp_dir.path().to_path_buf()));
529 let tool = MarkdownPlanTool::new(plan_store);
530
531 let mut input = HashMap::new();
532 input.insert(
533 "title".to_string(),
534 serde_json::Value::String("Noted".to_string()),
535 );
536 input.insert(
537 "steps".to_string(),
538 make_steps_with_notes(&[
539 ("Do something", Some("Watch out for edge cases")),
540 ("Do another thing", None),
541 ]),
542 );
543
544 let result = tool.execute(make_context("test-notes"), input).await;
545 assert!(result.is_ok());
546 let output = result.unwrap();
547 assert!(output.contains("Notes: Watch out for edge cases"));
548 }
549
550 #[test]
551 fn test_handles_own_permissions() {
552 let plan_store = Arc::new(PlanStore::new(std::path::PathBuf::from("/tmp")));
553 let tool = MarkdownPlanTool::new(plan_store);
554 assert!(tool.handles_own_permissions());
555 }
556}