wsx_core/
proc_tree.rs

1// Process tree snapshot for tmux pane classification.
2// One `ps -ax -o pid,ppid,comm` call per refresh; walk descendants of each
3// pane PID so an agent nested under shells/subprocesses still gets classified
4// as Agent (tmux's `pane_current_command` reports the deepest spawned child,
5// which can mask the real foreground — e.g. claude shows up as a node child
6// named "2.1.x" while the actual claude process sits mid-tree).
7
8use std::collections::HashMap;
9use std::process::Command;
10
11#[derive(Debug, Clone, Default)]
12pub struct ProcTree {
13    children: HashMap<u32, Vec<u32>>,
14    comm: HashMap<u32, String>,
15}
16
17impl ProcTree {
18    /// Snapshot the live process table. Returns an empty tree if `ps` fails.
19    pub fn snapshot() -> Self {
20        let Ok(out) = Command::new("ps")
21            .args(["-ax", "-o", "pid=,ppid=,comm="])
22            .output()
23        else {
24            return Self::default();
25        };
26        let text = String::from_utf8_lossy(&out.stdout);
27        Self::parse(&text)
28    }
29
30    /// Parse the raw `ps` output. Tolerates leading whitespace and `comm`
31    /// values that include path separators or spaces.
32    pub fn parse(raw: &str) -> Self {
33        let mut tree = Self::default();
34        for line in raw.lines() {
35            let mut iter = line.split_whitespace();
36            let Some(pid) = iter.next().and_then(|s| s.parse::<u32>().ok()) else {
37                continue;
38            };
39            let Some(ppid) = iter.next().and_then(|s| s.parse::<u32>().ok()) else {
40                continue;
41            };
42            let comm: String = iter.collect::<Vec<_>>().join(" ");
43            if comm.is_empty() {
44                continue;
45            }
46            tree.comm.insert(pid, comm);
47            tree.children.entry(ppid).or_default().push(pid);
48        }
49        tree
50    }
51
52    /// Pure builder for tests.
53    #[cfg(test)]
54    pub fn from_rows(rows: &[(u32, u32, &str)]) -> Self {
55        let mut tree = Self::default();
56        for (pid, ppid, comm) in rows {
57            tree.comm.insert(*pid, (*comm).to_string());
58            tree.children.entry(*ppid).or_default().push(*pid);
59        }
60        tree
61    }
62
63    /// Iterate `root` and all its descendants, yielding `(pid, comm)` per
64    /// process. `root` itself is yielded first if it exists in the tree.
65    pub fn descendants(&self, root: u32) -> Vec<(u32, &str)> {
66        let mut out: Vec<(u32, &str)> = Vec::new();
67        let mut stack = vec![root];
68        while let Some(pid) = stack.pop() {
69            if let Some(c) = self.comm.get(&pid) {
70                out.push((pid, c.as_str()));
71            }
72            if let Some(kids) = self.children.get(&pid) {
73                stack.extend(kids.iter().copied());
74            }
75        }
76        out
77    }
78}
79
80/// Strip a `comm` value down to its bare executable name for classification.
81/// macOS `ps -o comm` reports full paths (e.g. `/bin/zsh`); login shells
82/// arrive prefixed with `-` (e.g. `-zsh`). Both should match `Shell`.
83pub fn normalize_comm(comm: &str) -> &str {
84    let basename = comm.rsplit('/').next().unwrap_or(comm);
85    basename.strip_prefix('-').unwrap_or(basename)
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    // ── parse ────────────────────────────────────────────────────────────────
93
94    #[test]
95    fn given_single_line_when_parsed_then_pid_registered() {
96        let tree = ProcTree::parse("100 1 zsh");
97        let desc = tree.descendants(100);
98        assert_eq!(desc, vec![(100u32, "zsh")]);
99    }
100
101    #[test]
102    fn given_multiple_lines_when_parsed_then_all_pids_registered() {
103        let raw = "100 1 zsh\n200 100 claude\n300 200 node";
104        let tree = ProcTree::parse(raw);
105        let mut desc: Vec<(u32, &str)> = tree.descendants(100);
106        desc.sort_by_key(|(p, _)| *p);
107        assert_eq!(desc, vec![(100, "zsh"), (200, "claude"), (300, "node")]);
108    }
109
110    #[test]
111    fn given_leading_whitespace_when_parsed_then_pid_registered() {
112        let tree = ProcTree::parse("   100 1 zsh");
113        assert_eq!(tree.descendants(100), vec![(100u32, "zsh")]);
114    }
115
116    #[test]
117    fn given_trailing_newline_and_blank_lines_when_parsed_then_tolerated() {
118        let raw = "100 1 zsh\n\n   \n200 100 bash\n";
119        let tree = ProcTree::parse(raw);
120        let mut desc: Vec<(u32, &str)> = tree.descendants(100);
121        desc.sort_by_key(|(p, _)| *p);
122        assert_eq!(desc, vec![(100, "zsh"), (200, "bash")]);
123    }
124
125    #[test]
126    fn given_tab_separator_when_parsed_then_pid_registered() {
127        let tree = ProcTree::parse("100\t1\tzsh");
128        assert_eq!(tree.descendants(100), vec![(100u32, "zsh")]);
129    }
130
131    #[test]
132    fn given_comm_with_path_separator_when_parsed_then_full_path_kept() {
133        let tree = ProcTree::parse("100 1 /bin/zsh");
134        assert_eq!(tree.descendants(100), vec![(100u32, "/bin/zsh")]);
135    }
136
137    #[test]
138    fn given_comm_with_internal_space_when_parsed_then_space_preserved() {
139        let tree = ProcTree::parse("100 1 my prog");
140        assert_eq!(tree.descendants(100), vec![(100u32, "my prog")]);
141    }
142
143    #[test]
144    fn given_non_numeric_pid_when_parsed_then_line_skipped() {
145        let tree = ProcTree::parse("abc 1 zsh\n100 1 bash");
146        assert_eq!(tree.descendants(100), vec![(100u32, "bash")]);
147    }
148
149    #[test]
150    fn given_non_numeric_ppid_when_parsed_then_line_skipped() {
151        let tree = ProcTree::parse("100 xyz zsh");
152        assert!(tree.descendants(100).is_empty());
153    }
154
155    #[test]
156    fn given_empty_comm_when_parsed_then_line_skipped() {
157        let tree = ProcTree::parse("100 1");
158        assert!(tree.descendants(100).is_empty());
159    }
160
161    #[test]
162    fn given_empty_input_when_parsed_then_tree_is_empty() {
163        let tree = ProcTree::parse("");
164        assert!(tree.descendants(1).is_empty());
165        assert!(tree.descendants(0).is_empty());
166    }
167
168    // ── descendants ──────────────────────────────────────────────────────────
169
170    #[test]
171    fn given_chain_when_descendants_called_then_all_returned() {
172        let tree =
173            ProcTree::from_rows(&[(10, 1, "root"), (20, 10, "a"), (30, 20, "b"), (40, 30, "c")]);
174        let mut pids: Vec<u32> = tree.descendants(10).into_iter().map(|(p, _)| p).collect();
175        pids.sort();
176        assert_eq!(pids, vec![10, 20, 30, 40]);
177    }
178
179    #[test]
180    fn given_chain_when_descendants_called_then_root_returned_first() {
181        let tree = ProcTree::from_rows(&[(10, 1, "root"), (20, 10, "a"), (30, 20, "b")]);
182        let pids: Vec<u32> = tree.descendants(10).into_iter().map(|(p, _)| p).collect();
183        assert_eq!(pids.first().copied(), Some(10));
184    }
185
186    #[test]
187    fn given_root_with_no_children_when_descendants_called_then_only_root_returned() {
188        let tree = ProcTree::from_rows(&[(10, 1, "root")]);
189        assert_eq!(tree.descendants(10), vec![(10u32, "root")]);
190    }
191
192    #[test]
193    fn given_absent_root_when_descendants_called_then_empty() {
194        let tree = ProcTree::from_rows(&[(10, 1, "root")]);
195        assert!(tree.descendants(999).is_empty());
196    }
197
198    #[test]
199    fn given_multi_fanout_when_descendants_called_then_all_children_returned() {
200        let tree =
201            ProcTree::from_rows(&[(10, 1, "root"), (20, 10, "a"), (30, 10, "b"), (40, 10, "c")]);
202        let mut pids: Vec<u32> = tree.descendants(10).into_iter().map(|(p, _)| p).collect();
203        pids.sort();
204        assert_eq!(pids, vec![10, 20, 30, 40]);
205    }
206
207    #[test]
208    fn given_orphan_ppid_row_when_descendants_called_from_orphan_then_subtree_returned() {
209        // orphan(200) claims ppid=999 which has no comm — reachable from its
210        // own pid; descendants(999) traverses the children link but yields no
211        // entry for the phantom parent.
212        let tree = ProcTree::from_rows(&[(200, 999, "claude"), (300, 200, "node")]);
213        let mut pids: Vec<u32> = tree.descendants(200).into_iter().map(|(p, _)| p).collect();
214        pids.sort();
215        assert_eq!(pids, vec![200, 300]);
216        let mut phantom: Vec<u32> = tree.descendants(999).into_iter().map(|(p, _)| p).collect();
217        phantom.sort();
218        assert_eq!(phantom, vec![200, 300]);
219    }
220
221    // ── from_rows ────────────────────────────────────────────────────────────
222
223    #[test]
224    fn given_same_data_when_from_rows_and_parse_then_descendants_match() {
225        let rows: &[(u32, u32, &str)] =
226            &[(100, 1, "zsh"), (200, 100, "claude"), (300, 200, "node")];
227        let from_rows = ProcTree::from_rows(rows);
228        let parsed = ProcTree::parse("100 1 zsh\n200 100 claude\n300 200 node");
229        let mut a: Vec<(u32, String)> = from_rows
230            .descendants(100)
231            .into_iter()
232            .map(|(p, c)| (p, c.to_string()))
233            .collect();
234        let mut b: Vec<(u32, String)> = parsed
235            .descendants(100)
236            .into_iter()
237            .map(|(p, c)| (p, c.to_string()))
238            .collect();
239        a.sort();
240        b.sort();
241        assert_eq!(a, b);
242    }
243
244    // ── normalize_comm ───────────────────────────────────────────────────────
245
246    #[test]
247    fn given_bare_name_when_normalized_then_unchanged() {
248        assert_eq!(normalize_comm("claude"), "claude");
249    }
250
251    #[test]
252    fn given_full_path_when_normalized_then_basename_returned() {
253        assert_eq!(normalize_comm("/bin/zsh"), "zsh");
254    }
255
256    #[test]
257    fn given_dash_prefix_when_normalized_then_dash_stripped() {
258        assert_eq!(normalize_comm("-zsh"), "zsh");
259    }
260
261    #[test]
262    fn given_dash_and_path_when_normalized_then_both_stripped() {
263        assert_eq!(normalize_comm("-/usr/bin/bash"), "bash");
264    }
265
266    #[test]
267    fn given_double_dash_prefix_when_normalized_then_only_one_dash_stripped() {
268        assert_eq!(normalize_comm("--zsh"), "-zsh");
269    }
270
271    #[test]
272    fn given_empty_string_when_normalized_then_empty() {
273        assert_eq!(normalize_comm(""), "");
274    }
275}
wsx_core/proc_tree.rs

wsx_core/
proc_tree.rs
