tenant.openagents/omega

1use gpui::{Context, Task};
2use parking_lot::{MappedRwLockReadGuard, Mutex, RwLock, RwLockReadGuard};
3use std::{path::PathBuf, sync::Arc};
4
5#[cfg(target_os = "windows")]
6use windows::Win32::{Foundation::HANDLE, System::Threading::GetProcessId};
7
8use sysinfo::{Pid, Process, ProcessRefreshKind, ProcessesToUpdate, System, UpdateKind};
9
10use crate::{Event, Terminal};
11
12#[derive(Clone, Copy)]
13pub struct ProcessIdGetter {
14    handle: i32,
15    fallback_pid: u32,
16}
17
18impl ProcessIdGetter {
19    pub(crate) fn new(handle: i32, fallback_pid: u32) -> ProcessIdGetter {
20        ProcessIdGetter {
21            handle,
22            fallback_pid,
23        }
24    }
25
26    pub fn fallback_pid(&self) -> Pid {
27        Pid::from_u32(self.fallback_pid)
28    }
29}
30
31#[cfg(unix)]
32impl ProcessIdGetter {
33    fn pid(&self) -> Option<Pid> {
34        // Negative pid means error.
35        // Zero pid means no foreground process group is set on the PTY yet.
36        // Avoid killing the current process by returning a zero pid.
37        let pid = unsafe { libc::tcgetpgrp(self.handle) };
38        if pid > 0 {
39            return Some(Pid::from_u32(pid as u32));
40        }
41
42        if self.fallback_pid > 0 {
43            return Some(Pid::from_u32(self.fallback_pid));
44        }
45
46        None
47    }
48}
49
50#[cfg(windows)]
51impl ProcessIdGetter {
52    fn pid(&self) -> Option<Pid> {
53        let pid = unsafe { GetProcessId(HANDLE(self.handle as _)) };
54        // the GetProcessId may fail and returns zero, which will lead to a stack overflow issue
55        if pid == 0 {
56            // in the builder process, there is a small chance, almost negligible,
57            // that this value could be zero, which means child_watcher returns None,
58            // GetProcessId returns 0.
59            if self.fallback_pid == 0 {
60                return None;
61            }
62            return Some(Pid::from_u32(self.fallback_pid));
63        }
64        Some(Pid::from_u32(pid))
65    }
66}
67
68#[derive(Clone, Debug)]
69pub(crate) struct ProcessInfo {
70    pub(crate) name: String,
71    pub(crate) cwd: PathBuf,
72    pub(crate) argv: Vec<String>,
73}
74
75/// Process (group) ids of a terminal's shell and foreground job, snapshotted
76/// while the PTY master is still open: reading the foreground process group
77/// requires `tcgetpgrp` on the PTY fd, which the event loop closes when the
78/// terminal shuts down, so these ids must be captured before shutdown and
79/// signalled afterwards.
80#[derive(Clone, Copy)]
81pub(crate) struct TerminalProcessIds {
82    #[cfg_attr(not(unix), allow(dead_code))]
83    foreground: Option<Pid>,
84    #[cfg_attr(not(unix), allow(dead_code))]
85    child: Pid,
86}
87
88#[cfg(unix)]
89impl TerminalProcessIds {
90    /// The spawned child (the shell) leads its own process group, but under
91    /// job control a foreground job runs in a separate process group that
92    /// `killpg` on the shell's group never reaches, so both are signalled
93    /// (see #47412).
94    fn process_group_ids(self) -> impl Iterator<Item = i32> {
95        std::iter::once(self.child)
96            .chain(
97                self.foreground
98                    .filter(|foreground| *foreground != self.child),
99            )
100            .map(|pid| pid.as_u32() as i32)
101            // `killpg(0, ...)` signals the caller's own process group, i.e.
102            // Zed itself, so never let a zero id (or a negative one from an
103            // implausibly large pid wrapping the cast) through.
104            .filter(|process_group_id| *process_group_id > 0)
105    }
106
107    /// Returns whether at least one process group was signalled successfully;
108    /// `killpg` failing with `ESRCH` (the group already exited) is expected and
109    /// reported as an unsuccessful signal.
110    fn signal_process_groups(&self, signal: i32) -> bool {
111        let mut signalled = false;
112        for process_group_id in self.process_group_ids() {
113            signalled |= unsafe { libc::killpg(process_group_id, signal) } == 0;
114        }
115        signalled
116    }
117
118    pub(crate) fn terminate(&self) -> bool {
119        self.signal_process_groups(libc::SIGTERM)
120    }
121
122    pub(crate) fn kill(&self) -> bool {
123        self.signal_process_groups(libc::SIGKILL)
124    }
125}
126
127#[cfg(not(unix))]
128impl TerminalProcessIds {
129    pub(crate) fn terminate(&self) -> bool {
130        false
131    }
132
133    // Windows has no process groups to escalate on; killing the child relies
134    // on [`PtyProcessInfo::kill_child_process`] instead.
135    pub(crate) fn kill(&self) -> bool {
136        false
137    }
138}
139
140/// Fetches Zed-relevant Pseudo-Terminal (PTY) process information
141pub(crate) struct PtyProcessInfo {
142    system: RwLock<System>,
143    refresh_kind: ProcessRefreshKind,
144    pid_getter: ProcessIdGetter,
145    last_foreground_pid: Mutex<Option<Pid>>,
146    pub(crate) current: RwLock<Option<ProcessInfo>>,
147    task: Mutex<Option<Task<()>>>,
148}
149
150impl PtyProcessInfo {
151    pub(crate) fn new(pid_getter: ProcessIdGetter) -> PtyProcessInfo {
152        // Task enumeration is on by default and would retain a `Process` entry
153        // per thread, each pinning an open `/proc/<pid>/task/<tid>/stat` handle
154        // on Linux (#58651).
155        let process_refresh_kind = ProcessRefreshKind::nothing()
156            .with_cmd(UpdateKind::Always)
157            .with_cwd(UpdateKind::Always)
158            .with_exe(UpdateKind::Always)
159            .without_tasks();
160        // `System::new_with_specifics` with a process refresh kind would
161        // snapshot every process on the machine into this terminal's `System`,
162        // retaining one open procfs handle per process for the lifetime of the
163        // terminal (#58651). Refresh only the spawned child so that
164        // `kill_child_process` works before the first foreground refresh.
165        let mut system = System::new();
166        system.refresh_processes_specifics(
167            ProcessesToUpdate::Some(&[pid_getter.fallback_pid()]),
168            true,
169            process_refresh_kind,
170        );
171
172        PtyProcessInfo {
173            system: RwLock::new(system),
174            refresh_kind: process_refresh_kind,
175            pid_getter,
176            last_foreground_pid: Mutex::new(None),
177            current: RwLock::new(None),
178            task: Mutex::new(None),
179        }
180    }
181
182    pub(crate) fn pid_getter(&self) -> &ProcessIdGetter {
183        &self.pid_getter
184    }
185
186    pub(crate) fn capture_process_ids(&self) -> TerminalProcessIds {
187        TerminalProcessIds {
188            foreground: self.pid_getter.pid(),
189            child: self.pid_getter.fallback_pid(),
190        }
191    }
192
193    fn refresh(&self) -> Option<MappedRwLockReadGuard<'_, Process>> {
194        let pid = self.pid_getter.pid()?;
195        let fallback_pid = self.pid_getter.fallback_pid();
196        let mut system = self.system.write();
197        // sysinfo never evicts processes that are absent from the refreshed pid
198        // set, so entries for former foreground processes (each pinning an open
199        // `/proc/<pid>/stat` handle on Linux) would otherwise accumulate for as
200        // long as this terminal lives (#58651). Rebuild the `System` whenever
201        // the foreground process changes to keep the map bounded.
202        if self.last_foreground_pid.lock().replace(pid) != Some(pid) {
203            *system = System::new();
204        }
205        let pids = [pid, fallback_pid];
206        let pids = if pid == fallback_pid {
207            &pids[..1]
208        } else {
209            &pids[..]
210        };
211        system.refresh_processes_specifics(ProcessesToUpdate::Some(pids), true, self.refresh_kind);
212        drop(system);
213        RwLockReadGuard::try_map(self.system.read(), |system| system.process(pid)).ok()
214    }
215
216    fn get_child(&self) -> Option<MappedRwLockReadGuard<'_, Process>> {
217        let pid = self.pid_getter.fallback_pid();
218        RwLockReadGuard::try_map(self.system.read(), |system| system.process(pid)).ok()
219    }
220
221    #[cfg(unix)]
222    pub(crate) fn kill_current_process(&self) -> bool {
223        let Some(pid) = self.pid_getter.pid() else {
224            return false;
225        };
226        unsafe { libc::killpg(pid.as_u32() as i32, libc::SIGKILL) == 0 }
227    }
228
229    #[cfg(not(unix))]
230    pub(crate) fn kill_current_process(&self) -> bool {
231        self.refresh().is_some_and(|process| process.kill())
232    }
233
234    pub(crate) fn kill_child_process(&self) -> bool {
235        self.get_child().is_some_and(|process| process.kill())
236    }
237
238    fn load(&self) -> Option<ProcessInfo> {
239        let process = self.refresh()?;
240        let cwd = process.cwd().map_or(PathBuf::new(), |p| p.to_owned());
241
242        let info = ProcessInfo {
243            name: process.name().to_str()?.to_owned(),
244            cwd,
245            argv: process
246                .cmd()
247                .iter()
248                .filter_map(|s| s.to_str().map(ToOwned::to_owned))
249                .collect(),
250        };
251        *self.current.write() = Some(info.clone());
252        Some(info)
253    }
254
255    #[cfg(all(test, unix))]
256    pub(crate) fn load_for_test(&self) -> Option<ProcessInfo> {
257        self.load()
258    }
259
260    /// Updates the cached process info, emitting a [`Event::TitleChanged`] event if the Zed-relevant info has changed
261    pub(crate) fn emit_title_changed_if_changed(self: &Arc<Self>, cx: &mut Context<'_, Terminal>) {
262        if self.task.lock().is_some() {
263            return;
264        }
265        let this = self.clone();
266        let has_changed = cx.background_executor().spawn(async move {
267            let previous = this.current.read().clone();
268            let current = this.load();
269            let has_changed = match (previous.as_ref(), current.as_ref()) {
270                (None, None) => false,
271                (Some(prev), Some(now)) => prev.cwd != now.cwd || prev.name != now.name,
272                _ => true,
273            };
274            if has_changed {
275                *this.current.write() = current;
276            }
277            has_changed
278        });
279        let this = Arc::downgrade(self);
280        *self.task.lock() = Some(cx.spawn(async move |term, cx| {
281            if has_changed.await {
282                term.update(cx, |_, cx| cx.emit(Event::TitleChanged)).ok();
283            }
284            if let Some(this) = this.upgrade() {
285                this.task.lock().take();
286            }
287        }));
288    }
289
290    pub(crate) fn pid(&self) -> Option<Pid> {
291        self.pid_getter.pid()
292    }
293}
294
295#[cfg(all(test, unix))]
296mod tests {
297    use super::*;
298
299    /// Regression test for <https://github.com/zed-industries/zed/issues/58651>:
300    /// on Linux, sysinfo keeps an open `/proc/<pid>/stat` handle for every
301    /// `Process` entry retained in a `System`, and never evicts entries that are
302    /// absent from the refreshed pid set. The per-terminal `System` must
303    /// therefore not snapshot every process on the machine, nor accumulate an
304    /// entry per foreground process that has ever run in this terminal.
305    #[test]
306    #[allow(
307        clippy::disallowed_methods,
308        reason = "the test needs real short-lived child processes and may block"
309    )]
310    fn process_map_stays_bounded() {
311        let mut info = PtyProcessInfo::new(ProcessIdGetter::new(-1, std::process::id()));
312        assert!(
313            info.get_child().is_some(),
314            "the spawned child must be inspectable for kill_child_process \
315             before the first foreground refresh"
316        );
317        assert!(info.load_for_test().is_some());
318        let initial_len = info.system.read().processes().len();
319        assert!(
320            initial_len <= 2,
321            "creating a terminal retained {initial_len} process entries"
322        );
323
324        for _ in 0..3 {
325            let mut child = std::process::Command::new("sleep")
326                .arg("30")
327                .spawn()
328                .expect("failed to spawn child process");
329            info.pid_getter = ProcessIdGetter::new(-1, child.id());
330            assert!(info.load_for_test().is_some());
331            child.kill().expect("failed to kill child process");
332            child.wait().expect("failed to wait for child process");
333        }
334
335        let churned_len = info.system.read().processes().len();
336        assert!(
337            churned_len <= 2,
338            "foreground process churn retained {churned_len} process entries"
339        );
340    }
341}
342