telos_agent/knowledge/tasks/
mod.rs

1mod persistence;
2pub mod task;
3use std::collections::HashMap;
4use std::path::PathBuf;
5use std::sync::Mutex;
6pub use task::{Task, TaskStatus};
7
8use crate::config::CancellationState;
9
10// Compatibility re-export: the task tool implementations now live in
11// `crate::tools` so they live next to the other built-in tools.
12pub use crate::tools::{
13    TaskCreateTool, TaskGetTool, TaskListTool, TaskOutputTool, TaskStopTool, TaskUpdateTool,
14};
15
16pub struct TaskManager {
17    tasks: Mutex<HashMap<String, Task>>,
18    cancellations: Mutex<HashMap<String, CancellationState>>,
19    dir: PathBuf,
20}
21
22impl TaskManager {
23    pub fn new(dir: PathBuf) -> Self {
24        std::fs::create_dir_all(&dir).ok();
25        let tasks =
26            persistence::load_all(&dir).into_iter().map(|task| (task.id.clone(), task)).collect();
27        Self { tasks: Mutex::new(tasks), cancellations: Mutex::new(HashMap::new()), dir }
28    }
29    pub fn create(&self, task: Task) {
30        let id = task.id.clone();
31        self.tasks.lock().unwrap().insert(id.clone(), task.clone());
32        let _ = persistence::save(&self.dir, &task);
33    }
34    pub fn get(&self, id: &str) -> Option<Task> {
35        self.tasks.lock().unwrap().get(id).cloned()
36    }
37    pub fn list(&self) -> Vec<Task> {
38        self.tasks.lock().unwrap().values().cloned().collect()
39    }
40    pub fn update(&self, id: &str, status: TaskStatus) {
41        let mut tasks = self.tasks.lock().unwrap();
42        if let Some(task) = tasks.get_mut(id) {
43            task.status = status;
44            let task = task.clone();
45            drop(tasks);
46            let _ = persistence::save(&self.dir, &task);
47        }
48    }
49
50    pub fn update_task(&self, id: &str, update: impl FnOnce(&mut Task)) {
51        let mut tasks = self.tasks.lock().unwrap();
52        if let Some(task) = tasks.get_mut(id) {
53            update(task);
54            let task = task.clone();
55            drop(tasks);
56            let _ = persistence::save(&self.dir, &task);
57        }
58    }
59
60    pub fn set_output(&self, id: &str, output: String) {
61        self.update_task(id, |task| {
62            task.output = Some(output);
63        });
64    }
65
66    pub fn complete(&self, id: &str, output: Option<String>) {
67        self.update_task(id, |task| {
68            task.status = TaskStatus::Completed;
69            task.error = None;
70            if let Some(output) = output {
71                task.output = Some(output);
72            }
73        });
74    }
75
76    pub fn fail(&self, id: &str, error: String) {
77        self.update_task(id, |task| {
78            task.status = TaskStatus::Failed;
79            task.error = Some(error);
80        });
81    }
82
83    pub fn cancel(&self, id: &str, reason: String) {
84        if let Some(cancellation) = self.cancellations.lock().unwrap().get(id).cloned() {
85            cancellation.cancel();
86        }
87        self.update_task(id, |task| {
88            task.status = TaskStatus::Cancelled;
89            task.error = Some(reason);
90        });
91    }
92
93    pub fn register_cancellation(&self, id: impl Into<String>, cancellation: CancellationState) {
94        self.cancellations.lock().unwrap().insert(id.into(), cancellation);
95    }
96
97    pub fn unregister_cancellation(&self, id: &str) {
98        self.cancellations.lock().unwrap().remove(id);
99    }
100
101    pub fn request_cancel(&self, id: &str) -> bool {
102        if let Some(cancellation) = self.cancellations.lock().unwrap().get(id).cloned() {
103            cancellation.cancel();
104            true
105        } else {
106            false
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn new_loads_existing_persisted_tasks() {
117        let dir = tempfile::tempdir().unwrap();
118        let task = Task {
119            id: "task_existing".into(),
120            subject: "Persisted task".into(),
121            description: "Loaded on startup".into(),
122            status: TaskStatus::InProgress,
123            blocked_by: vec![],
124            blocks: vec![],
125            output: None,
126            kind: None,
127            agent_id: None,
128            agent_type: None,
129            worktree_path: None,
130            error: None,
131        };
132        persistence::save(dir.path(), &task).unwrap();
133
134        let manager = TaskManager::new(dir.path().to_path_buf());
135
136        assert_eq!(manager.get("task_existing").unwrap().subject, "Persisted task");
137        assert_eq!(manager.list().len(), 1);
138    }
139
140    #[test]
141    fn new_ignores_malformed_and_non_json_task_files() {
142        let dir = tempfile::tempdir().unwrap();
143        let task = Task {
144            id: "task_valid".into(),
145            subject: "Valid task".into(),
146            description: "Loaded from valid JSON".into(),
147            status: TaskStatus::Pending,
148            blocked_by: vec![],
149            blocks: vec![],
150            output: None,
151            kind: None,
152            agent_id: None,
153            agent_type: None,
154            worktree_path: None,
155            error: None,
156        };
157        persistence::save(dir.path(), &task).unwrap();
158        std::fs::write(dir.path().join("bad.json"), "{not valid json").unwrap();
159        std::fs::write(dir.path().join("notes.txt"), "ignore me").unwrap();
160
161        let manager = TaskManager::new(dir.path().to_path_buf());
162
163        assert_eq!(manager.list().len(), 1);
164        assert_eq!(manager.get("task_valid").unwrap().subject, "Valid task");
165    }
166}
telos_agent/knowledge/tasks/mod.rs

telos_agent/knowledge/tasks/
mod.rs
