fix(linux): scope foreground process discovery to panes

refs #1399

1 files changed, 191 insertions, 328 deletions

diff --git a/src/platform/linux.rs b/src/platform/linux.rs
index 2013873f..e700a7ff 100644
--- a/src/platform/linux.rs
+++ b/src/platform/linux.rs
@@ -1,11 +1,9 @@
 use std::{
-    collections::HashMap,
+    collections::{HashSet, VecDeque},
     io::Write,
     os::fd::RawFd,
     path::PathBuf,
     process::{Command, Stdio},
-    sync::Mutex,
-    time::{Duration, Instant},
 };
 
 use super::{
@@ -13,7 +11,6 @@ use super::{
     LimitedRead, Signal,
 };
 
-const FOREGROUND_MEMBERS_CACHE_TTL: Duration = Duration::from_millis(250);
 const WSL_MARKER_ENV_VARS: &[&str] = &["WSL_DISTRO_NAME", "WSL_INTEROP"];
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -22,22 +19,6 @@ struct ProcGroupMember {
     comm: String,
 }
 
-type ForegroundMembersByGroup = HashMap<u32, Vec<ProcGroupMember>>;
-
-#[derive(Debug, Clone)]
-struct CachedForegroundMembers {
-    built_at: Instant,
-    by_group: ForegroundMembersByGroup,
-}
-
-#[derive(Debug, Default)]
-struct ForegroundMembersCache {
-    cached: Option<CachedForegroundMembers>,
-}
-
-static FOREGROUND_MEMBERS_CACHE: Mutex<ForegroundMembersCache> =
-    Mutex::new(ForegroundMembersCache { cached: None });
-
 pub fn raise_server_nofile_limit() {}
 
 pub(crate) fn should_draw_host_cursor_by_default() -> bool {
@@ -108,7 +89,7 @@ fn shell_quote(value: &str) -> String {
 /// Collect the foreground terminal job for a given child PID.
 pub fn foreground_job(child_pid: u32) -> Option<ForegroundJob> {
     let tpgid = foreground_process_group_id(child_pid)?;
-    let members = foreground_process_group_members(tpgid)?;
+    let members = foreground_process_group_members(child_pid, tpgid)?;
     let processes = members
         .into_iter()
         .map(|member| {
@@ -133,19 +114,90 @@ pub fn foreground_job(child_pid: u32) -> Option<ForegroundJob> {
     })
 }
 
-fn foreground_process_group_members(process_group_id: u32) -> Option<Vec<ProcGroupMember>> {
-    let mut cache = FOREGROUND_MEMBERS_CACHE
-        .lock()
-        .unwrap_or_else(|err| err.into_inner());
-    cache.members(
+fn foreground_process_group_members(
+    child_pid: u32,
+    process_group_id: u32,
+) -> Option<Vec<ProcGroupMember>> {
+    foreground_process_group_members_with(
+        child_pid,
         process_group_id,
-        Instant::now(),
-        FOREGROUND_MEMBERS_CACHE_TTL,
-        build_foreground_members_by_group,
+        process_task_ids,
+        process_task_children,
         live_process_group_member,
     )
 }
 
+fn foreground_process_group_members_with(
+    child_pid: u32,
+    process_group_id: u32,
+    task_ids: impl FnMut(u32) -> Vec<u32>,
+    task_children: impl FnMut(u32, u32) -> Vec<u32>,
+    mut live_member: impl FnMut(u32, u32) -> Option<ProcGroupMember>,
+) -> Option<Vec<ProcGroupMember>> {
+    let mut members = process_tree_pids([child_pid, process_group_id], task_ids, task_children)
+        .into_iter()
+        .filter_map(|pid| live_member(process_group_id, pid))
+        .collect::<Vec<_>>();
+    members.sort_unstable_by_key(|member| member.pid);
+    (!members.is_empty()).then_some(members)
+}
+
+fn process_tree_pids(
+    roots: impl IntoIterator<Item = u32>,
+    mut task_ids: impl FnMut(u32) -> Vec<u32>,
+    mut task_children: impl FnMut(u32, u32) -> Vec<u32>,
+) -> Vec<u32> {
+    let mut pending = VecDeque::new();
+    let mut visited = HashSet::new();
+    for pid in roots {
+        if pid > 0 && visited.insert(pid) {
+            pending.push_back(pid);
+        }
+    }
+
+    let mut pids = Vec::new();
+    while let Some(pid) = pending.pop_front() {
+        pids.push(pid);
+        for tid in task_ids(pid) {
+            for child_pid in task_children(pid, tid) {
+                if child_pid > 0 && visited.insert(child_pid) {
+                    pending.push_back(child_pid);
+                }
+            }
+        }
+    }
+    pids
+}
+
+fn process_task_ids(pid: u32) -> Vec<u32> {
+    std::fs::read_dir(format!("/proc/{pid}/task"))
+        .into_iter()
+        .flatten()
+        .flatten()
+        .filter_map(|entry| numeric_file_name(&entry))
+        .collect()
+}
+
+fn process_task_children(pid: u32, tid: u32) -> Vec<u32> {
+    let Some(children) = std::fs::read_to_string(format!("/proc/{pid}/task/{tid}/children")).ok()
+    else {
+        return Vec::new();
+    };
+    children
+        .split_whitespace()
+        .filter_map(|child| child.parse::<u32>().ok())
+        .collect()
+}
+
+fn numeric_file_name(entry: &std::fs::DirEntry) -> Option<u32> {
+    let file_name = entry.file_name();
+    let value = file_name.to_str()?;
+    if !value.bytes().all(|byte| byte.is_ascii_digit()) {
+        return None;
+    }
+    value.parse().ok()
+}
+
 fn live_process_group_member(process_group_id: u32, pid: u32) -> Option<ProcGroupMember> {
     let (pgrp, comm) = process_pgrp_and_comm(pid)?;
     (pgrp > 0 && pgrp as u32 == process_group_id).then_some(ProcGroupMember { pid, comm })
@@ -200,94 +252,6 @@ fn process_pgrp_and_comm_from_stat(stat: &str) -> Option<(i32, String)> {
     Some((pgrp, comm))
 }
 
-fn build_foreground_members_by_group() -> ForegroundMembersByGroup {
-    let entries = std::fs::read_dir("/proc")
-        .into_iter()
-        .flatten()
-        .flatten()
-        .filter_map(|entry| {
-            let file_name = entry.file_name();
-            let pid_str = file_name.to_str()?;
-            if !pid_str.bytes().all(|b| b.is_ascii_digit()) {
-                return None;
-            }
-            let pid = pid_str.parse::<u32>().ok()?;
-            let (pgrp, comm) = process_pgrp_and_comm(pid)?;
-            Some(ProcStatEntry { pid, pgrp, comm })
-        });
-    foreground_members_by_group_from_entries(entries)
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-struct ProcStatEntry {
-    pid: u32,
-    pgrp: i32,
-    comm: String,
-}
-
-fn foreground_members_by_group_from_entries(
-    entries: impl IntoIterator<Item = ProcStatEntry>,
-) -> ForegroundMembersByGroup {
-    let mut by_group = ForegroundMembersByGroup::default();
-    for entry in entries {
-        if entry.pgrp <= 0 {
-            continue;
-        }
-        by_group
-            .entry(entry.pgrp as u32)
-            .or_default()
-            .push(ProcGroupMember {
-                pid: entry.pid,
-                comm: entry.comm,
-            });
-    }
-    by_group
-}
-
-impl ForegroundMembersCache {
-    fn members(
-        &mut self,
-        process_group_id: u32,
-        now: Instant,
-        max_age: Duration,
-        build: impl FnMut() -> ForegroundMembersByGroup,
-        mut validate: impl FnMut(u32, u32) -> Option<ProcGroupMember>,
-    ) -> Option<Vec<ProcGroupMember>> {
-        if let Some(cached) = &self.cached {
-            if now.duration_since(cached.built_at) < max_age {
-                if let Some(members) = cached.by_group.get(&process_group_id) {
-                    let members = members
-                        .iter()
-                        .filter_map(|member| validate(process_group_id, member.pid))
-                        .collect::<Vec<_>>();
-                    if !members.is_empty() {
-                        return Some(members);
-                    }
-                }
-                return self.refresh_and_get(process_group_id, now, build);
-            }
-        }
-        self.refresh_and_get(process_group_id, now, build)
-    }
-
-    fn refresh_and_get(
-        &mut self,
-        process_group_id: u32,
-        now: Instant,
-        build: impl FnOnce() -> ForegroundMembersByGroup,
-    ) -> Option<Vec<ProcGroupMember>> {
-        self.cached = Some(CachedForegroundMembers {
-            built_at: now,
-            by_group: build(),
-        });
-        self.cached
-            .as_ref()?
-            .by_group
-            .get(&process_group_id)
-            .cloned()
-    }
-}
-
 fn process_argv(pid: u32) -> Option<Vec<String>> {
     let bytes = std::fs::read(format!("/proc/{pid}/cmdline")).ok()?;
     if bytes.is_empty() {
@@ -688,55 +652,14 @@ fn process_session_id(pid: u32) -> Option<i32> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::{Mutex, OnceLock};
+    use std::{cell::RefCell, collections::HashMap};
 
     fn env_lock() -> &'static Mutex<()> {
         static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
         LOCK.get_or_init(|| Mutex::new(()))
     }
 
-    fn proc_entry(pid: u32, pgrp: i32, comm: &str) -> ProcStatEntry {
-        ProcStatEntry {
-            pid,
-            pgrp,
-            comm: comm.to_string(),
-        }
-    }
-
-    fn foreground_members(groups: &[(u32, &str, i32)]) -> ForegroundMembersByGroup {
-        foreground_members_by_group_from_entries(
-            groups
-                .iter()
-                .map(|(pid, comm, pgrp)| proc_entry(*pid, *pgrp, comm)),
-        )
-    }
-
-    fn validate_from<'a>(
-        groups: &'a [(u32, &'a str, i32)],
-    ) -> impl FnMut(u32, u32) -> Option<ProcGroupMember> + 'a {
-        move |process_group_id, pid| {
-            groups.iter().find_map(|(member_pid, comm, pgrp)| {
-                (*member_pid == pid && *pgrp > 0 && *pgrp as u32 == process_group_id).then(|| {
-                    ProcGroupMember {
-                        pid,
-                        comm: (*comm).to_string(),
-                    }
-                })
-            })
-        }
-    }
-
-    fn member_names(by_group: &ForegroundMembersByGroup, process_group_id: u32) -> Vec<String> {
-        by_group
-            .get(&process_group_id)
-            .cloned()
-            .unwrap_or_default()
-            .into_iter()
-            .map(|member| member.comm)
-            .collect()
-    }
-
     #[test]
     fn wsl_marker_detection_matches_kernel_release_text() {
         assert!(text_indicates_wsl("5.15.167.4-microsoft-standard-WSL2"));
@@ -746,195 +669,135 @@ mod tests {
     }
 
     #[test]
-    fn foreground_members_indexes_by_process_group() {
-        let by_group = foreground_members(&[
-            (10, "shell", 10),
-            (11, "claude", 11),
-            (12, "node", 11),
-            (13, "ignored", -1),
+    fn foreground_members_follow_the_pane_tree_and_filter_by_process_group() {
+        let tasks = HashMap::from([
+            (100, vec![100, 101]),
+            (200, vec![200]),
+            (201, vec![201]),
+            (210, vec![210]),
+            (220, vec![220]),
+            (221, vec![221]),
+            (300, vec![300]),
         ]);
-
-        assert_eq!(member_names(&by_group, 10), vec!["shell"]);
-        assert_eq!(member_names(&by_group, 11), vec!["claude", "node"]);
-        assert_eq!(by_group.get(&13), None);
-    }
-
-    #[test]
-    fn foreground_members_cache_reuses_snapshot_inside_ttl() {
-        let mut cache = ForegroundMembersCache::default();
-        let now = Instant::now();
-        let builds = AtomicUsize::new(0);
-
-        let first = cache.members(
-            10,
-            now,
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(10, "shell", 10)])
-            },
-            validate_from(&[(10, "shell", 10)]),
-        );
-        let second = cache.members(
-            10,
-            now + Duration::from_millis(100),
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(20, "new", 20)])
+        let children = HashMap::from([
+            ((100, 100), vec![200, 201, 300]),
+            ((100, 101), vec![210]),
+            ((200, 200), vec![220]),
+            ((220, 220), vec![221]),
+        ]);
+        let processes = HashMap::from([
+            (100, (100, "shell")),
+            (200, (200, "leader")),
+            (201, (200, "pipeline")),
+            (210, (200, "thread-child")),
+            (220, (220, "intermediate")),
+            (221, (200, "nested-agent")),
+            (300, (300, "background")),
+            (9999, (200, "unrelated-host-process")),
+        ]);
+        let task_reads = RefCell::new(Vec::new());
+        let child_reads = RefCell::new(Vec::new());
+        let member_reads = RefCell::new(Vec::new());
+
+        let members = foreground_process_group_members_with(
+            100,
+            200,
+            |pid| {
+                task_reads.borrow_mut().push(pid);
+                tasks.get(&pid).cloned().unwrap_or_default()
             },
-            validate_from(&[(10, "shell-live", 10)]),
-        );
-
-        assert_eq!(
-            first
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["shell"]
-        );
-        assert_eq!(
-            second
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["shell-live"]
-        );
-        assert_eq!(builds.load(Ordering::Relaxed), 1);
-    }
-
-    #[test]
-    fn foreground_members_cache_rebuilds_after_ttl() {
-        let mut cache = ForegroundMembersCache::default();
-        let now = Instant::now();
-        let builds = AtomicUsize::new(0);
-
-        let _ = cache.members(
-            10,
-            now,
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(10, "shell", 10)])
+            |pid, tid| {
+                child_reads.borrow_mut().push((pid, tid));
+                children.get(&(pid, tid)).cloned().unwrap_or_default()
             },
-            validate_from(&[(10, "shell", 10)]),
-        );
-        let second = cache.members(
-            20,
-            now + FOREGROUND_MEMBERS_CACHE_TTL,
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(20, "new", 20)])
+            |process_group_id, pid| {
+                member_reads.borrow_mut().push(pid);
+                let (pgrp, comm) = processes.get(&pid)?;
+                (*pgrp == process_group_id).then(|| ProcGroupMember {
+                    pid,
+                    comm: (*comm).to_string(),
+                })
             },
-            validate_from(&[(20, "new", 20)]),
-        );
+        )
+        .unwrap();
 
-        assert_eq!(builds.load(Ordering::Relaxed), 2);
         assert_eq!(
-            second
-                .unwrap()
+            members
                 .into_iter()
-                .map(|member| member.comm)
+                .map(|member| (member.pid, member.comm))
                 .collect::<Vec<_>>(),
-            vec!["new"]
+            vec![
+                (200, "leader".to_string()),
+                (201, "pipeline".to_string()),
+                (210, "thread-child".to_string()),
+                (221, "nested-agent".to_string()),
+            ]
         );
+        assert!(child_reads.borrow().contains(&(100, 101)));
+        assert!(task_reads.borrow().contains(&220));
+        assert!(!task_reads.borrow().contains(&9999));
+        assert!(!member_reads.borrow().contains(&9999));
     }
 
     #[test]
-    fn foreground_members_cache_refreshes_reused_snapshot_on_group_miss() {
-        let mut cache = ForegroundMembersCache::default();
-        let now = Instant::now();
-        let builds = AtomicUsize::new(0);
-
-        let stale = cache.members(
-            10,
-            now,
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(10, "shell", 10)])
-            },
-            validate_from(&[(10, "shell", 10)]),
-        );
-        assert_eq!(
-            stale
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["shell"]
-        );
-
-        let refreshed = cache.members(
-            42,
-            now + Duration::from_millis(10),
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(42, "claude", 42)])
+    fn foreground_members_degrade_to_the_direct_group_leader() {
+        let members = foreground_process_group_members_with(
+            100,
+            200,
+            |_| Vec::new(),
+            |_, _| Vec::new(),
+            |process_group_id, pid| {
+                (pid == process_group_id).then(|| ProcGroupMember {
+                    pid,
+                    comm: "leader".to_string(),
+                })
             },
-            validate_from(&[(42, "claude", 42)]),
-        );
+        )
+        .unwrap();
 
-        assert_eq!(builds.load(Ordering::Relaxed), 2);
         assert_eq!(
-            refreshed
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["claude"]
+            members,
+            vec![ProcGroupMember {
+                pid: 200,
+                comm: "leader".to_string()
+            }]
         );
     }
 
     #[test]
-    fn foreground_members_cache_refreshes_existing_group_when_cached_members_are_stale() {
-        let mut cache = ForegroundMembersCache::default();
-        let now = Instant::now();
-        let builds = AtomicUsize::new(0);
-
-        let stale = cache.members(
-            42,
-            now,
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(10, "old", 42)])
-            },
-            validate_from(&[(10, "old", 42)]),
-        );
-        assert_eq!(
-            stale
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["old"]
-        );
-
-        let refreshed = cache.members(
-            42,
-            now + Duration::from_millis(10),
-            FOREGROUND_MEMBERS_CACHE_TTL,
-            || {
-                builds.fetch_add(1, Ordering::Relaxed);
-                foreground_members(&[(20, "new", 42)])
-            },
-            validate_from(&[(10, "old", 7), (20, "new", 42)]),
-        );
+    fn foreground_members_observe_new_children_without_a_snapshot_cache() {
+        let children = RefCell::new(HashMap::from([((100, 100), vec![200])]));
+        let discover = || {
+            foreground_process_group_members_with(
+                100,
+                200,
+                |pid| vec![pid],
+                |pid, tid| {
+                    children
+                        .borrow()
+                        .get(&(pid, tid))
+                        .cloned()
+                        .unwrap_or_default()
+                },
+                |process_group_id, pid| {
+                    [200, 201]
+                        .contains(&pid)
+                        .then(|| ProcGroupMember {
+                            pid,
+                            comm: format!("member-{pid}"),
+                        })
+                        .filter(|_| process_group_id == 200)
+                },
+            )
+            .unwrap()
+            .into_iter()
+            .map(|member| member.pid)
+            .collect::<Vec<_>>()
+        };
 
-        assert_eq!(builds.load(Ordering::Relaxed), 2);
-        assert_eq!(
-            refreshed
-                .unwrap()
-                .into_iter()
-                .map(|member| member.comm)
-                .collect::<Vec<_>>(),
-            vec!["new"]
-        );
+        assert_eq!(discover(), vec![200]);
+        children.borrow_mut().insert((100, 100), vec![200, 201]);
+        assert_eq!(discover(), vec![200, 201]);
     }
 
     #[test]