Skip to main content

brainwires_agents/task_manager/
mod.rs

1//! Task Manager - Manages a tree-structured task list
2//!
3//! Provides functionality for the AI to create, update, and manage tasks
4//! in a hierarchical structure with dependencies.
5
6mod dependency_ops;
7mod query_ops;
8mod status_ops;
9mod time_tracking;
10
11#[cfg(test)]
12mod tests;
13
14use anyhow::{Context, Result};
15use std::collections::HashMap;
16use std::sync::Arc;
17use tokio::sync::RwLock;
18
19use brainwires_core::{Task, TaskPriority};
20
21// Re-export public types
22pub use time_tracking::{TaskStats, TaskTimeInfo, TimeStats, format_duration_secs};
23
24/// Manages a tree of tasks with dependencies
25#[derive(Debug, Clone)]
26pub struct TaskManager {
27    /// All tasks indexed by ID
28    pub(crate) tasks: Arc<RwLock<HashMap<String, Task>>>,
29}
30
31impl TaskManager {
32    /// Create a new task manager
33    pub fn new() -> Self {
34        Self {
35            tasks: Arc::new(RwLock::new(HashMap::new())),
36        }
37    }
38
39    /// Create a new task
40    #[tracing::instrument(name = "agent.task.create", skip(self, description))]
41    pub async fn create_task(
42        &self,
43        description: String,
44        parent_id: Option<String>,
45        priority: TaskPriority,
46    ) -> Result<String> {
47        let task_id = uuid::Uuid::new_v4().to_string();
48        let mut task = Task::new(task_id.clone(), description);
49        task.priority = priority;
50
51        let mut tasks = self.tasks.write().await;
52
53        // If parent_id is provided, validate and update parent
54        if let Some(ref pid) = parent_id {
55            let parent = tasks
56                .get_mut(pid)
57                .context(format!("Parent task '{}' not found", pid))?;
58            parent.add_child(task_id.clone());
59            task.parent_id = Some(pid.clone());
60        }
61
62        tasks.insert(task_id.clone(), task);
63        Ok(task_id)
64    }
65
66    /// Add a subtask to an existing task
67    pub async fn add_subtask(&self, parent_id: String, description: String) -> Result<String> {
68        self.create_task(description, Some(parent_id), TaskPriority::Normal)
69            .await
70    }
71
72    /// Get a task by ID
73    pub async fn get_task(&self, task_id: &str) -> Option<Task> {
74        let tasks = self.tasks.read().await;
75        tasks.get(task_id).cloned()
76    }
77
78    /// Clear all tasks
79    pub async fn clear(&self) {
80        let mut tasks = self.tasks.write().await;
81        tasks.clear();
82    }
83
84    /// Get task count
85    pub async fn count(&self) -> usize {
86        let tasks = self.tasks.read().await;
87        tasks.len()
88    }
89
90    /// Load tasks from storage
91    pub async fn load_tasks(&self, tasks_to_load: Vec<Task>) {
92        let mut tasks = self.tasks.write().await;
93        tasks.clear();
94        for task in tasks_to_load {
95            tasks.insert(task.id.clone(), task);
96        }
97    }
98
99    /// Export all tasks for persistence
100    pub async fn export_tasks(&self) -> Vec<Task> {
101        self.get_all_tasks().await
102    }
103
104    /// Assign a task to an agent (sets the `assigned_to` field).
105    pub async fn assign_task(&self, task_id: &str, agent_id: &str) -> Result<()> {
106        let mut tasks = self.tasks.write().await;
107        let task = tasks
108            .get_mut(task_id)
109            .context(format!("Task '{}' not found", task_id))?;
110
111        task.assigned_to = Some(agent_id.to_string());
112        task.updated_at = chrono::Utc::now().timestamp();
113        Ok(())
114    }
115}
116
117impl Default for TaskManager {
118    fn default() -> Self {
119        Self::new()
120    }
121}