tenant.openagents/omega

1use anyhow::{Context as _, Result};
2use std::process::Stdio;
3
4/// A wrapper around `smol::process::Child` that ensures all subprocesses
5/// are killed when the process is terminated: on Unix by using process
6/// groups, and on Windows by using job objects.
7///
8/// On Windows, dropping this struct closes the job object handle, which
9/// terminates all processes in the job. This also applies when the Zed
10/// process exits for any reason (including crashes), since the OS closes
11/// its handles, so spawned process trees can never outlive Zed.
12pub struct Child {
13    process: smol::process::Child,
14    #[cfg(windows)]
15    job: Option<windows_job::JobObject>,
16}
17
18impl std::ops::Deref for Child {
19    type Target = smol::process::Child;
20
21    fn deref(&self) -> &Self::Target {
22        &self.process
23    }
24}
25
26impl std::ops::DerefMut for Child {
27    fn deref_mut(&mut self) -> &mut Self::Target {
28        &mut self.process
29    }
30}
31
32impl Child {
33    #[cfg(not(windows))]
34    pub fn spawn(
35        mut command: std::process::Command,
36        stdin: Stdio,
37        stdout: Stdio,
38        stderr: Stdio,
39    ) -> Result<Self> {
40        crate::set_pre_exec_to_start_new_session(&mut command);
41        let mut command = smol::process::Command::from(command);
42        let process = command
43            .stdin(stdin)
44            .stdout(stdout)
45            .stderr(stderr)
46            .spawn()
47            .with_context(|| {
48                format!(
49                    "failed to spawn command {}",
50                    crate::redact::redact_command(&format!("{command:?}"))
51                )
52            })?;
53        Ok(Self { process })
54    }
55
56    #[cfg(windows)]
57    pub fn spawn(
58        command: std::process::Command,
59        stdin: Stdio,
60        stdout: Stdio,
61        stderr: Stdio,
62    ) -> Result<Self> {
63        let mut command = smol::process::Command::from(command);
64        let process = command
65            .stdin(stdin)
66            .stdout(stdout)
67            .stderr(stderr)
68            .spawn()
69            .with_context(|| {
70                format!(
71                    "failed to spawn command {}",
72                    crate::redact::redact_command(&format!("{command:?}"))
73                )
74            })?;
75
76        // Assign the child to a job object configured to kill the entire
77        // process tree when the last job handle is closed, so descendants
78        // (e.g. node workers and MCP servers spawned by agent servers) are
79        // reaped even if the direct child doesn't clean them up. Any process
80        // the child spawns after this assignment is automatically part of the
81        // job.
82        //
83        // There is a small race: descendants the child spawns between the
84        // `spawn()` call returning and the assignment below escape the job.
85        // Closing it fully would require creating the process suspended
86        // (`CREATE_SUSPENDED`), assigning it, then resuming it, which the
87        // std/smol process APIs don't support without reimplementing process
88        // creation. The window is microseconds, and the children we care
89        // about (`npx`, `node`, etc.) take far longer to load their runtime
90        // and spawn anything, so in practice nothing escapes.
91        let job = windows_job::JobObject::new()
92            .and_then(|job| {
93                job.assign_process(process.id())?;
94                Ok(job)
95            })
96            .map_err(|error| {
97                log::error!("failed to assign spawned process to a job object: {error:#}");
98            })
99            .ok();
100
101        Ok(Self { process, job })
102    }
103
104    /// Consumes the child, draining its stdout/stderr and waiting for it to
105    /// exit, then returns the collected output.
106    pub async fn output(self) -> Result<std::process::Output> {
107        // NOTE: Keep `self` alive across this await, do not destructure it to
108        // pull `process` out first. On Windows that drops the job object early,
109        // which triggers `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` and kills the
110        // child before `output()` finishes collecting its stdout/stderr.
111        Ok(self.process.output().await?)
112    }
113
114    #[cfg(not(windows))]
115    pub fn kill(&mut self) -> Result<()> {
116        let pid = self.process.id();
117        unsafe {
118            libc::killpg(pid as i32, libc::SIGKILL);
119        }
120        Ok(())
121    }
122
123    #[cfg(windows)]
124    pub fn kill(&mut self) -> Result<()> {
125        if let Some(job) = &self.job {
126            job.terminate()
127        } else {
128            self.process.kill()?;
129            Ok(())
130        }
131    }
132}
133
134#[cfg(windows)]
135mod windows_job {
136    use crate::ResultExt as _;
137    use anyhow::{Context as _, Result};
138    use windows::Win32::{
139        Foundation::{CloseHandle, HANDLE},
140        System::{
141            JobObjects::{
142                AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
143                JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation,
144                SetInformationJobObject, TerminateJobObject,
145            },
146            Threading::{OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE},
147        },
148    };
149
150    /// A Win32 job object configured with `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE`:
151    /// all processes assigned to the job (and their descendants) are terminated
152    /// when the last handle to the job is closed, which happens when this struct
153    /// is dropped, or when the OS closes the owning process's handles after it
154    /// exits for any reason.
155    pub(crate) struct JobObject(HANDLE);
156
157    // SAFETY: Job object handles can be used from any thread.
158    unsafe impl Send for JobObject {}
159    unsafe impl Sync for JobObject {}
160
161    impl JobObject {
162        pub(crate) fn new() -> Result<Self> {
163            unsafe {
164                let job =
165                    Self(CreateJobObjectW(None, None).context("failed to create job object")?);
166                let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default();
167                info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
168                SetInformationJobObject(
169                    job.0,
170                    JobObjectExtendedLimitInformation,
171                    &info as *const _ as *const _,
172                    size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
173                )
174                .context("failed to set job object limits")?;
175                Ok(job)
176            }
177        }
178
179        pub(crate) fn assign_process(&self, pid: u32) -> Result<()> {
180            unsafe {
181                let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, pid)
182                    .context("failed to open process")?;
183                let result = AssignProcessToJobObject(self.0, process)
184                    .context("failed to assign process to job object");
185                CloseHandle(process).log_err();
186                result
187            }
188        }
189
190        pub(crate) fn terminate(&self) -> Result<()> {
191            unsafe { TerminateJobObject(self.0, 1).context("failed to terminate job object") }
192        }
193    }
194
195    impl Drop for JobObject {
196        fn drop(&mut self) {
197            unsafe {
198                CloseHandle(self.0).log_err();
199            }
200        }
201    }
202}
203
204#[cfg(all(test, windows))]
205mod windows_tests {
206    use super::*;
207    use std::time::{Duration, Instant};
208
209    /// Spawns a process tree `powershell -> ping` via `Child::spawn` and
210    /// returns the `Child` along with the pid of the grandchild (`ping`).
211    fn spawn_process_tree(temp_dir: &std::path::Path) -> (Child, u32) {
212        let pid_file = temp_dir.join("grandchild_pid");
213        let mut command = std::process::Command::new("powershell.exe");
214        command.args(["-NoProfile", "-Command"]).arg(format!(
215            "$p = Start-Process -FilePath ping.exe -ArgumentList @('-n','60','127.0.0.1') -PassThru -WindowStyle Hidden; \
216             Set-Content -LiteralPath '{}' -Value $p.Id; \
217             Wait-Process -Id $p.Id",
218            pid_file.display()
219        ));
220        let child = Child::spawn(command, Stdio::null(), Stdio::null(), Stdio::null())
221            .expect("failed to spawn powershell");
222
223        let deadline = Instant::now() + Duration::from_secs(5);
224        let grandchild_pid = loop {
225            if let Ok(contents) = std::fs::read_to_string(&pid_file)
226                && let Ok(pid) = contents.trim().parse::<u32>()
227            {
228                break pid;
229            }
230            assert!(
231                Instant::now() < deadline,
232                "timed out waiting for grandchild pid file"
233            );
234            std::thread::sleep(Duration::from_millis(50));
235        };
236        assert!(
237            process_is_alive(grandchild_pid),
238            "grandchild should be alive after spawning"
239        );
240        (child, grandchild_pid)
241    }
242
243    fn process_is_alive(pid: u32) -> bool {
244        use windows::Win32::{
245            Foundation::{CloseHandle, STILL_ACTIVE},
246            System::Threading::{
247                GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
248            },
249        };
250
251        unsafe {
252            let Ok(handle) = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid) else {
253                return false;
254            };
255            let mut exit_code = 0u32;
256            let alive = GetExitCodeProcess(handle, &mut exit_code).is_ok()
257                && exit_code == STILL_ACTIVE.0 as u32;
258            CloseHandle(handle).expect("failed to close process handle");
259            alive
260        }
261    }
262
263    fn assert_process_exits(pid: u32, message: &str) {
264        let deadline = Instant::now() + Duration::from_secs(2);
265        while process_is_alive(pid) {
266            assert!(Instant::now() < deadline, "{message} (pid {pid})");
267            std::thread::sleep(Duration::from_millis(100));
268        }
269    }
270
271    #[test]
272    fn test_kill_terminates_grandchildren() {
273        let temp_dir = tempfile::tempdir().unwrap();
274        let (mut child, grandchild_pid) = spawn_process_tree(temp_dir.path());
275
276        child.kill().expect("failed to kill child");
277
278        assert_process_exits(
279            grandchild_pid,
280            "grandchild should be terminated after killing the child",
281        );
282    }
283
284    #[test]
285    fn test_drop_terminates_grandchildren() {
286        let temp_dir = tempfile::tempdir().unwrap();
287        let (child, grandchild_pid) = spawn_process_tree(temp_dir.path());
288
289        drop(child);
290
291        assert_process_exits(
292            grandchild_pid,
293            "grandchild should be terminated after dropping the child",
294        );
295    }
296}
297