From 25aab81aded66d0edaf186e4c53ce69249bf05ae Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:43:15 +0700 Subject: [PATCH 1/2] Port upstream 0.69.0: reap Antigravity CLI usage probe process tree --- rust/src/managed_process.rs | 38 ++++- .../src/providers/antigravity/cli_fallback.rs | 143 +++++++++++++++++- 2 files changed, 170 insertions(+), 11 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index d6a9627914..2471112994 100644 --- a/rust/src/managed_process.rs +++ b/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 { 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 { + 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, diff --git a/rust/src/providers/antigravity/cli_fallback.rs b/rust/src/providers/antigravity/cli_fallback.rs index ed3572358d..ad33474bb8 100644 --- a/rust/src/providers/antigravity/cli_fallback.rs +++ b/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,44 @@ fn prepare_command(binary: &Path, args: &[&str], working_dir: &Path) -> AsyncCom command } +/// Owns the 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)] +struct ProbeContainment( + #[expect(dead_code, reason = "held only so its Drop closes the job")] Option, +); + +#[cfg(not(windows))] +struct ProbeContainment; + +/// 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, ProbeContainment), ProviderError> { + let child = command + .spawn() + .map_err(|_| ProviderError::Other("Failed to start Antigravity CLI".into()))?; + #[cfg(windows)] + let containment = ProbeContainment(contain_child(&child)); + #[cfg(not(windows))] + let containment = ProbeContainment; + Ok((child, containment)) +} + +#[cfg(windows)] +fn contain_child(child: &tokio::process::Child) -> Option { + let handle = child.raw_handle()?; + ProcessJob::create("agy-probe") + .and_then(|job| job.contain(handle).map(|()| job)) + .inspect_err(|error| { + tracing::debug!(%error, "Antigravity CLI probe could not be job-contained"); + }) + .ok() +} + async fn run_cli_command( binary: &Path, args: &[&str], @@ -198,9 +238,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 +351,105 @@ 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)] + 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 mut bystander = 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"); + + wait_until(|| marker.exists(), "descendant pid marker").await; + let mut descendant = None; + wait_until( + || { + descendant = std::fs::read_to_string(&marker) + .ok() + .and_then(|text| text.trim().parse::().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.id())); + + drop(bystander.kill()); + drop(bystander.wait()); + } + #[test] fn usage_fallback_uses_exact_noninteractive_command_and_bounded_timeouts() { assert_eq!(VERSION_ARGS, ["--version"]); From cecf41e588136eff91a37e44fdcc9780b6bec369 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 22:25:03 +0700 Subject: [PATCH 2/2] Address thermo review --- .../src/providers/antigravity/cli_fallback.rs | 75 +++++++++++-------- 1 file changed, 43 insertions(+), 32 deletions(-) diff --git a/rust/src/providers/antigravity/cli_fallback.rs b/rust/src/providers/antigravity/cli_fallback.rs index ad33474bb8..86d02fc099 100644 --- a/rust/src/providers/antigravity/cli_fallback.rs +++ b/rust/src/providers/antigravity/cli_fallback.rs @@ -192,44 +192,46 @@ fn prepare_command(binary: &Path, args: &[&str], working_dir: &Path) -> AsyncCom command } -/// Owns the 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. +/// 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)] -struct ProbeContainment( - #[expect(dead_code, reason = "held only so its Drop closes the job")] Option, -); - +type ProbeJob = ProcessJob; #[cfg(not(windows))] -struct ProbeContainment; +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, ProbeContainment), ProviderError> { +) -> Result<(tokio::process::Child, Option), ProviderError> { let child = command .spawn() .map_err(|_| ProviderError::Other("Failed to start Antigravity CLI".into()))?; - #[cfg(windows)] - let containment = ProbeContainment(contain_child(&child)); - #[cfg(not(windows))] - let containment = ProbeContainment; - Ok((child, containment)) + let job = contain_child(&child); + Ok((child, job)) } #[cfg(windows)] -fn contain_child(child: &tokio::process::Child) -> Option { - let handle = child.raw_handle()?; +fn contain_child(child: &tokio::process::Child) -> Option { + 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::debug!(%error, "Antigravity CLI probe could not be job-contained"); + tracing::warn!(%error, "Antigravity CLI probe could not be job-contained"); }) .ok() } +#[cfg(not(windows))] +fn contain_child(_child: &tokio::process::Child) -> Option { + None +} + async fn run_cli_command( binary: &Path, args: &[&str], @@ -378,6 +380,17 @@ mod tests { } } + #[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); @@ -399,15 +412,17 @@ mod tests { marker.display() ); // An unrelated process started outside the probe must survive the reap. - let mut bystander = std::process::Command::new("powershell.exe") - .args([ - "-NoLogo", - "-NoProfile", - "-Command", - "Start-Sleep -Seconds 120", - ]) - .spawn() - .expect("start bystander"); + 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 @@ -419,7 +434,6 @@ mod tests { .kill_on_drop(true); let (child, containment) = spawn_contained(&mut command).expect("spawn contained probe"); - wait_until(|| marker.exists(), "descendant pid marker").await; let mut descendant = None; wait_until( || { @@ -444,10 +458,7 @@ mod tests { ); drop(containment); wait_until(|| !process_is_alive(descendant), "descendant reaped").await; - assert!(process_is_alive(bystander.id())); - - drop(bystander.kill()); - drop(bystander.wait()); + assert!(process_is_alive(bystander.0.id())); } #[test]