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..86d02fc099 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,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), 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 { + 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 { + 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::().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"]);