diff --git a/zellij-server/src/os_input_output.rs b/zellij-server/src/os_input_output.rs index dd7225ad8b..5e11ab524e 100644 --- a/zellij-server/src/os_input_output.rs +++ b/zellij-server/src/os_input_output.rs @@ -548,9 +548,12 @@ impl ServerOsApi for ServerOsInputOutput { #[cfg(unix)] fn get_all_cmds_by_ppid(&self, post_hook: &Option) -> HashMap> { // the key is the stringified ppid - let mut cmds = HashMap::new(); + // We collect all children per ppid and prefer the foreground process + // (STAT contains '+') to avoid the last child in ps output silently + // overwriting earlier ones via HashMap::insert. + let mut candidates: HashMap)>> = HashMap::new(); if let Some(output) = Command::new("ps") - .args(vec!["-ao", "ppid,args"]) + .args(vec!["-ao", "ppid,stat,args"]) .output() .ok() { @@ -564,7 +567,9 @@ impl ServerOsApi for ServerOsInputOutput { .collect(); let mut line_parts = line_parts.into_iter(); let ppid = line_parts.next(); - if let Some(ppid) = ppid { + let stat = line_parts.next(); + if let (Some(ppid), Some(stat)) = (ppid, stat) { + let is_foreground = stat.contains('+'); match &post_hook { Some(post_hook) => { let command: Vec = line_parts.clone().collect(); @@ -581,16 +586,21 @@ impl ServerOsApi for ServerOsInputOutput { .split_ascii_whitespace() .map(|p| p.to_owned()) .collect(); - cmds.insert(ppid.into(), line_parts); + if !line_parts.is_empty() { + candidates.entry(ppid).or_default().push((is_foreground, line_parts)); + } }, None => { - cmds.insert(ppid.into(), line_parts.collect()); + let parts: Vec = line_parts.collect(); + if !parts.is_empty() { + candidates.entry(ppid).or_default().push((is_foreground, parts)); + } }, } } } } - cmds + select_best_candidates(candidates) } #[cfg(not(unix))] @@ -598,7 +608,10 @@ impl ServerOsApi for ServerOsInputOutput { let mut system_info = System::new(); let refresh_kind = ProcessRefreshKind::nothing().with_cmd(UpdateKind::Always); system_info.refresh_processes_specifics(ProcessesToUpdate::All, true, refresh_kind); - let mut cmds = HashMap::new(); + // Use Vec to collect all children per ppid, same as the unix path, + // to avoid HashMap::insert silently overwriting earlier entries. + // STAT/foreground info is unavailable via sysinfo, so we use false. + let mut candidates: HashMap)>> = HashMap::new(); for (_pid, process) in system_info.processes() { if let Some(parent_pid) = process.parent() { let ppid_str = format!("{}", parent_pid); @@ -625,15 +638,17 @@ impl ServerOsApi for ServerOsInputOutput { .split_ascii_whitespace() .map(|p| p.to_owned()) .collect(); - cmds.insert(ppid_str, line_parts); + if !line_parts.is_empty() { + candidates.entry(ppid_str).or_default().push((false, line_parts)); + } }, None => { - cmds.insert(ppid_str, command); + candidates.entry(ppid_str).or_default().push((false, command)); }, } } } - cmds + select_best_candidates(candidates) } fn write_to_file(&mut self, buf: String, name: Option) -> Result<()> { @@ -759,6 +774,23 @@ fn run_command_hook( Ok(String::from_utf8(output.stdout)?.trim().to_string()) } +/// Given a map of ppid -> [(is_foreground, command_args)], select the best +/// candidate per ppid: prefer a foreground process, fall back to first child. +fn select_best_candidates( + candidates: HashMap)>>, +) -> HashMap> { + let mut cmds = HashMap::new(); + for (ppid, children) in candidates { + if children.is_empty() { + continue; + } + let idx = children.iter().position(|(fg, _)| *fg).unwrap_or(0); + let (_, cmd) = children.into_iter().nth(idx).unwrap(); + cmds.insert(ppid, cmd); + } + cmds +} + #[cfg(test)] #[path = "./unit/os_input_output_tests.rs"] mod os_input_output_tests; diff --git a/zellij-server/src/unit/os_input_output_tests.rs b/zellij-server/src/unit/os_input_output_tests.rs index 994ed98421..d73153fc88 100644 --- a/zellij-server/src/unit/os_input_output_tests.rs +++ b/zellij-server/src/unit/os_input_output_tests.rs @@ -1,4 +1,5 @@ use super::*; +use super::select_best_candidates; use zellij_utils::input::command::RunCommand; fn make_server() -> ServerOsInputOutput { @@ -201,3 +202,81 @@ fn spawn_and_read_output() { output_str ); } + +// --- select_best_candidates tests --- + +fn candidates_from(entries: Vec<(&str, Vec<(bool, Vec<&str>)>)>) -> HashMap)>> { + entries + .into_iter() + .map(|(ppid, children)| { + let children = children + .into_iter() + .map(|(fg, args)| (fg, args.into_iter().map(String::from).collect())) + .collect(); + (ppid.to_string(), children) + }) + .collect() +} + +#[test] +fn foreground_process_preferred_over_background_children() { + let candidates = candidates_from(vec![ + ("1234", vec![ + (true, vec!["claude", "--resume", "my-session"]), + (false, vec!["node", "/path/to/mcp-server"]), + ]), + ]); + let cmds = select_best_candidates(candidates); + assert_eq!(cmds.get("1234").unwrap(), &["claude", "--resume", "my-session"]); +} + +#[test] +fn background_process_listed_first_does_not_win() { + let candidates = candidates_from(vec![ + ("1234", vec![ + (false, vec!["node", "/path/to/mcp-server-1"]), + (false, vec!["node", "/path/to/mcp-server-2"]), + (true, vec!["claude", "--resume", "my-session"]), + ]), + ]); + let cmds = select_best_candidates(candidates); + assert_eq!(cmds.get("1234").unwrap(), &["claude", "--resume", "my-session"]); +} + +#[test] +fn single_child_returned_regardless_of_foreground() { + let candidates = candidates_from(vec![ + ("5678", vec![(false, vec!["nvim", "main.rs"])]), + ]); + let cmds = select_best_candidates(candidates); + assert_eq!(cmds.get("5678").unwrap(), &["nvim", "main.rs"]); +} + +#[test] +fn multiple_ppids_handled_independently() { + let candidates = candidates_from(vec![ + ("100", vec![ + (false, vec!["node", "mcp-server"]), + (true, vec!["claude", "--resume", "foo"]), + ]), + ("200", vec![ + (true, vec!["nvim", "bar.rs"]), + (false, vec!["node", "lsp-server"]), + ]), + ]); + let cmds = select_best_candidates(candidates); + assert_eq!(cmds.get("100").unwrap(), &["claude", "--resume", "foo"]); + assert_eq!(cmds.get("200").unwrap(), &["nvim", "bar.rs"]); +} + +#[test] +fn no_foreground_falls_back_to_first_child() { + let candidates = candidates_from(vec![ + ("300", vec![ + (false, vec!["node", "server-a"]), + (false, vec!["node", "server-b"]), + ]), + ]); + let cmds = select_best_candidates(candidates); + assert_eq!(cmds.get("300").unwrap(), &["node", "server-a"]); +}