1use std::collections::{BTreeMap, HashSet};
26
27use crate::{SessionId, SessionView};
28
29#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub enum TreeNodeId {
36 Project(String),
38 Worktree(String),
40 Team(String),
42 Agent(SessionId),
44}
45
46#[allow(clippy::large_enum_variant)]
50#[derive(Debug, Clone)]
51pub enum TreeNode {
52 Project {
54 name: String,
56 root: String,
58 children: Vec<TreeNode>,
60 },
61 Worktree {
63 path: String,
65 branch: Option<String>,
67 children: Vec<TreeNode>,
69 },
70 Team {
72 name: String,
74 children: Vec<TreeNode>,
76 },
77 Agent {
79 session: SessionView,
81 subagents: Vec<TreeNode>,
83 },
84}
85
86impl TreeNode {
87 pub fn agent_count(&self) -> usize {
89 match self {
90 TreeNode::Project { children, .. }
91 | TreeNode::Worktree { children, .. }
92 | TreeNode::Team { children, .. } => children.iter().map(|c| c.agent_count()).sum(),
93 TreeNode::Agent { subagents, .. } => {
94 1 + subagents.iter().map(|s| s.agent_count()).sum::<usize>()
95 }
96 }
97 }
98
99 pub fn needs_attention(&self) -> bool {
101 match self {
102 TreeNode::Project { children, .. }
103 | TreeNode::Worktree { children, .. }
104 | TreeNode::Team { children, .. } => children.iter().any(|c| c.needs_attention()),
105 TreeNode::Agent { session, subagents } => {
106 session.needs_attention || subagents.iter().any(|s| s.needs_attention())
107 }
108 }
109 }
110
111 pub fn node_id(&self) -> TreeNodeId {
113 match self {
114 TreeNode::Project { root, .. } => TreeNodeId::Project(root.clone()),
115 TreeNode::Worktree { path, .. } => TreeNodeId::Worktree(path.clone()),
116 TreeNode::Team { name, .. } => TreeNodeId::Team(name.clone()),
117 TreeNode::Agent { session, .. } => TreeNodeId::Agent(session.id.clone()),
118 }
119 }
120}
121
122#[allow(clippy::large_enum_variant)]
128#[derive(Debug, Clone)]
129pub enum TreeRowKind {
130 Project {
131 name: String,
132 root: String,
133 },
134 Worktree {
135 path: String,
136 branch: Option<String>,
137 },
138 Team {
139 name: String,
140 },
141 Agent {
142 session: SessionView,
143 },
144}
145
146#[derive(Debug, Clone)]
150pub struct TreeRow {
151 pub depth: u8,
153 pub node_id: TreeNodeId,
155 pub kind: TreeRowKind,
157 pub is_expanded: bool,
159 pub agent_count: usize,
161 pub needs_attention: bool,
163 pub has_children: bool,
165}
166
167const UNGROUPED_PROJECT_NAME: &str = "Other";
173const UNGROUPED_PROJECT_ROOT: &str = "__ungrouped__";
174
175pub fn build_tree(sessions: &[SessionView]) -> Vec<TreeNode> {
181 if sessions.is_empty() {
182 return Vec::new();
183 }
184
185 let child_ids: HashSet<&SessionId> = sessions
187 .iter()
188 .filter(|s| s.parent_session_id.is_some())
189 .map(|s| &s.id)
190 .collect();
191
192 let by_id: BTreeMap<&str, &SessionView> = sessions.iter().map(|s| (s.id.as_str(), s)).collect();
194
195 let mut by_project: BTreeMap<&str, Vec<&SessionView>> = BTreeMap::new();
198
199 for session in sessions {
200 if child_ids.contains(&session.id) {
202 continue;
203 }
204 let project_key = session
205 .project_root
206 .as_deref()
207 .unwrap_or(UNGROUPED_PROJECT_ROOT);
208 by_project.entry(project_key).or_default().push(session);
209 }
210
211 let mut project_nodes = Vec::new();
212
213 for (project_root, project_sessions) in &by_project {
214 let project_name = if *project_root == UNGROUPED_PROJECT_ROOT {
215 UNGROUPED_PROJECT_NAME.to_string()
216 } else {
217 extract_project_name(project_root)
218 };
219
220 let make_agent_node = |session: &SessionView| -> TreeNode {
222 let subagents: Vec<TreeNode> = session
223 .child_session_ids
224 .iter()
225 .filter_map(|child_id| by_id.get(child_id.as_str()))
226 .map(|child| TreeNode::Agent {
227 session: (*child).clone(),
228 subagents: Vec::new(), })
230 .collect();
231 TreeNode::Agent {
232 session: session.clone(),
233 subagents,
234 }
235 };
236
237 let mut by_worktree: BTreeMap<Option<&str>, Vec<&SessionView>> = BTreeMap::new();
239 for session in project_sessions {
240 let wt_key = session.worktree_path.as_deref();
241 by_worktree.entry(wt_key).or_default().push(session);
242 }
243
244 let skip_worktree = by_worktree.len() <= 1;
246
247 if skip_worktree {
248 let mut agents: Vec<TreeNode> = project_sessions
250 .iter()
251 .map(|s| make_agent_node(s))
252 .collect();
253 sort_agent_nodes(&mut agents);
254
255 project_nodes.push(TreeNode::Project {
256 name: project_name,
257 root: project_root.to_string(),
258 children: agents,
259 });
260 } else {
261 let mut worktree_nodes: Vec<TreeNode> = Vec::new();
263 for (wt_path, wt_sessions) in &by_worktree {
264 let branch = wt_sessions.first().and_then(|s| s.worktree_branch.clone());
265 let path = wt_path
266 .map(|p| p.to_string())
267 .unwrap_or_else(|| "unknown".to_string());
268
269 let mut agents: Vec<TreeNode> =
270 wt_sessions.iter().map(|s| make_agent_node(s)).collect();
271 sort_agent_nodes(&mut agents);
272
273 worktree_nodes.push(TreeNode::Worktree {
274 path,
275 branch,
276 children: agents,
277 });
278 }
279
280 project_nodes.push(TreeNode::Project {
281 name: project_name,
282 root: project_root.to_string(),
283 children: worktree_nodes,
284 });
285 }
286 }
287
288 project_nodes
289}
290
291fn sort_agent_nodes(nodes: &mut [TreeNode]) {
293 nodes.sort_by(|a, b| {
294 let a_time = match a {
295 TreeNode::Agent { session, .. } => session.started_at.as_str(),
296 _ => "",
297 };
298 let b_time = match b {
299 TreeNode::Agent { session, .. } => session.started_at.as_str(),
300 _ => "",
301 };
302 b_time.cmp(a_time) });
304}
305
306fn extract_project_name(path: &str) -> String {
308 path.rsplit('/')
309 .find(|s| !s.is_empty())
310 .unwrap_or(path)
311 .to_string()
312}
313
314pub fn flatten_tree(tree: &[TreeNode], expanded: &HashSet<TreeNodeId>) -> Vec<TreeRow> {
325 let mut rows = Vec::new();
326 for node in tree {
327 flatten_node(node, 0, expanded, &mut rows);
328 }
329 rows
330}
331
332fn flatten_node(
333 node: &TreeNode,
334 depth: u8,
335 expanded: &HashSet<TreeNodeId>,
336 rows: &mut Vec<TreeRow>,
337) {
338 let node_id = node.node_id();
339 let is_expanded = expanded.contains(&node_id);
340 let has_children = match node {
341 TreeNode::Project { children, .. }
342 | TreeNode::Worktree { children, .. }
343 | TreeNode::Team { children, .. } => !children.is_empty(),
344 TreeNode::Agent { subagents, .. } => !subagents.is_empty(),
345 };
346
347 let kind = match node {
348 TreeNode::Project { name, root, .. } => TreeRowKind::Project {
349 name: name.clone(),
350 root: root.clone(),
351 },
352 TreeNode::Worktree { path, branch, .. } => TreeRowKind::Worktree {
353 path: path.clone(),
354 branch: branch.clone(),
355 },
356 TreeNode::Team { name, .. } => TreeRowKind::Team { name: name.clone() },
357 TreeNode::Agent { session, .. } => TreeRowKind::Agent {
358 session: session.clone(),
359 },
360 };
361
362 rows.push(TreeRow {
363 depth,
364 node_id: node_id.clone(),
365 kind,
366 is_expanded,
367 agent_count: node.agent_count(),
368 needs_attention: node.needs_attention(),
369 has_children,
370 });
371
372 if is_expanded {
374 let children: &[TreeNode] = match node {
375 TreeNode::Project { children, .. }
376 | TreeNode::Worktree { children, .. }
377 | TreeNode::Team { children, .. } => children,
378 TreeNode::Agent { subagents, .. } => subagents,
379 };
380 for child in children {
381 flatten_node(child, depth.saturating_add(1), expanded, rows);
382 }
383 }
384}
385
386pub fn all_node_ids(tree: &[TreeNode]) -> HashSet<TreeNodeId> {
392 let mut ids = HashSet::new();
393 fn collect(node: &TreeNode, ids: &mut HashSet<TreeNodeId>) {
394 ids.insert(node.node_id());
395 match node {
396 TreeNode::Project { children, .. }
397 | TreeNode::Worktree { children, .. }
398 | TreeNode::Team { children, .. } => {
399 for child in children {
400 collect(child, ids);
401 }
402 }
403 TreeNode::Agent { subagents, .. } => {
404 for child in subagents {
405 collect(child, ids);
406 }
407 }
408 }
409 }
410 for node in tree {
411 collect(node, &mut ids);
412 }
413 ids
414}
415
416#[cfg(test)]
421mod tests {
422 use super::*;
423 use crate::SessionStatus;
424
425 fn make_session(id: &str) -> SessionView {
426 SessionView {
427 id: SessionId::new(id),
428 id_short: id.get(..8).unwrap_or(id).to_string(),
429 started_at: "2026-01-01T00:00:00Z".to_string(),
430 status: SessionStatus::Working,
431 ..Default::default()
432 }
433 }
434
435 fn make_session_in_project(
436 id: &str,
437 project_root: &str,
438 worktree_path: &str,
439 branch: &str,
440 started_at: &str,
441 ) -> SessionView {
442 SessionView {
443 id: SessionId::new(id),
444 id_short: id.get(..8).unwrap_or(id).to_string(),
445 project_root: Some(project_root.to_string()),
446 worktree_path: Some(worktree_path.to_string()),
447 worktree_branch: Some(branch.to_string()),
448 started_at: started_at.to_string(),
449 status: SessionStatus::Working,
450 ..Default::default()
451 }
452 }
453
454 #[test]
459 fn test_empty_sessions() {
460 let tree = build_tree(&[]);
461 assert!(tree.is_empty());
462 }
463
464 #[test]
465 fn test_ungrouped_sessions() {
466 let sessions = vec![make_session("a"), make_session("b")];
467 let tree = build_tree(&sessions);
468
469 assert_eq!(tree.len(), 1);
470 match &tree[0] {
471 TreeNode::Project { name, children, .. } => {
472 assert_eq!(name, "Other");
473 assert_eq!(children.len(), 2);
474 }
475 _ => panic!("expected Project node"),
476 }
477 }
478
479 #[test]
480 fn test_single_worktree_skips_nesting() {
481 let sessions = vec![
482 make_session_in_project(
483 "a",
484 "/home/user/myapp",
485 "/home/user/myapp",
486 "main",
487 "2026-01-01T00:00:00Z",
488 ),
489 make_session_in_project(
490 "b",
491 "/home/user/myapp",
492 "/home/user/myapp",
493 "main",
494 "2026-01-01T00:01:00Z",
495 ),
496 ];
497 let tree = build_tree(&sessions);
498
499 assert_eq!(tree.len(), 1);
500 match &tree[0] {
501 TreeNode::Project { name, children, .. } => {
502 assert_eq!(name, "myapp");
503 assert_eq!(children.len(), 2);
505 assert!(matches!(&children[0], TreeNode::Agent { .. }));
506 assert!(matches!(&children[1], TreeNode::Agent { .. }));
507 }
508 _ => panic!("expected Project node"),
509 }
510 }
511
512 #[test]
513 fn test_multiple_worktrees_adds_nesting() {
514 let sessions = vec![
515 make_session_in_project(
516 "a",
517 "/home/user/myapp",
518 "/home/user/myapp",
519 "main",
520 "2026-01-01T00:00:00Z",
521 ),
522 make_session_in_project(
523 "b",
524 "/home/user/myapp",
525 "/home/user/myapp-auth",
526 "feature/auth",
527 "2026-01-01T00:01:00Z",
528 ),
529 ];
530 let tree = build_tree(&sessions);
531
532 assert_eq!(tree.len(), 1);
533 match &tree[0] {
534 TreeNode::Project { name, children, .. } => {
535 assert_eq!(name, "myapp");
536 assert_eq!(children.len(), 2);
538 assert!(matches!(&children[0], TreeNode::Worktree { .. }));
539 assert!(matches!(&children[1], TreeNode::Worktree { .. }));
540 }
541 _ => panic!("expected Project node"),
542 }
543 }
544
545 #[test]
546 fn test_multiple_projects() {
547 let sessions = vec![
548 make_session_in_project(
549 "a",
550 "/home/user/app-a",
551 "/home/user/app-a",
552 "main",
553 "2026-01-01T00:00:00Z",
554 ),
555 make_session_in_project(
556 "b",
557 "/home/user/app-b",
558 "/home/user/app-b",
559 "main",
560 "2026-01-01T00:00:00Z",
561 ),
562 ];
563 let tree = build_tree(&sessions);
564
565 assert_eq!(tree.len(), 2);
567 match &tree[0] {
568 TreeNode::Project { name, .. } => assert_eq!(name, "app-a"),
569 _ => panic!("expected Project"),
570 }
571 match &tree[1] {
572 TreeNode::Project { name, .. } => assert_eq!(name, "app-b"),
573 _ => panic!("expected Project"),
574 }
575 }
576
577 #[test]
578 fn test_subagent_nesting() {
579 let mut parent = make_session_in_project(
580 "parent-1",
581 "/home/user/myapp",
582 "/home/user/myapp",
583 "main",
584 "2026-01-01T00:00:00Z",
585 );
586 parent.child_session_ids = vec![SessionId::new("child-1")];
587
588 let mut child = make_session_in_project(
589 "child-1",
590 "/home/user/myapp",
591 "/home/user/myapp",
592 "main",
593 "2026-01-01T00:00:01Z",
594 );
595 child.parent_session_id = Some(SessionId::new("parent-1"));
596
597 let sessions = vec![parent, child];
598 let tree = build_tree(&sessions);
599
600 assert_eq!(tree.len(), 1);
601 match &tree[0] {
602 TreeNode::Project { children, .. } => {
603 assert_eq!(children.len(), 1, "child should be nested, not top-level");
605 match &children[0] {
606 TreeNode::Agent {
607 session, subagents, ..
608 } => {
609 assert_eq!(session.id.as_str(), "parent-1");
610 assert_eq!(subagents.len(), 1);
611 match &subagents[0] {
612 TreeNode::Agent { session, .. } => {
613 assert_eq!(session.id.as_str(), "child-1");
614 }
615 _ => panic!("expected Agent subagent"),
616 }
617 }
618 _ => panic!("expected Agent"),
619 }
620 }
621 _ => panic!("expected Project"),
622 }
623 }
624
625 #[test]
626 fn test_agent_count() {
627 let sessions = vec![
628 make_session_in_project(
629 "a",
630 "/home/user/myapp",
631 "/home/user/myapp",
632 "main",
633 "2026-01-01T00:00:00Z",
634 ),
635 make_session_in_project(
636 "b",
637 "/home/user/myapp",
638 "/home/user/myapp",
639 "main",
640 "2026-01-01T00:01:00Z",
641 ),
642 make_session_in_project(
643 "c",
644 "/home/user/myapp",
645 "/home/user/myapp-wt",
646 "dev",
647 "2026-01-01T00:02:00Z",
648 ),
649 ];
650 let tree = build_tree(&sessions);
651
652 assert_eq!(tree[0].agent_count(), 3);
653 }
654
655 #[test]
656 fn test_attention_bubbles_up() {
657 let mut session = make_session_in_project(
658 "alert-1",
659 "/home/user/myapp",
660 "/home/user/myapp",
661 "main",
662 "2026-01-01T00:00:00Z",
663 );
664 session.needs_attention = true;
665
666 let normal = make_session_in_project(
667 "normal-1",
668 "/home/user/myapp",
669 "/home/user/myapp",
670 "main",
671 "2026-01-01T00:01:00Z",
672 );
673
674 let tree = build_tree(&[session, normal]);
675 assert!(
676 tree[0].needs_attention(),
677 "project should bubble up attention"
678 );
679 }
680
681 #[test]
682 fn test_agents_sorted_newest_first() {
683 let sessions = vec![
684 make_session_in_project(
685 "old",
686 "/home/user/app",
687 "/home/user/app",
688 "main",
689 "2026-01-01T00:00:00Z",
690 ),
691 make_session_in_project(
692 "new",
693 "/home/user/app",
694 "/home/user/app",
695 "main",
696 "2026-01-01T00:05:00Z",
697 ),
698 make_session_in_project(
699 "mid",
700 "/home/user/app",
701 "/home/user/app",
702 "main",
703 "2026-01-01T00:02:00Z",
704 ),
705 ];
706 let tree = build_tree(&sessions);
707
708 match &tree[0] {
709 TreeNode::Project { children, .. } => {
710 let ids: Vec<&str> = children
711 .iter()
712 .filter_map(|c| match c {
713 TreeNode::Agent { session, .. } => Some(session.id.as_str()),
714 _ => None,
715 })
716 .collect();
717 assert_eq!(ids, vec!["new", "mid", "old"]);
718 }
719 _ => panic!("expected Project"),
720 }
721 }
722
723 #[test]
728 fn test_flatten_empty() {
729 let rows = flatten_tree(&[], &HashSet::new());
730 assert!(rows.is_empty());
731 }
732
733 #[test]
734 fn test_flatten_collapsed_project() {
735 let sessions = vec![
736 make_session_in_project(
737 "a",
738 "/home/user/app",
739 "/home/user/app",
740 "main",
741 "2026-01-01T00:00:00Z",
742 ),
743 make_session_in_project(
744 "b",
745 "/home/user/app",
746 "/home/user/app",
747 "main",
748 "2026-01-01T00:01:00Z",
749 ),
750 ];
751 let tree = build_tree(&sessions);
752 let rows = flatten_tree(&tree, &HashSet::new()); assert_eq!(rows.len(), 1);
756 assert!(!rows[0].is_expanded);
757 assert_eq!(rows[0].agent_count, 2);
758 assert!(rows[0].has_children);
759 }
760
761 #[test]
762 fn test_flatten_expanded_project() {
763 let sessions = vec![
764 make_session_in_project(
765 "a",
766 "/home/user/app",
767 "/home/user/app",
768 "main",
769 "2026-01-01T00:00:00Z",
770 ),
771 make_session_in_project(
772 "b",
773 "/home/user/app",
774 "/home/user/app",
775 "main",
776 "2026-01-01T00:01:00Z",
777 ),
778 ];
779 let tree = build_tree(&sessions);
780
781 let mut expanded = HashSet::new();
782 expanded.insert(TreeNodeId::Project("/home/user/app".to_string()));
783
784 let rows = flatten_tree(&tree, &expanded);
785
786 assert_eq!(rows.len(), 3);
788 assert!(rows[0].is_expanded);
789 assert_eq!(rows[0].depth, 0);
790 assert_eq!(rows[1].depth, 1);
791 assert_eq!(rows[2].depth, 1);
792 }
793
794 #[test]
795 fn test_flatten_with_worktrees() {
796 let sessions = vec![
797 make_session_in_project(
798 "a",
799 "/home/user/app",
800 "/home/user/app",
801 "main",
802 "2026-01-01T00:00:00Z",
803 ),
804 make_session_in_project(
805 "b",
806 "/home/user/app",
807 "/home/user/app-wt",
808 "dev",
809 "2026-01-01T00:01:00Z",
810 ),
811 ];
812 let tree = build_tree(&sessions);
813
814 let expanded = all_node_ids(&tree);
816 let rows = flatten_tree(&tree, &expanded);
817
818 assert_eq!(rows.len(), 5);
820 assert_eq!(rows[0].depth, 0); assert_eq!(rows[1].depth, 1); assert_eq!(rows[2].depth, 2); assert_eq!(rows[3].depth, 1); assert_eq!(rows[4].depth, 2); }
826
827 #[test]
828 fn test_flatten_subagent_nesting() {
829 let mut parent = make_session_in_project(
830 "parent",
831 "/home/user/app",
832 "/home/user/app",
833 "main",
834 "2026-01-01T00:00:00Z",
835 );
836 parent.child_session_ids = vec![SessionId::new("child")];
837
838 let mut child = make_session_in_project(
839 "child",
840 "/home/user/app",
841 "/home/user/app",
842 "main",
843 "2026-01-01T00:00:01Z",
844 );
845 child.parent_session_id = Some(SessionId::new("parent"));
846
847 let tree = build_tree(&[parent, child]);
848 let expanded = all_node_ids(&tree);
849 let rows = flatten_tree(&tree, &expanded);
850
851 assert_eq!(rows.len(), 3);
853 assert_eq!(rows[0].depth, 0); assert_eq!(rows[1].depth, 1); assert_eq!(rows[2].depth, 2); }
857
858 #[test]
863 fn test_extract_project_name() {
864 assert_eq!(extract_project_name("/home/user/myapp"), "myapp");
865 assert_eq!(extract_project_name("/home/user/my-project"), "my-project");
866 assert_eq!(extract_project_name("single"), "single");
867 assert_eq!(extract_project_name("/trailing/slash/"), "slash");
868 }
869
870 #[test]
871 fn test_all_node_ids() {
872 let sessions = vec![make_session_in_project(
873 "a",
874 "/home/user/app",
875 "/home/user/app",
876 "main",
877 "2026-01-01T00:00:00Z",
878 )];
879 let tree = build_tree(&sessions);
880 let ids = all_node_ids(&tree);
881
882 assert!(ids.contains(&TreeNodeId::Project("/home/user/app".to_string())));
883 assert!(ids.contains(&TreeNodeId::Agent(SessionId::new("a"))));
884 assert_eq!(ids.len(), 2); }
886
887 #[test]
888 fn test_mixed_grouped_and_ungrouped() {
889 let sessions = vec![
890 make_session_in_project(
891 "a",
892 "/home/user/app",
893 "/home/user/app",
894 "main",
895 "2026-01-01T00:00:00Z",
896 ),
897 make_session("orphan"),
898 ];
899 let tree = build_tree(&sessions);
900
901 assert_eq!(tree.len(), 2);
903 let names: Vec<&str> = tree
904 .iter()
905 .filter_map(|n| match n {
906 TreeNode::Project { name, .. } => Some(name.as_str()),
907 _ => None,
908 })
909 .collect();
910 assert!(names.contains(&"Other"));
911 assert!(names.contains(&"app"));
912 }
913}