Skip to main content

atm_core/
tree.rs

1//! Tree model for grouping sessions by project, worktree, and team.
2//!
3//! Transforms a flat list of [`SessionView`] into a hierarchical tree
4//! structure for TUI rendering. The grouping hierarchy is:
5//!
6//! ```text
7//! Project (git repo root)
8//! ├── Worktree (branch / checkout path)
9//! │   ├── Agent (session)
10//! │   │   └── Subagent (child session)
11//! │   └── Agent
12//! └── Worktree
13//!     └── ...
14//! ```
15//!
16//! **Conditional worktree nesting:** When a project has only one worktree,
17//! the worktree level is skipped — agents appear directly under the project.
18//!
19//! **Ungrouped sessions:** Sessions without a `project_root` are collected
20//! under a synthetic "Other" project node.
21//!
22//! This module is pure logic with no TUI dependency, enabling reuse
23//! in the future web UI.
24
25use std::collections::{BTreeMap, HashSet};
26
27use crate::{SessionId, SessionView};
28
29// ============================================================================
30// Tree Node Types
31// ============================================================================
32
33/// Unique identifier for a tree node, used for expand/collapse state tracking.
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub enum TreeNodeId {
36    /// A project node, keyed by project root path.
37    Project(String),
38    /// A worktree node, keyed by worktree path.
39    Worktree(String),
40    /// A team node (placeholder for CC Agent Teams).
41    Team(String),
42    /// An agent (session) node, keyed by session ID.
43    Agent(SessionId),
44}
45
46/// A node in the session grouping tree.
47// Agent variant embeds SessionView directly. Boxing it is a deferred refactor
48// that touches every consumer; revisit if/when SessionView stops being a leaf.
49#[allow(clippy::large_enum_variant)]
50#[derive(Debug, Clone)]
51pub enum TreeNode {
52    /// A git project (repo root). Groups all agents across its worktrees.
53    Project {
54        /// Display name (last path component of project_root).
55        name: String,
56        /// Full project root path.
57        root: String,
58        /// Child nodes (Worktree or Agent when single-worktree).
59        children: Vec<TreeNode>,
60    },
61    /// A git worktree within a project.
62    Worktree {
63        /// Full worktree path.
64        path: String,
65        /// Branch name (e.g., "main", "feature/auth").
66        branch: Option<String>,
67        /// Child nodes (Agent).
68        children: Vec<TreeNode>,
69    },
70    /// Placeholder for CC Agent Teams integration (Phase 6).
71    Team {
72        /// Team name.
73        name: String,
74        /// Child nodes (Agent).
75        children: Vec<TreeNode>,
76    },
77    /// A single session (leaf node).
78    Agent {
79        /// The session view data.
80        session: SessionView,
81        /// Child subagent sessions.
82        subagents: Vec<TreeNode>,
83    },
84}
85
86impl TreeNode {
87    /// Returns the total number of agent sessions in this subtree.
88    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    /// Returns true if any agent in this subtree needs attention.
100    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    /// Returns the [`TreeNodeId`] for this node.
112    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// ============================================================================
123// Flattened Row (for cursor navigation)
124// ============================================================================
125
126/// The kind of row in the flattened tree.
127#[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/// A single row in the flattened, navigable tree view.
147///
148/// Created by [`flatten_tree`] from a `Vec<TreeNode>` + expand/collapse state.
149#[derive(Debug, Clone)]
150pub struct TreeRow {
151    /// Nesting depth (0 = top-level project, 1 = worktree/agent, etc.).
152    pub depth: u8,
153    /// Unique node identifier for expand/collapse state tracking.
154    pub node_id: TreeNodeId,
155    /// What kind of row this is.
156    pub kind: TreeRowKind,
157    /// Whether this node is currently expanded (only meaningful for non-leaf nodes).
158    pub is_expanded: bool,
159    /// Total agent count in this subtree (for collapsed groups).
160    pub agent_count: usize,
161    /// Whether any agent in the subtree needs attention (bubble-up).
162    pub needs_attention: bool,
163    /// Whether this node has children (is expandable).
164    pub has_children: bool,
165}
166
167// ============================================================================
168// Tree Building
169// ============================================================================
170
171/// Label for sessions that don't belong to any git project.
172const UNGROUPED_PROJECT_NAME: &str = "Other";
173const UNGROUPED_PROJECT_ROOT: &str = "__ungrouped__";
174
175/// Builds a tree of [`TreeNode`] from a flat slice of sessions.
176///
177/// Grouping hierarchy: Project > Worktree (conditional) > Agent.
178/// Sessions without `project_root` are grouped under an "Other" project.
179/// Sessions with `parent_session_id` are nested under their parent.
180pub fn build_tree(sessions: &[SessionView]) -> Vec<TreeNode> {
181    if sessions.is_empty() {
182        return Vec::new();
183    }
184
185    // Separate parent-level sessions from subagents
186    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    // Index sessions by ID for subagent lookup
193    let by_id: BTreeMap<&str, &SessionView> = sessions.iter().map(|s| (s.id.as_str(), s)).collect();
194
195    // Group top-level sessions by project_root
196    // BTreeMap for deterministic alphabetical ordering
197    let mut by_project: BTreeMap<&str, Vec<&SessionView>> = BTreeMap::new();
198
199    for session in sessions {
200        // Skip subagents — they'll be nested under their parent
201        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        // Build agent nodes (with subagent nesting)
221        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(), // No recursive subagent nesting for now
229                })
230                .collect();
231            TreeNode::Agent {
232                session: session.clone(),
233                subagents,
234            }
235        };
236
237        // Group by worktree within this project
238        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        // Conditional worktree nesting: skip worktree level if only one
245        let skip_worktree = by_worktree.len() <= 1;
246
247        if skip_worktree {
248            // Agents directly under project
249            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            // Worktree grouping
262            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
291/// Sort agent nodes by started_at descending (newest first).
292fn 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) // newest first
303    });
304}
305
306/// Extracts a short project name from a path (last component).
307fn extract_project_name(path: &str) -> String {
308    path.rsplit('/')
309        .find(|s| !s.is_empty())
310        .unwrap_or(path)
311        .to_string()
312}
313
314// ============================================================================
315// Tree Flattening
316// ============================================================================
317
318/// Flattens a tree into navigable rows, respecting expand/collapse state.
319///
320/// # Arguments
321/// * `tree` — The tree nodes built by [`build_tree`].
322/// * `expanded` — Set of [`TreeNodeId`]s that are currently expanded.
323///   Nodes not in this set are collapsed (children hidden).
324pub 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    // Only recurse into children if expanded
373    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
386// ============================================================================
387// Helpers
388// ============================================================================
389
390/// Returns a set of all [`TreeNodeId`]s in the tree, useful for "expand all".
391pub 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// ============================================================================
417// Tests
418// ============================================================================
419
420#[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    // ------------------------------------------------------------------
455    // build_tree tests
456    // ------------------------------------------------------------------
457
458    #[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                // Should be Agent nodes directly (no Worktree nesting)
504                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                // Should have 2 Worktree nodes
537                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        // BTreeMap ordering: app-a before app-b
566        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                // Single worktree → agents directly under project
604                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    // ------------------------------------------------------------------
724    // flatten_tree tests
725    // ------------------------------------------------------------------
726
727    #[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()); // nothing expanded
753
754        // Should only see the project row (collapsed)
755        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        // Project (expanded) + 2 agents = 3 rows
787        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        // Expand everything
815        let expanded = all_node_ids(&tree);
816        let rows = flatten_tree(&tree, &expanded);
817
818        // Project + Worktree(main) + Agent(a) + Worktree(dev) + Agent(b) = 5
819        assert_eq!(rows.len(), 5);
820        assert_eq!(rows[0].depth, 0); // Project
821        assert_eq!(rows[1].depth, 1); // Worktree
822        assert_eq!(rows[2].depth, 2); // Agent
823        assert_eq!(rows[3].depth, 1); // Worktree
824        assert_eq!(rows[4].depth, 2); // Agent
825    }
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        // Project + Parent Agent + Child Agent = 3
852        assert_eq!(rows.len(), 3);
853        assert_eq!(rows[0].depth, 0); // Project
854        assert_eq!(rows[1].depth, 1); // Parent agent
855        assert_eq!(rows[2].depth, 2); // Child subagent
856    }
857
858    // ------------------------------------------------------------------
859    // Helper tests
860    // ------------------------------------------------------------------
861
862    #[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); // project + agent (no worktree since single)
885    }
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        // __ungrouped__ sorts before /home/... in BTreeMap
902        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}