Skip to main content

codetether_agent/tool/
task.rs

1//! Task Tool - Spawn sub-tasks for parallel or sequential execution.
2
3use anyhow::{Context, Result};
4use async_trait::async_trait;
5use serde::Deserialize;
6use serde_json::{json, Value};
7use super::{Tool, ToolResult};
8use std::sync::atomic::{AtomicUsize, Ordering};
9use std::collections::HashMap;
10use parking_lot::RwLock;
11
12static TASK_COUNTER: AtomicUsize = AtomicUsize::new(1);
13
14lazy_static::lazy_static! {
15    static ref TASK_STORE: RwLock<HashMap<String, TaskInfo>> = RwLock::new(HashMap::new());
16}
17
18#[derive(Debug, Clone)]
19struct TaskInfo {
20    id: String,
21    description: String,
22    status: TaskStatus,
23    result: Option<String>,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq)]
27enum TaskStatus {
28    Pending,
29    #[allow(dead_code)]
30    Running,
31    Complete,
32    Failed,
33}
34
35pub struct TaskTool;
36
37impl Default for TaskTool {
38    fn default() -> Self { Self::new() }
39}
40
41impl TaskTool {
42    pub fn new() -> Self { Self }
43}
44
45#[derive(Deserialize)]
46struct Params {
47    action: String,  // create, status, complete, list, cancel
48    #[serde(default)]
49    id: Option<String>,
50    #[serde(default)]
51    description: Option<String>,
52    #[serde(default)]
53    result: Option<String>,
54}
55
56#[async_trait]
57impl Tool for TaskTool {
58    fn id(&self) -> &str { "task" }
59    fn name(&self) -> &str { "Task Manager" }
60    fn description(&self) -> &str { "Manage sub-tasks: create, query status, complete, or list tasks for tracking complex workflows." }
61    fn parameters(&self) -> Value {
62        json!({
63            "type": "object",
64            "properties": {
65                "action": {
66                    "type": "string",
67                    "enum": ["create", "status", "complete", "list", "cancel"],
68                    "description": "Action to perform"
69                },
70                "id": {"type": "string", "description": "Task ID (for status/complete/cancel)"},
71                "description": {"type": "string", "description": "Task description (for create)"},
72                "result": {"type": "string", "description": "Task result (for complete)"}
73            },
74            "required": ["action"]
75        })
76    }
77
78    async fn execute(&self, params: Value) -> Result<ToolResult> {
79        let p: Params = serde_json::from_value(params).context("Invalid params")?;
80        
81        match p.action.as_str() {
82            "create" => {
83                let description = p.description.ok_or_else(|| anyhow::anyhow!("description required"))?;
84                let id = format!("task_{}", TASK_COUNTER.fetch_add(1, Ordering::SeqCst));
85                
86                let task = TaskInfo {
87                    id: id.clone(),
88                    description: description.clone(),
89                    status: TaskStatus::Pending,
90                    result: None,
91                };
92                
93                TASK_STORE.write().insert(id.clone(), task);
94                Ok(ToolResult::success(format!("Created task: {} - {}", id, description))
95                    .with_metadata("task_id", json!(id)))
96            }
97            "status" => {
98                let id = p.id.ok_or_else(|| anyhow::anyhow!("id required"))?;
99                let store = TASK_STORE.read();
100                
101                match store.get(&id) {
102                    Some(task) => {
103                        let status_str = match task.status {
104                            TaskStatus::Pending => "pending",
105                            TaskStatus::Running => "running",
106                            TaskStatus::Complete => "complete",
107                            TaskStatus::Failed => "failed",
108                        };
109                        Ok(ToolResult::success(format!(
110                            "Task {} [{}]: {}\n{}", 
111                            task.id, status_str, task.description,
112                            task.result.as_deref().unwrap_or("")
113                        )))
114                    }
115                    None => Ok(ToolResult::error(format!("Task not found: {}", id)))
116                }
117            }
118            "complete" => {
119                let id = p.id.ok_or_else(|| anyhow::anyhow!("id required"))?;
120                let result = p.result.unwrap_or_else(|| "Completed".to_string());
121                
122                let mut store = TASK_STORE.write();
123                match store.get_mut(&id) {
124                    Some(task) => {
125                        task.status = TaskStatus::Complete;
126                        task.result = Some(result.clone());
127                        Ok(ToolResult::success(format!("Completed task: {}", id)))
128                    }
129                    None => Ok(ToolResult::error(format!("Task not found: {}", id)))
130                }
131            }
132            "cancel" => {
133                let id = p.id.ok_or_else(|| anyhow::anyhow!("id required"))?;
134                
135                let mut store = TASK_STORE.write();
136                match store.get_mut(&id) {
137                    Some(task) => {
138                        task.status = TaskStatus::Failed;
139                        task.result = Some("Cancelled".to_string());
140                        Ok(ToolResult::success(format!("Cancelled task: {}", id)))
141                    }
142                    None => Ok(ToolResult::error(format!("Task not found: {}", id)))
143                }
144            }
145            "list" => {
146                let store = TASK_STORE.read();
147                if store.is_empty() {
148                    return Ok(ToolResult::success("No active tasks".to_string()));
149                }
150                
151                let output = store.values().map(|t| {
152                    let icon = match t.status {
153                        TaskStatus::Pending => "○",
154                        TaskStatus::Running => "◐",
155                        TaskStatus::Complete => "●",
156                        TaskStatus::Failed => "✗",
157                    };
158                    format!("{} {} - {}", icon, t.id, t.description)
159                }).collect::<Vec<_>>().join("\n");
160                
161                Ok(ToolResult::success(output).with_metadata("count", json!(store.len())))
162            }
163            _ => Ok(ToolResult::error(format!("Unknown action: {}", p.action)))
164        }
165    }
166}