Port upstream 0.69.0: reap Antigravity CLI usage probe process tree

rust/src/managed_process.rs

@@ -19,7 +19,6 @@ use std::ffi::{OsStr, OsString};
 use std::fs::File;
 use std::io::{Read, Write};
 use std::os::windows::ffi::OsStrExt;
-#[cfg(test)]
 use std::os::windows::io::RawHandle;
 use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
 use std::path::{Path, PathBuf};
@@ -32,12 +31,10 @@ use windows::Win32::NetworkManagement::IpHelper::{
 use windows::Win32::System::Console::{
     ClosePseudoConsole, CreatePseudoConsole, HPCON, PSEUDOCONSOLE_INHERIT_CURSOR,
 };
-#[cfg(test)]
-use windows::Win32::System::JobObjects::AssignProcessToJobObject;
 use windows::Win32::System::JobObjects::{
-    CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOBOBJECT_BASIC_LIMIT_INFORMATION,
-    JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation,
-    SetInformationJobObject, TerminateJobObject,
+    AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+    JOBOBJECT_BASIC_LIMIT_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
+    JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject,
 };
 use windows::Win32::System::Pipes::CreatePipe;
 use windows::Win32::System::Threading::{
@@ -557,8 +554,33 @@ fn create_managed_job(label: &str) -> ManagedProcessResult<OwnedHandle> {
     Ok(job)
 }
 
-/// Post-creation assignment retained only for the job-termination unit test.
-#[cfg(test)]
+/// Kill-on-close Job Object for a child the caller spawned itself (for example
+/// a piped, non-PTY probe). Every process in the job, including descendants
+/// started after the child was assigned, is terminated when this value drops.
+/// Only processes explicitly assigned here (and their descendants) are touched.
+pub struct ProcessJob {
+    job: OwnedHandle,
+    label: String,
+}
+
+impl ProcessJob {
+    /// Create an empty kill-on-close job.
+    pub fn create(label: &str) -> ManagedProcessResult<Self> {
+        Ok(Self {
+            job: create_managed_job(label)?,
+            label: label.to_string(),
+        })
+    }
+
+    /// Place a freshly spawned process (and any descendants it starts from now
+    /// on) into the job. Assign immediately after spawning; unlike
+    /// [`ManagedProcess`], which joins the job atomically at creation, this is
+    /// a post-creation assignment.
+    pub fn contain(&self, process: RawHandle) -> ManagedProcessResult<()> {
+        assign_process_to_job(&self.job, process, &self.label)
+    }
+}
+
 fn assign_process_to_job(
     job: &OwnedHandle,
     process: RawHandle,

rust/src/providers/antigravity/cli_fallback.rs

@@ -11,6 +11,8 @@ use uuid::Uuid;
 use super::AntigravityProvider;
 use super::quota_summary;
 use crate::core::{ProviderError, ProviderFetchResult};
+#[cfg(windows)]
+use crate::managed_process::ProcessJob;
 
 const REPORT_TIMEOUT: Duration = Duration::from_secs(90);
 const VERSION_ARGS: [&str; 1] = ["--version"];
@@ -190,6 +192,46 @@ fn prepare_command(binary: &Path, args: &[&str], working_dir: &Path) -> AsyncCom
     command
 }
 
+/// Kill-on-close job for one probe. Dropping it (success, error or timeout)
+/// terminates the probe's whole process tree, not only the direct child that
+/// `kill_on_drop` reaches; `agy` starts MCP server descendants that would
+/// otherwise outlive the probe.
+#[cfg(windows)]
+type ProbeJob = ProcessJob;
+#[cfg(not(windows))]
+type ProbeJob = ();
+
+/// Spawn the probe and place it in its own job. Only this probe's process tree
+/// is ever in the job, so unrelated `agy` processes are never touched.
+fn spawn_contained(
+    command: &mut AsyncCommand,
+) -> Result<(tokio::process::Child, Option<ProbeJob>), ProviderError> {
+    let child = command
+        .spawn()
+        .map_err(|_| ProviderError::Other("Failed to start Antigravity CLI".into()))?;
+    let job = contain_child(&child);
+    Ok((child, job))
+}
+
+#[cfg(windows)]
+fn contain_child(child: &tokio::process::Child) -> Option<ProbeJob> {
+    let Some(handle) = child.raw_handle() else {
+        tracing::warn!("Antigravity CLI probe exited before it could be job-contained");
+        return None;
+    };
+    ProcessJob::create("agy-probe")
+        .and_then(|job| job.contain(handle).map(|()| job))
+        .inspect_err(|error| {
+            tracing::warn!(%error, "Antigravity CLI probe could not be job-contained");
+        })
+        .ok()
+}
+
+#[cfg(not(windows))]
+fn contain_child(_child: &tokio::process::Child) -> Option<ProbeJob> {
+    None
+}
+
 async fn run_cli_command(
     binary: &Path,
     args: &[&str],
@@ -198,9 +240,7 @@ async fn run_cli_command(
     let working_dir = PrivateWorkdir::create()?;
     let mut command = prepare_command(binary, args, working_dir.path());
 
-    let mut child = command
-        .spawn()
-        .map_err(|_| ProviderError::Other("Failed to start Antigravity CLI".into()))?;
+    let (mut child, _containment) = spawn_contained(&mut command)?;
     let stdout = child
         .stdout
         .take()
@@ -313,6 +353,114 @@ mod tests {
         assert!(!path.exists());
     }
 
+    #[cfg(windows)]
+    fn process_is_alive(pid: u32) -> bool {
+        use windows::Win32::Foundation::{CloseHandle, WAIT_OBJECT_0};
+        use windows::Win32::System::Threading::{
+            OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_SYNCHRONIZE,
+            WaitForSingleObject,
+        };
+
+        // SAFETY: OpenProcess returns a handle owned by this function and closed below.
+        match unsafe {
+            OpenProcess(
+                PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_SYNCHRONIZE,
+                false,
+                pid,
+            )
+        } {
+            Ok(handle) => {
+                // SAFETY: `handle` is a valid process handle.
+                let exited = unsafe { WaitForSingleObject(handle, 0) } == WAIT_OBJECT_0;
+                // SAFETY: closing the handle returned by OpenProcess exactly once.
+                drop(unsafe { CloseHandle(handle) });
+                !exited
+            }
+            Err(_) => false,
+        }
+    }
+
+    #[cfg(windows)]
+    struct KillOnDrop(std::process::Child);
+
+    #[cfg(windows)]
+    impl Drop for KillOnDrop {
+        fn drop(&mut self) {
+            drop(self.0.kill());
+            drop(self.0.wait());
+        }
+    }
+
+    #[cfg(windows)]
+    async fn wait_until(mut condition: impl FnMut() -> bool, what: &str) {
+        let deadline = std::time::Instant::now() + Duration::from_secs(15);
+        while !condition() {
+            assert!(std::time::Instant::now() < deadline, "timed out: {what}");
+            tokio::time::sleep(Duration::from_millis(50)).await;
+        }
+    }
+
+    #[cfg(windows)]
+    #[tokio::test]
+    async fn dropping_the_probe_reaps_descendants_but_not_unrelated_processes() {
+        let workdir = PrivateWorkdir::create().expect("private working directory");
+        let marker = workdir.path().join("descendant.pid");
+        let script = format!(
+            "$p = Start-Process -PassThru -WindowStyle Hidden powershell.exe \
+             -ArgumentList '-NoLogo','-NoProfile','-Command','Start-Sleep -Seconds 120'; \
+             Set-Content -LiteralPath '{}' -Value $p.Id; Start-Sleep -Seconds 120",
+            marker.display()
+        );
+        // An unrelated process started outside the probe must survive the reap.
+        let bystander = KillOnDrop(
+            std::process::Command::new("powershell.exe")
+                .args([
+                    "-NoLogo",
+                    "-NoProfile",
+                    "-Command",
+                    "Start-Sleep -Seconds 120",
+                ])
+                .spawn()
+                .expect("start bystander"),
+        );
+
+        let mut command = AsyncCommand::new("powershell.exe");
+        command
+            .args(["-NoLogo", "-NoProfile", "-Command", &script])
+            .current_dir(workdir.path())
+            .stdin(Stdio::null())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .kill_on_drop(true);
+        let (child, containment) = spawn_contained(&mut command).expect("spawn contained probe");
+
+        let mut descendant = None;
+        wait_until(
+            || {
+                descendant = std::fs::read_to_string(&marker)
+                    .ok()
+                    .and_then(|text| text.trim().parse::<u32>().ok());
+                descendant.is_some()
+            },
+            "descendant pid",
+        )
+        .await;
+        let descendant = descendant.expect("descendant pid");
+        assert!(process_is_alive(descendant));
+
+        // kill_on_drop only kills the direct child: the descendant survives it
+        // and is reaped only when the job closes.
+        drop(child);
+        tokio::time::sleep(Duration::from_millis(500)).await;
+        assert!(
+            process_is_alive(descendant),
+            "kill_on_drop alone must not be what reaps the descendant"
+        );
+        drop(containment);
+        wait_until(|| !process_is_alive(descendant), "descendant reaped").await;
+        assert!(process_is_alive(bystander.0.id()));
+    }
+
     #[test]
     fn usage_fallback_uses_exact_noninteractive_command_and_bounded_timeouts() {
         assert_eq!(VERSION_ARGS, ["--version"]);