From ef0b4bc6e201c137b22da4604590d28ebae7c2bd Mon Sep 17 00:00:00 2001 From: Ogulcan Celik Date: Sat, 20 Jun 2026 03:15:21 +0300 Subject: fix: support agent env hints refs #679 related #699 --- src/pane.rs | 306 ++++++++++++++++++++++++++++++++++++++++++-------- src/platform/linux.rs | 38 +++++++ src/platform/mod.rs | 5 + 3 files changed, 302 insertions(+), 47 deletions(-) (limited to 'src') diff --git a/src/pane.rs b/src/pane.rs index 468eaad6..d616e5a1 100644 --- a/src/pane.rs +++ b/src/pane.rs @@ -234,14 +234,6 @@ struct AgentDetectionPresence { consecutive_misses: u8, } -fn should_clear_agent_for_foreground_shell( - previous_agent: Option, - new_agent: Option, - foreground_is_pane_shell: bool, -) -> bool { - previous_agent.is_some() && new_agent.is_none() && foreground_is_pane_shell -} - #[cfg(unix)] fn usable_process_cwd(pid: u32) -> Option { crate::platform::process_cwd(pid).filter(|cwd| cwd.is_absolute() && cwd.is_dir()) @@ -280,19 +272,22 @@ fn foreground_shell_agent_action( foreground_is_pane_shell: bool, process_exit_reported: bool, ) -> ForegroundShellAgentAction { - if !should_clear_agent_for_foreground_shell(previous_agent, new_agent, foreground_is_pane_shell) - { + if previous_agent.is_none() || new_agent.is_some() { return ForegroundShellAgentAction::ObserveProbe; } - // Do not clear identity immediately. First publish an idle process-exit - // transition for the previous agent so notifications and wait-agent callers - // observe completion before the pane becomes unknown. if process_exit_reported { - ForegroundShellAgentAction::ClearAgent - } else { - ForegroundShellAgentAction::ReportProcessExit + return ForegroundShellAgentAction::ClearAgent; } + + if foreground_is_pane_shell { + // Do not clear identity immediately. First publish an idle process-exit + // transition for the previous agent so notifications and wait-agent callers + // observe completion before the pane becomes unknown. + return ForegroundShellAgentAction::ReportProcessExit; + } + + ForegroundShellAgentAction::ObserveProbe } #[derive(Debug, Clone, Copy)] @@ -412,23 +407,108 @@ struct ProcessProbeResult { process_name: Option, } -fn probe_foreground_process(pid: u32, foreground_pgid: Option) -> ProcessProbeResult { - if let Some(job) = foreground_pgid.and_then(crate::detect::foreground_group_leader_job) { - if let Some((agent, process_name)) = crate::detect::identify_agent_in_job(&job) { - return ProcessProbeResult { - process_group_id: Some(job.process_group_id), - foreground_is_pane_shell: job.processes.iter().any(|p| p.pid == pid), - agent: Some(agent), - process_name: Some(process_name), - }; +fn agent_hint_for_foreground_job_members( + job: &crate::platform::ForegroundJob, + read_hint: impl Fn(u32) -> Option, +) -> Option { + read_hint(job.process_group_id) + .or_else(|| agent_hint_for_non_leader_foreground_job_members(job, read_hint)) +} + +fn agent_hint_for_non_leader_foreground_job_members( + job: &crate::platform::ForegroundJob, + read_hint: impl Fn(u32) -> Option, +) -> Option { + job.processes + .iter() + .filter(|process| process.pid != job.process_group_id) + .find_map(|process| read_hint(process.pid)) +} + +fn identify_process_group_leader_in_job( + job: &crate::platform::ForegroundJob, +) -> Option<(Agent, String)> { + let leader = job + .processes + .iter() + .find(|process| process.pid == job.process_group_id)?; + let leader_job = crate::platform::ForegroundJob { + process_group_id: job.process_group_id, + processes: vec![leader.clone()], + }; + crate::detect::identify_agent_in_job(&leader_job) +} + +fn process_probe_result( + job: &crate::platform::ForegroundJob, + pid: u32, + agent: Agent, + process_name: String, +) -> ProcessProbeResult { + ProcessProbeResult { + process_group_id: Some(job.process_group_id), + foreground_is_pane_shell: job.processes.iter().any(|process| process.pid == pid), + agent: Some(agent), + process_name: Some(process_name), + } +} + +fn hinted_process_probe_result( + job: &crate::platform::ForegroundJob, + pid: u32, + read_hint: impl Fn(u32) -> Option, +) -> Option { + let agent = agent_hint_for_foreground_job_members(job, read_hint)?; + Some(process_probe_result( + job, + pid, + agent, + crate::detect::agent_label(agent).to_string(), + )) +} + +fn probe_foreground_process_from_jobs( + pid: u32, + foreground_pgid: Option, + leader_job: Option, + foreground_job: impl FnOnce() -> Option, + read_hint: impl Fn(u32) -> Option + Copy, +) -> ProcessProbeResult { + if let Some(job) = leader_job.as_ref() { + if let Some(hinted) = hinted_process_probe_result(job, pid, read_hint) { + return hinted; + } + if let Some((agent, process_name)) = crate::detect::identify_agent_in_job(job) { + return process_probe_result(job, pid, agent, process_name); } } - if let Some(job) = crate::detect::foreground_job(pid) { - let identified = crate::detect::identify_agent_in_job(&job); + let foreground_job = foreground_job(); + if let Some(job) = foreground_job.as_ref() { + if let Some(agent) = read_hint(job.process_group_id) { + return process_probe_result( + job, + pid, + agent, + crate::detect::agent_label(agent).to_string(), + ); + } + if let Some((agent, process_name)) = identify_process_group_leader_in_job(job) { + return process_probe_result(job, pid, agent, process_name); + } + if let Some(agent) = agent_hint_for_non_leader_foreground_job_members(job, read_hint) { + return process_probe_result( + job, + pid, + agent, + crate::detect::agent_label(agent).to_string(), + ); + } + + let identified = crate::detect::identify_agent_in_job(job); return ProcessProbeResult { process_group_id: Some(job.process_group_id), - foreground_is_pane_shell: job.processes.iter().any(|p| p.pid == pid), + foreground_is_pane_shell: job.processes.iter().any(|process| process.pid == pid), agent: identified.as_ref().map(|(agent, _)| *agent), process_name: identified.map(|(_, process_name)| process_name), }; @@ -442,6 +522,16 @@ fn probe_foreground_process(pid: u32, foreground_pgid: Option) -> ProcessPr } } +fn probe_foreground_process(pid: u32, foreground_pgid: Option) -> ProcessProbeResult { + probe_foreground_process_from_jobs( + pid, + foreground_pgid, + foreground_pgid.and_then(crate::detect::foreground_group_leader_job), + || crate::detect::foreground_job(pid), + crate::platform::process_agent_hint, + ) +} + #[cfg(unix)] fn spawn_basic_detection_task( pane_id: PaneId, @@ -2980,15 +3070,6 @@ mod tests { ); } - #[test] - fn foreground_shell_without_agent_is_immediate_clear_signal() { - assert!(should_clear_agent_for_foreground_shell( - Some(Agent::Claude), - None, - true - )); - } - #[test] fn foreground_shell_reports_process_exit_before_clearing_agent() { assert_eq!( @@ -3003,20 +3084,151 @@ mod tests { #[test] fn unknown_non_shell_foreground_job_is_not_immediate_clear_signal() { - assert!(!should_clear_agent_for_foreground_shell( - Some(Agent::Claude), - None, - false - )); + assert_eq!( + foreground_shell_agent_action(Some(Agent::Claude), None, false, false), + ForegroundShellAgentAction::ObserveProbe + ); + } + + #[test] + fn reported_process_exit_clears_before_unknown_foreground_probe() { + assert_eq!( + foreground_shell_agent_action(Some(Agent::Claude), None, false, true), + ForegroundShellAgentAction::ClearAgent + ); } #[test] fn foreground_agent_job_is_not_clear_signal() { - assert!(!should_clear_agent_for_foreground_shell( - Some(Agent::Claude), - Some(Agent::OpenCode), - true - )); + assert_eq!( + foreground_shell_agent_action(Some(Agent::Claude), Some(Agent::OpenCode), true, false), + ForegroundShellAgentAction::ObserveProbe + ); + } + + fn foreground_process(pid: u32, name: &str) -> crate::platform::ForegroundProcess { + crate::platform::ForegroundProcess { + pid, + name: name.to_string(), + argv0: None, + argv: None, + cmdline: None, + } + } + + #[test] + fn foreground_agent_hint_accepts_pane_shell_environment() { + let job = crate::platform::ForegroundJob { + process_group_id: 42, + processes: vec![foreground_process(42, "bash")], + }; + + assert_eq!( + agent_hint_for_foreground_job_members(&job, |pid| { + (pid == 42).then_some(Agent::Claude) + }), + Some(Agent::Claude) + ); + } + + #[test] + fn foreground_agent_hint_accepts_non_leader_foreground_process_environment() { + let job = crate::platform::ForegroundJob { + process_group_id: 99, + processes: vec![ + foreground_process(99, "fence"), + foreground_process(100, "pi"), + ], + }; + + assert_eq!( + agent_hint_for_foreground_job_members(&job, |pid| { + (pid == 100).then_some(Agent::Codex) + }), + Some(Agent::Codex) + ); + } + + #[test] + fn foreground_agent_hint_wins_over_process_name_detection() { + let job = crate::platform::ForegroundJob { + process_group_id: 99, + processes: vec![foreground_process(99, "codex")], + }; + + let result = probe_foreground_process_from_jobs( + 42, + Some(99), + Some(job), + || None, + |pid| (pid == 99).then_some(Agent::Claude), + ); + + assert_eq!(result.agent, Some(Agent::Claude)); + assert_eq!(result.process_name.as_deref(), Some("claude")); + } + + #[test] + fn foreground_agent_hint_on_inherited_child_environment_is_authoritative() { + let job = crate::platform::ForegroundJob { + process_group_id: 99, + processes: vec![foreground_process(99, "vim")], + }; + + let result = probe_foreground_process_from_jobs( + 42, + Some(99), + None, + || Some(job), + |pid| (pid == 99).then_some(Agent::Claude), + ); + + assert_eq!(result.agent, Some(Agent::Claude)); + assert_eq!(result.process_name.as_deref(), Some("claude")); + } + + #[test] + fn non_leader_agent_hint_does_not_override_identifiable_leader() { + let job = crate::platform::ForegroundJob { + process_group_id: 99, + processes: vec![ + foreground_process(99, "codex"), + foreground_process(100, "vim"), + ], + }; + + let result = probe_foreground_process_from_jobs( + 42, + Some(99), + None, + || Some(job), + |pid| (pid == 100).then_some(Agent::Claude), + ); + + assert_eq!(result.agent, Some(Agent::Codex)); + assert_eq!(result.process_name.as_deref(), Some("codex")); + } + + #[test] + fn non_leader_agent_hint_wins_when_leader_is_unidentified() { + let job = crate::platform::ForegroundJob { + process_group_id: 99, + processes: vec![ + foreground_process(99, "some_vm"), + foreground_process(100, "vim"), + ], + }; + + let result = probe_foreground_process_from_jobs( + 42, + Some(99), + None, + || Some(job), + |pid| (pid == 100).then_some(Agent::Claude), + ); + + assert_eq!(result.agent, Some(Agent::Claude)); + assert_eq!(result.process_name.as_deref(), Some("claude")); } fn process_probe_input() -> ProcessProbeInput { diff --git a/src/platform/linux.rs b/src/platform/linux.rs index 9a44d02e..3cf643e5 100644 --- a/src/platform/linux.rs +++ b/src/platform/linux.rs @@ -124,6 +124,26 @@ pub fn process_cwd(pid: u32) -> Option { std::fs::read_link(format!("/proc/{pid}/cwd")).ok() } +/// Read a Herdr agent identity hint from a process environment. +pub fn process_agent_hint(pid: u32) -> Option { + if pid == 0 { + return None; + } + let environ = std::fs::read(format!("/proc/{pid}/environ")).ok()?; + parse_agent_env_hint(&environ) +} + +fn parse_agent_env_hint(environ: &[u8]) -> Option { + for record in environ.split(|&byte| byte == 0) { + let Some(value) = record.strip_prefix(b"HERDR_AGENT=") else { + continue; + }; + let value = std::str::from_utf8(value).ok()?; + return crate::detect::parse_agent_label(value); + } + None +} + pub fn session_processes(child_pid: u32) -> Vec { let Some(session_id) = process_session_id(child_pid) else { return Vec::new(); @@ -489,6 +509,24 @@ mod tests { LOCK.get_or_init(|| Mutex::new(())) } + #[test] + fn parse_agent_env_hint_accepts_known_agents() { + assert_eq!( + parse_agent_env_hint(b"PATH=/bin\0HERDR_AGENT=claude\0TERM=xterm\0"), + Some(crate::detect::Agent::Claude) + ); + assert_eq!( + parse_agent_env_hint(b"HERDR_AGENT=codex"), + Some(crate::detect::Agent::Codex) + ); + } + + #[test] + fn parse_agent_env_hint_ignores_missing_or_unknown_agents() { + assert_eq!(parse_agent_env_hint(b"PATH=/bin\0TERM=xterm\0"), None); + assert_eq!(parse_agent_env_hint(b"HERDR_AGENT=not-an-agent\0"), None); + } + #[test] fn clipboard_commands_prefer_wayland_when_available() { let _guard = env_lock().lock().unwrap(); diff --git a/src/platform/mod.rs b/src/platform/mod.rs index ec9a497f..0b5f9ecb 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -121,6 +121,11 @@ mod fallback; #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] pub use fallback::*; +#[cfg(not(target_os = "linux"))] +pub fn process_agent_hint(_pid: u32) -> Option { + None +} + #[cfg(not(target_os = "macos"))] #[derive(Debug)] pub(crate) struct InputSourceRestore; -- cgit v1.2.3