Skip to main content

perspt_tui/
task_tree.rs

1//! Task Tree Component
2//!
3//! Displays the SRBN Task DAG as an interactive tree view with expand/collapse support.
4
5use crate::theme::Theme;
6use ratatui::{
7    layout::Rect,
8    style::{Color, Modifier, Style},
9    text::{Line, Span},
10    widgets::{Block, Borders, List, ListItem, ListState},
11    Frame,
12};
13use std::collections::{HashMap, HashSet};
14
15/// Node status for display
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum TaskStatus {
18    Queued,
19    Planning,
20    Pending,
21    Coding,
22    Running,
23    Verifying,
24    Retrying,
25    SheafCheck,
26    Committing,
27    Completed,
28    Failed,
29    Escalated,
30    Aborted,
31}
32
33impl TaskStatus {
34    pub fn icon(&self) -> &'static str {
35        match self {
36            TaskStatus::Queued => "◇",
37            TaskStatus::Planning => "◈",
38            TaskStatus::Pending => "○",
39            TaskStatus::Coding => "◉",
40            TaskStatus::Running => "◐",
41            TaskStatus::Verifying => "◑",
42            TaskStatus::Retrying => "↻",
43            TaskStatus::SheafCheck => "⊘",
44            TaskStatus::Committing => "⊙",
45            TaskStatus::Completed => "●",
46            TaskStatus::Failed => "✗",
47            TaskStatus::Escalated => "⚠",
48            TaskStatus::Aborted => "⊘",
49        }
50    }
51
52    pub fn color(&self) -> Color {
53        match self {
54            TaskStatus::Queued => Color::Rgb(158, 158, 158), // Lighter gray
55            TaskStatus::Planning => Color::Rgb(179, 157, 219), // Light purple
56            TaskStatus::Pending => Color::Rgb(120, 144, 156), // Gray
57            TaskStatus::Coding => Color::Rgb(255, 213, 79),  // Yellow
58            TaskStatus::Running => Color::Rgb(255, 183, 77), // Amber
59            TaskStatus::Verifying => Color::Rgb(129, 212, 250), // Light blue
60            TaskStatus::Retrying => Color::Rgb(255, 152, 0), // Orange
61            TaskStatus::SheafCheck => Color::Rgb(77, 208, 225), // Cyan
62            TaskStatus::Committing => Color::Rgb(165, 214, 167), // Light green
63            TaskStatus::Completed => Color::Rgb(102, 187, 106), // Green
64            TaskStatus::Failed => Color::Rgb(239, 83, 80),   // Red
65            TaskStatus::Escalated => Color::Rgb(186, 104, 200), // Purple
66            TaskStatus::Aborted => Color::Rgb(255, 152, 0),  // Orange
67        }
68    }
69}
70
71impl From<perspt_core::NodeStatus> for TaskStatus {
72    fn from(status: perspt_core::NodeStatus) -> Self {
73        match status {
74            perspt_core::NodeStatus::Queued => TaskStatus::Queued,
75            perspt_core::NodeStatus::Planning => TaskStatus::Planning,
76            perspt_core::NodeStatus::Pending => TaskStatus::Pending,
77            perspt_core::NodeStatus::Coding => TaskStatus::Coding,
78            perspt_core::NodeStatus::Running => TaskStatus::Running,
79            perspt_core::NodeStatus::Verifying => TaskStatus::Verifying,
80            perspt_core::NodeStatus::Retrying => TaskStatus::Retrying,
81            perspt_core::NodeStatus::SheafCheck => TaskStatus::SheafCheck,
82            perspt_core::NodeStatus::Committing => TaskStatus::Committing,
83            perspt_core::NodeStatus::Completed => TaskStatus::Completed,
84            perspt_core::NodeStatus::Failed => TaskStatus::Failed,
85            perspt_core::NodeStatus::Escalated => TaskStatus::Escalated,
86            perspt_core::NodeStatus::Aborted => TaskStatus::Aborted,
87        }
88    }
89}
90
91/// A task node for the tree view
92#[derive(Debug, Clone)]
93pub struct TaskNode {
94    /// Unique identifier
95    pub id: String,
96    /// Task goal/description
97    pub goal: String,
98    /// Current status
99    pub status: TaskStatus,
100    /// Depth in tree (for indentation)
101    pub depth: usize,
102    /// Parent node ID (None for root)
103    pub parent_id: Option<String>,
104    /// Whether this node has children
105    pub has_children: bool,
106    /// Lyapunov energy (if available)
107    pub energy: Option<f32>,
108    /// Retry count (incremented on Retrying status)
109    pub retry_count: usize,
110}
111
112/// Task tree viewer state with expand/collapse support
113#[derive(Default)]
114pub struct TaskTree {
115    /// All task nodes indexed by ID
116    nodes: HashMap<String, TaskNode>,
117    /// Root node IDs (top-level tasks)
118    roots: Vec<String>,
119    /// Currently collapsed node IDs
120    collapsed: HashSet<String>,
121    /// Flattened visible list for display
122    visible_tasks: Vec<String>,
123    /// Selection state
124    pub state: ListState,
125    /// Theme for styling
126    theme: Theme,
127}
128
129impl TaskTree {
130    /// Create a new task tree
131    pub fn new() -> Self {
132        Self::default()
133    }
134
135    /// Add a task to the tree (legacy API for compatibility)
136    pub fn add_task(&mut self, id: String, goal: String, depth: usize) {
137        let node = TaskNode {
138            id: id.clone(),
139            goal,
140            status: TaskStatus::Pending,
141            depth,
142            parent_id: None,
143            has_children: false,
144            energy: None,
145            retry_count: 0,
146        };
147
148        if depth == 0 {
149            self.roots.push(id.clone());
150        }
151
152        self.nodes.insert(id, node);
153        self.rebuild_visible();
154    }
155
156    /// Populate tree from TaskPlan using dependency information for tree structure
157    pub fn populate_from_plan(&mut self, plan: perspt_core::types::TaskPlan) {
158        self.clear();
159
160        // Build a map of task ID to dependencies for depth calculation
161        let mut depth_map: HashMap<String, usize> = HashMap::new();
162
163        // First pass: insert all tasks with initial depth based on dependencies
164        for task in &plan.tasks {
165            // Calculate depth: max depth of dependencies + 1, or 0 if no deps
166            let depth = if task.dependencies.is_empty() {
167                0
168            } else {
169                task.dependencies
170                    .iter()
171                    .filter_map(|dep_id| depth_map.get(dep_id))
172                    .max()
173                    .map(|d| d + 1)
174                    .unwrap_or(0)
175            };
176            depth_map.insert(task.id.clone(), depth);
177
178            // Use first dependency as parent (for tree visualization)
179            // This creates a logical parent-child relationship
180            let parent_id = task.dependencies.first().cloned();
181
182            self.add_task_with_parent(task.id.clone(), task.goal.clone(), parent_id, depth);
183        }
184
185        // Selection reset
186        if !self.visible_tasks.is_empty() {
187            self.state.select(Some(0));
188        }
189    }
190
191    /// Add a task with parent relationship
192    pub fn add_task_with_parent(
193        &mut self,
194        id: String,
195        goal: String,
196        parent_id: Option<String>,
197        depth: usize,
198    ) {
199        // Mark parent as having children
200        if let Some(ref pid) = parent_id {
201            if let Some(parent) = self.nodes.get_mut(pid) {
202                parent.has_children = true;
203            }
204        }
205
206        let is_root = parent_id.is_none();
207        let node = TaskNode {
208            id: id.clone(),
209            goal,
210            status: TaskStatus::Pending,
211            depth,
212            parent_id,
213            has_children: false,
214            energy: None,
215            retry_count: 0,
216        };
217
218        if is_root {
219            self.roots.push(id.clone());
220        }
221
222        self.nodes.insert(id, node);
223        self.rebuild_visible();
224    }
225
226    /// Clear all tasks
227    pub fn clear(&mut self) {
228        self.nodes.clear();
229        self.roots.clear();
230        self.collapsed.clear();
231        self.visible_tasks.clear();
232        self.state.select(None);
233    }
234
235    /// Update task status
236    pub fn update_status(&mut self, id: &str, status: TaskStatus) {
237        if let Some(task) = self.nodes.get_mut(id) {
238            if status == TaskStatus::Retrying {
239                task.retry_count += 1;
240            }
241            task.status = status;
242        }
243    }
244
245    /// PSP-5 Phase 8: Add a node or update its status if already present.
246    ///
247    /// Used during resume to pre-populate the tree from persisted state
248    /// without requiring a full TaskPlan.
249    pub fn add_or_update_node(&mut self, id: &str, goal: &str, status: TaskStatus) {
250        if let Some(task) = self.nodes.get_mut(id) {
251            task.status = status;
252        } else {
253            let node = TaskNode {
254                id: id.to_string(),
255                goal: goal.to_string(),
256                status,
257                depth: 0,
258                parent_id: None,
259                has_children: false,
260                energy: None,
261                retry_count: 0,
262            };
263            self.roots.push(id.to_string());
264            self.nodes.insert(id.to_string(), node);
265            self.rebuild_visible();
266        }
267    }
268
269    /// Update task energy
270    pub fn update_energy(&mut self, id: &str, energy: f32) {
271        if let Some(task) = self.nodes.get_mut(id) {
272            task.energy = Some(energy);
273        }
274    }
275
276    /// Rebuild the visible task list based on collapse state
277    fn rebuild_visible(&mut self) {
278        self.visible_tasks.clear();
279
280        // Sort tasks by depth for proper tree structure
281        let mut sorted: Vec<_> = self.nodes.values().collect();
282        sorted.sort_by(|a, b| a.depth.cmp(&b.depth).then_with(|| a.id.cmp(&b.id)));
283
284        // Build parent-children map
285        let mut children_map: HashMap<Option<String>, Vec<String>> = HashMap::new();
286        for node in sorted {
287            children_map
288                .entry(node.parent_id.clone())
289                .or_default()
290                .push(node.id.clone());
291        }
292
293        // DFS traversal respecting collapse state
294        fn dfs(
295            node_id: &str,
296            nodes: &HashMap<String, TaskNode>,
297            children_map: &HashMap<Option<String>, Vec<String>>,
298            collapsed: &HashSet<String>,
299            result: &mut Vec<String>,
300        ) {
301            result.push(node_id.to_string());
302
303            if collapsed.contains(node_id) {
304                return; // Skip children if collapsed
305            }
306
307            if let Some(children) = children_map.get(&Some(node_id.to_string())) {
308                for child_id in children {
309                    if nodes.contains_key(child_id) {
310                        dfs(child_id, nodes, children_map, collapsed, result);
311                    }
312                }
313            }
314        }
315
316        // Start from roots
317        if let Some(root_children) = children_map.get(&None) {
318            for root_id in root_children {
319                dfs(
320                    root_id,
321                    &self.nodes,
322                    &children_map,
323                    &self.collapsed,
324                    &mut self.visible_tasks,
325                );
326            }
327        }
328    }
329
330    /// Select next task
331    pub fn next(&mut self) {
332        let len = self.visible_tasks.len();
333        if len == 0 {
334            return;
335        }
336        let i = match self.state.selected() {
337            Some(i) => {
338                if i >= len - 1 {
339                    0
340                } else {
341                    i + 1
342                }
343            }
344            None => 0,
345        };
346        self.state.select(Some(i));
347    }
348
349    /// Select previous task
350    pub fn previous(&mut self) {
351        let len = self.visible_tasks.len();
352        if len == 0 {
353            return;
354        }
355        let i = match self.state.selected() {
356            Some(i) => {
357                if i == 0 {
358                    len - 1
359                } else {
360                    i - 1
361                }
362            }
363            None => 0,
364        };
365        self.state.select(Some(i));
366    }
367
368    /// Get the currently selected task
369    pub fn selected_task(&self) -> Option<&TaskNode> {
370        self.state
371            .selected()
372            .and_then(|i| self.visible_tasks.get(i))
373            .and_then(|id| self.nodes.get(id))
374    }
375
376    /// Render the task tree
377    pub fn render(&mut self, frame: &mut Frame, area: Rect) {
378        let items: Vec<ListItem> = self
379            .visible_tasks
380            .iter()
381            .filter_map(|id| self.nodes.get(id))
382            .map(|task| {
383                // Build tree connector characters
384                let indent = "  ".repeat(task.depth);
385                let collapse_indicator = if task.has_children {
386                    if self.collapsed.contains(&task.id) {
387                        "▶ " // Collapsed
388                    } else {
389                        "▼ " // Expanded
390                    }
391                } else {
392                    "  " // Leaf node
393                };
394
395                let icon = task.status.icon();
396                let color = task.status.color();
397                let goal = truncate(&task.goal, 35);
398
399                // Build styled spans
400                let mut spans = vec![
401                    Span::styled(indent, Style::default().fg(Color::DarkGray)),
402                    Span::styled(collapse_indicator, Style::default().fg(Color::Cyan)),
403                    Span::styled(format!("{} ", icon), Style::default().fg(color)),
404                ];
405
406                // Add energy if available
407                if let Some(energy) = task.energy {
408                    let energy_style = self.theme.energy_style(energy);
409                    spans.push(Span::styled(format!("[{:.2}] ", energy), energy_style));
410                }
411
412                // Add retry count if > 0
413                if task.retry_count > 0 {
414                    spans.push(Span::styled(
415                        format!("↻{} ", task.retry_count),
416                        Style::default().fg(Color::Rgb(255, 152, 0)),
417                    ));
418                }
419
420                spans.push(Span::styled(
421                    format!("{}: ", task.id),
422                    Style::default().fg(color).add_modifier(Modifier::BOLD),
423                ));
424                spans.push(Span::styled(goal, Style::default().fg(Color::White)));
425
426                ListItem::new(Line::from(spans))
427            })
428            .collect();
429
430        let title = format!(
431            "🌳 Task DAG ({} nodes{})",
432            self.visible_tasks.len(),
433            if !self.collapsed.is_empty() {
434                format!(", {} collapsed", self.collapsed.len())
435            } else {
436                String::new()
437            }
438        );
439
440        let list = List::new(items)
441            .block(
442                Block::default()
443                    .title(title)
444                    .borders(Borders::ALL)
445                    .border_style(Style::default().fg(Color::Rgb(96, 125, 139))),
446            )
447            .highlight_style(
448                Style::default()
449                    .bg(Color::Rgb(55, 71, 79))
450                    .add_modifier(Modifier::BOLD),
451            )
452            .highlight_symbol("→ ");
453
454        frame.render_stateful_widget(list, area, &mut self.state);
455    }
456}
457
458/// Truncate a string to max length with ellipsis
459fn truncate(s: &str, max: usize) -> String {
460    if s.chars().count() > max {
461        format!(
462            "{}...",
463            s.chars().take(max.saturating_sub(3)).collect::<String>()
464        )
465    } else {
466        s.to_string()
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    #[test]
475    fn test_add_tasks() {
476        let mut tree = TaskTree::new();
477        tree.add_task("root".to_string(), "Root task".to_string(), 0);
478        tree.add_task("child1".to_string(), "Child 1".to_string(), 1);
479
480        assert_eq!(tree.nodes.len(), 2);
481        assert_eq!(tree.visible_tasks.len(), 2);
482    }
483
484    #[test]
485    fn test_update_status() {
486        let mut tree = TaskTree::new();
487        tree.add_task("task1".to_string(), "Test".to_string(), 0);
488        tree.update_status("task1", TaskStatus::Running);
489
490        assert_eq!(tree.nodes.get("task1").unwrap().status, TaskStatus::Running);
491    }
492
493    #[test]
494    fn test_navigation() {
495        let mut tree = TaskTree::new();
496        tree.add_task("t1".to_string(), "Task 1".to_string(), 0);
497        tree.add_task("t2".to_string(), "Task 2".to_string(), 0);
498        tree.add_task("t3".to_string(), "Task 3".to_string(), 0);
499
500        assert!(tree.state.selected().is_none());
501
502        tree.next();
503        assert_eq!(tree.state.selected(), Some(0));
504
505        tree.next();
506        assert_eq!(tree.state.selected(), Some(1));
507
508        tree.previous();
509        assert_eq!(tree.state.selected(), Some(0));
510    }
511
512    #[test]
513    fn test_lifecycle_mapping_all_variants() {
514        // Verify all NodeStatus variants map to the expected TaskStatus
515        use perspt_core::NodeStatus;
516        let mappings = vec![
517            (NodeStatus::Queued, TaskStatus::Queued),
518            (NodeStatus::Planning, TaskStatus::Planning),
519            (NodeStatus::Pending, TaskStatus::Pending),
520            (NodeStatus::Coding, TaskStatus::Coding),
521            (NodeStatus::Running, TaskStatus::Running),
522            (NodeStatus::Verifying, TaskStatus::Verifying),
523            (NodeStatus::Retrying, TaskStatus::Retrying),
524            (NodeStatus::SheafCheck, TaskStatus::SheafCheck),
525            (NodeStatus::Committing, TaskStatus::Committing),
526            (NodeStatus::Completed, TaskStatus::Completed),
527            (NodeStatus::Failed, TaskStatus::Failed),
528            (NodeStatus::Escalated, TaskStatus::Escalated),
529            (NodeStatus::Aborted, TaskStatus::Aborted),
530        ];
531        for (node_status, expected) in mappings {
532            let result: TaskStatus = node_status.into();
533            assert_eq!(
534                result, expected,
535                "NodeStatus::{:?} should map to TaskStatus::{:?}",
536                node_status, expected
537            );
538        }
539    }
540
541    #[test]
542    fn test_retry_count_increments_on_retrying() {
543        let mut tree = TaskTree::new();
544        tree.add_task("t1".to_string(), "Task".to_string(), 0);
545        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 0);
546
547        tree.update_status("t1", TaskStatus::Retrying);
548        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 1);
549
550        tree.update_status("t1", TaskStatus::Verifying);
551        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 1);
552
553        tree.update_status("t1", TaskStatus::Retrying);
554        assert_eq!(tree.nodes.get("t1").unwrap().retry_count, 2);
555    }
556
557    #[test]
558    fn test_status_icons_and_colors_unique() {
559        let statuses = vec![
560            TaskStatus::Queued,
561            TaskStatus::Planning,
562            TaskStatus::Pending,
563            TaskStatus::Coding,
564            TaskStatus::Running,
565            TaskStatus::Verifying,
566            TaskStatus::Retrying,
567            TaskStatus::SheafCheck,
568            TaskStatus::Committing,
569            TaskStatus::Completed,
570            TaskStatus::Failed,
571            TaskStatus::Escalated,
572        ];
573        // Every status should have a non-empty icon
574        for s in &statuses {
575            assert!(!s.icon().is_empty(), "{:?} should have an icon", s);
576        }
577    }
578}