motosan_agent_loop/session/
tree.rs

1//! Tree projection — turn the raw entry log into a navigable branch tree.
2//!
3//! This is the read-only counterpart of [`crate::session::projection`]: where
4//! the message projection renders one branch for the LLM, this renders the
5//! whole tree for a UI to navigate (`/tree`). Pure and side-effect free.
6
7use crate::session::entry::{EntryId, SessionEntry, StoredEntry};
8
9/// One node of a [`BranchTree`], index-aligned with the entry log.
10#[derive(Debug, Clone)]
11pub struct BranchNode {
12    /// Id of the entry this node represents.
13    pub id: EntryId,
14    /// Index of the parent node, or `None` for the root.
15    pub parent: Option<usize>,
16    /// Indices of child nodes, in log order.
17    pub children: Vec<usize>,
18    /// A short human label — the first line of a message, or `[kind]` for a
19    /// custom entry.
20    pub label: String,
21}
22
23/// A navigable view of a session log as a branch tree.
24#[derive(Debug, Clone)]
25pub struct BranchTree {
26    /// Nodes, index-aligned with the `entries` slice passed to
27    /// [`project_branches`].
28    pub nodes: Vec<BranchNode>,
29    /// Index of the root node, or `None` for an empty log.
30    pub root: Option<usize>,
31    /// Index of the active leaf — always the last log entry (pi semantics:
32    /// the active branch is whichever one was appended to last).
33    pub active_leaf: Option<usize>,
34}
35
36fn branch_label(entry: &SessionEntry) -> String {
37    match entry {
38        SessionEntry::Message { message } => message
39            .text()
40            .lines()
41            .next()
42            .unwrap_or("")
43            .chars()
44            .take(80)
45            .collect(),
46        SessionEntry::Custom { kind, .. } => format!("[{kind}]"),
47    }
48}
49
50/// Build a [`BranchTree`] from a raw entry log.
51///
52/// Parent resolution matches the projection walk: `Some(id)` resolves to that
53/// entry; `None` resolves to the positional predecessor (`None` at index 0 is
54/// the root). A dangling `Some(id)` yields a `None` parent (an extra root) and
55/// is otherwise tolerated.
56pub fn project_branches(entries: &[StoredEntry]) -> BranchTree {
57    let mut nodes: Vec<BranchNode> = Vec::with_capacity(entries.len());
58    for (i, stored) in entries.iter().enumerate() {
59        let parent = match &stored.parent_id {
60            Some(pid) => entries[..i].iter().rposition(|e| &e.id == pid),
61            None => {
62                if i == 0 {
63                    None
64                } else {
65                    Some(i - 1)
66                }
67            }
68        };
69        nodes.push(BranchNode {
70            id: stored.id.clone(),
71            parent,
72            children: Vec::new(),
73            label: branch_label(&stored.entry),
74        });
75    }
76    for i in 0..nodes.len() {
77        if let Some(p) = nodes[i].parent {
78            nodes[p].children.push(i);
79        }
80    }
81    BranchTree {
82        root: if entries.is_empty() { None } else { Some(0) },
83        active_leaf: if entries.is_empty() {
84            None
85        } else {
86            Some(entries.len() - 1)
87        },
88        nodes,
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use crate::message::Message;
96
97    fn se(id: &str, entry: SessionEntry) -> StoredEntry {
98        StoredEntry::new(id.to_string(), entry)
99    }
100
101    #[test]
102    fn empty_log_yields_empty_tree() {
103        let tree = project_branches(&[]);
104        assert!(tree.root.is_none());
105        assert!(tree.active_leaf.is_none());
106        assert!(tree.nodes.is_empty());
107    }
108
109    #[test]
110    fn flat_log_is_a_linear_chain() {
111        let entries = vec![
112            se(
113                "1",
114                SessionEntry::message(Message::user("first line\nsecond")),
115            ),
116            se("2", SessionEntry::message(Message::assistant("reply"))),
117        ];
118        let tree = project_branches(&entries);
119        assert_eq!(tree.root, Some(0));
120        assert_eq!(tree.active_leaf, Some(1));
121        assert_eq!(tree.nodes[0].parent, None);
122        assert_eq!(tree.nodes[0].children, vec![1]);
123        assert_eq!(tree.nodes[1].parent, Some(0));
124        assert!(tree.nodes[1].children.is_empty());
125        assert_eq!(tree.nodes[0].label, "first line");
126    }
127
128    #[test]
129    fn branch_point_has_two_children() {
130        // A is the branch point: B continues it linearly, D forks from it.
131        let entries = vec![
132            se("A", SessionEntry::message(Message::user("a"))),
133            se("B", SessionEntry::message(Message::assistant("b"))),
134            StoredEntry::with_parent(
135                "D".to_string(),
136                Some("A".to_string()),
137                SessionEntry::message(Message::assistant("d")),
138            ),
139        ];
140        let tree = project_branches(&entries);
141        assert_eq!(tree.nodes[0].children, vec![1, 2]);
142        assert_eq!(tree.nodes[2].parent, Some(0));
143        assert_eq!(tree.active_leaf, Some(2));
144    }
145
146    #[test]
147    fn custom_entry_is_labelled_by_kind() {
148        let entries = vec![se(
149            "c",
150            SessionEntry::custom("compaction", serde_json::json!({})),
151        )];
152        let tree = project_branches(&entries);
153        assert_eq!(tree.nodes[0].label, "[compaction]");
154    }
155}
motosan_agent_loop/session/tree.rs

motosan_agent_loop/session/
tree.rs
