From 2425db3ec09f11df956527ec17d63b6daef67d57 Mon Sep 17 00:00:00 2001 From: iiiMohammed <311510718+iiiMohammed@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:57:13 +0300 Subject: [PATCH 01/10] fix(antigravity): launch agy for offline refresh --- rust/src/providers/antigravity/mod.rs | 277 +++++++++++++++++++++++- rust/src/providers/antigravity/tests.rs | 53 +++++ 2 files changed, 321 insertions(+), 9 deletions(-) diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 43592b6f16..dad3819628 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -12,10 +12,13 @@ mod quota_summary; use async_trait::async_trait; use regex_lite::Regex; use serde::Deserialize; +use std::io::{Read, Write}; #[cfg(windows)] use std::os::windows::process::CommandExt; +use std::path::PathBuf; use std::process::Command; use std::sync::{LazyLock, OnceLock}; +use std::time::{Duration, Instant}; use crate::core::{ FetchContext, NamedRateWindow, Provider, ProviderError, ProviderFetchResult, ProviderId, @@ -24,10 +27,18 @@ use crate::core::{ const NOT_RUNNING_MESSAGE: &str = "Antigravity language server not running. Start Google Antigravity and sign in, then retry."; +const AGY_NOT_FOUND_MESSAGE: &str = + "Antigravity is not running and the signed-in agy CLI was not found."; +const AGY_READY_TIMEOUT: Duration = Duration::from_secs(20); +const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; +/// Serialize task-owned `agy` launches so concurrent app surfaces never start +/// multiple interactive CLI servers at the same time. +static MANAGED_AGY_FETCH: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + /// Antigravity provider pub struct AntigravityProvider { metadata: ProviderMetadata, @@ -71,6 +82,15 @@ fn is_agy_cli_command(command_line: &str) -> bool { CLI_PATH_RE.is_match(&lower) || AGY_RE.is_match(&lower) } +fn terminal_requested_cursor_position(tail: &mut Vec, chunk: &[u8]) -> bool { + tail.extend_from_slice(chunk); + let requested = tail.windows(4).any(|bytes| bytes == b"\x1b[6n"); + if tail.len() > 3 { + tail.drain(..tail.len() - 3); + } + requested +} + impl AntigravityProvider { pub fn new() -> Self { Self { @@ -329,7 +349,15 @@ impl AntigravityProvider { async fn fetch_user_status(&self) -> Result { let process_info = Self::detect_process_info()?; let api_port = Self::find_api_port(process_info.extension_port, process_info.pid).await?; + self.fetch_user_status_at_port(&process_info, api_port) + .await + } + async fn fetch_user_status_at_port( + &self, + process_info: &ProcessInfo, + api_port: u16, + ) -> Result { // SECURITY: TLS verification disabled only for this loopback language server. let client = crate::core::credentialed_http_client_builder() .no_proxy() @@ -342,7 +370,7 @@ impl AntigravityProvider { let quota_body = serde_json::json!({ "forceRefresh": true }); match Self::fetch_local_payload( &client, - &process_info, + process_info, api_port, QUOTA_SUMMARY_PATH, "a_body, @@ -364,7 +392,7 @@ impl AntigravityProvider { }); if let Ok(identity_bytes) = Self::fetch_local_payload( &client, - &process_info, + process_info, api_port, GET_USER_STATUS_PATH, &identity_body, @@ -399,7 +427,7 @@ impl AntigravityProvider { }); let bytes = Self::fetch_local_payload( &client, - &process_info, + process_info, api_port, GET_USER_STATUS_PATH, &body, @@ -411,6 +439,123 @@ impl AntigravityProvider { self.parse_user_status(response) } + /// Start a short-lived, headless `agy` session when neither the Antigravity + /// desktop app nor a user-owned CLI session is running. The child is kept + /// alive only for this fetch and is always reaped by `ManagedAgyProcess`. + async fn fetch_with_managed_agy(&self) -> Result { + let _launch_guard = MANAGED_AGY_FETCH.lock().await; + + let binary = Self::locate_agy_binary() + .ok_or_else(|| ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.to_string()))?; + let mut managed = ManagedAgyProcess::spawn(&binary)?; + let pid = managed.id(); + let process_info = ProcessInfo { + csrf_token: String::new(), + extension_server_csrf_token: None, + extension_port: None, + pid: Some(pid), + source: ProcessSource::Cli, + }; + let deadline = Instant::now() + AGY_READY_TIMEOUT; + let probe_client = crate::core::credentialed_http_client_builder() + .no_proxy() + .timeout(Duration::from_secs(2)) + .danger_accept_invalid_certs(true) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|error| ProviderError::Other(error.to_string()))?; + let mut last_error = None; + + loop { + if let Some(status) = managed.try_wait()? { + return Err(ProviderError::NotInstalled(format!( + "agy exited before its local quota service was ready ({status}). Open Antigravity or run agy and sign in, then retry." + ))); + } + + let ports = Self::listening_ports_for_pid(pid); + for port in ports { + if !Self::probe_api_port(&probe_client, port).await { + continue; + } + + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match tokio::time::timeout( + remaining, + self.fetch_user_status_at_port(&process_info, port), + ) + .await + { + Ok(Ok(usage)) => return Ok(usage), + Ok(Err(ProviderError::AuthRequired)) => { + return Err(ProviderError::AuthRequired); + } + Ok(Err(error)) => last_error = Some(error), + Err(_) => break, + } + } + + if Instant::now() >= deadline { + break; + } + tokio::time::sleep( + AGY_READY_POLL_INTERVAL.min(deadline.saturating_duration_since(Instant::now())), + ) + .await; + } + + if let Some(error) = last_error { + tracing::debug!(%error, "managed agy quota service did not become ready"); + } + Err(ProviderError::Other( + "agy started but its quota service did not become ready. Open Antigravity or run agy and sign in, then retry." + .to_string(), + )) + } + + fn is_not_running_error(error: &ProviderError) -> bool { + matches!(error, ProviderError::NotInstalled(message) if message == NOT_RUNNING_MESSAGE) + } + + fn locate_agy_binary() -> Option { + let candidates = Self::agy_binary_candidates( + std::env::var_os("ANTIGRAVITY_CLI_PATH").map(PathBuf::from), + which::which("agy").ok(), + std::env::var_os("LOCALAPPDATA").map(PathBuf::from), + dirs::home_dir(), + ); + candidates.into_iter().find(|path| path.is_file()) + } + + fn agy_binary_candidates( + explicit: Option, + path_lookup: Option, + local_app_data: Option, + home: Option, + ) -> Vec { + let mut candidates = Vec::new(); + if let Some(path) = explicit { + candidates.push(path); + } + if let Some(path) = path_lookup { + candidates.push(path); + } + if let Some(root) = local_app_data { + candidates.push(root.join("agy").join("bin").join("agy.exe")); + } + if let Some(root) = home { + candidates.push(root.join(".local").join("bin").join(if cfg!(windows) { + "agy.exe" + } else { + "agy" + })); + } + candidates + } + async fn fetch_local_payload( client: &reqwest::Client, process_info: &ProcessInfo, @@ -625,9 +770,8 @@ impl Provider for AntigravityProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { // `oauth` is not supported (no remote API path is ported yet); surface it // explicitly instead of silently probing locally. Both `auto` and `cli` - // resolve to the same local language-server probe: `detect_process_info` - // prefers the CSRF-protected desktop IDE/app server and falls back to the - // tokenless `agy` CLI when only that is running. + // prefer an existing desktop/CLI language server. When neither is + // running, start a task-owned `agy` session for this fetch only. if ctx.source_mode == SourceMode::OAuth { return Err(ProviderError::UnsupportedSource(ctx.source_mode)); } @@ -639,7 +783,22 @@ impl Provider for AntigravityProvider { Self::with_cadence_labels(usage), "local", )), - Err(e) => { + Err(error) if Self::is_not_running_error(&error) => { + match self.fetch_with_managed_agy().await { + Ok(usage) => { + return Ok(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "cli", + )); + } + Err(ProviderError::NotInstalled(message)) + if message == AGY_NOT_FOUND_MESSAGE => {} + Err(error) => { + tracing::warn!(%error, "managed Antigravity CLI probe failed"); + return Err(error); + } + } + let count = local_sessions::offline_conversation_count(); if count > 0 { let noun = if count == 1 { @@ -653,8 +812,13 @@ impl Provider for AntigravityProvider { .with_login_method("offline"); return Ok(ProviderFetchResult::new(usage, "offline")); } - tracing::warn!("Antigravity probe failed: {}", e); - Err(e) + Err(ProviderError::NotInstalled( + AGY_NOT_FOUND_MESSAGE.to_string(), + )) + } + Err(error) => { + tracing::warn!(%error, "Antigravity local probe failed"); + Err(error) } } } @@ -693,6 +857,101 @@ struct ProcessInfo { source: ProcessSource, } +/// RAII owner for the exact `agy` process started by this provider. Dropping it +/// cannot affect Antigravity or CLI processes that were already running. +struct ManagedAgyProcess { + child: Box, + pid: u32, + master: Option>, + drain_thread: Option>, +} + +impl ManagedAgyProcess { + fn spawn(binary: &std::path::Path) -> Result { + let pty_system = portable_pty::native_pty_system(); + let pair = pty_system + .openpty(portable_pty::PtySize { + rows: 30, + cols: 120, + pixel_width: 0, + pixel_height: 0, + }) + .map_err(|error| { + ProviderError::Other(format!("Failed to create a terminal for agy: {error}")) + })?; + let mut command = portable_pty::CommandBuilder::new(binary.as_os_str()); + if let Some(home) = dirs::home_dir().filter(|path| path.is_dir()) { + command.cwd(home.as_os_str()); + } + command.env("TERM", "xterm-256color"); + command.env("COLORTERM", "truecolor"); + + let mut reader = pair.master.try_clone_reader().map_err(|error| { + ProviderError::Other(format!("Failed to read the agy terminal: {error}")) + })?; + let mut writer = pair.master.take_writer().map_err(|error| { + ProviderError::Other(format!("Failed to open the agy terminal: {error}")) + })?; + let mut child = pair.slave.spawn_command(command).map_err(|error| { + ProviderError::Other(format!("Failed to launch the agy CLI: {error}")) + })?; + drop(pair.slave); + + let Some(pid) = child.process_id() else { + drop(child.kill()); + drop(child.wait()); + return Err(ProviderError::Other( + "Failed to determine the managed agy process id".to_string(), + )); + }; + let drain_thread = std::thread::spawn(move || { + // Drain without logging: terminal output can contain account data. + // Windows ConPTY programs may request the cursor position and wait + // for a terminal response before continuing initialization. + let mut buffer = [0_u8; 4096]; + let mut tail = Vec::with_capacity(3); + while let Ok(read) = reader.read(&mut buffer) { + if read == 0 { + break; + } + if terminal_requested_cursor_position(&mut tail, &buffer[..read]) { + drop(writer.write_all(b"\x1b[1;1R")); + drop(writer.flush()); + } + } + }); + Ok(Self { + child, + pid, + master: Some(pair.master), + drain_thread: Some(drain_thread), + }) + } + + fn id(&self) -> u32 { + self.pid + } + + fn try_wait(&mut self) -> Result, ProviderError> { + self.child.try_wait().map_err(|error| { + ProviderError::Other(format!("Failed to inspect the agy CLI: {error}")) + }) + } +} + +impl Drop for ManagedAgyProcess { + fn drop(&mut self) { + if self.child.try_wait().ok().flatten().is_none() { + drop(self.child.kill()); + } + drop(self.child.wait()); + drop(self.master.take()); + if let Some(thread) = self.drain_thread.take() { + drop(thread.join()); + } + } +} + // API Response types #[derive(Debug, Deserialize)] diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index 9e6545bf24..bf2383ce72 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -225,6 +225,59 @@ fn not_running_error_tells_user_how_to_start() { assert!(error.contains("Start Google Antigravity and sign in")); } +#[test] +fn managed_agy_candidates_prefer_override_then_path_then_known_installs() { + let candidates = AntigravityProvider::agy_binary_candidates( + Some(PathBuf::from(r"D:\tools\agy.exe")), + Some(PathBuf::from(r"C:\path\agy.exe")), + Some(PathBuf::from(r"C:\Users\test\AppData\Local")), + Some(PathBuf::from(r"C:\Users\test")), + ); + + assert_eq!(candidates[0], PathBuf::from(r"D:\tools\agy.exe")); + assert_eq!(candidates[1], PathBuf::from(r"C:\path\agy.exe")); + assert_eq!( + candidates[2], + PathBuf::from(r"C:\Users\test\AppData\Local\agy\bin\agy.exe") + ); + assert_eq!( + candidates[3], + PathBuf::from(r"C:\Users\test\.local\bin").join(if cfg!(windows) { + "agy.exe" + } else { + "agy" + }) + ); +} + +#[test] +fn managed_agy_not_running_check_is_exact() { + assert!(AntigravityProvider::is_not_running_error( + &ProviderError::NotInstalled(NOT_RUNNING_MESSAGE.to_string()) + )); + assert!(!AntigravityProvider::is_not_running_error( + &ProviderError::NotInstalled("Failed to detect Antigravity process".to_string()) + )); + assert!(!AntigravityProvider::is_not_running_error( + &ProviderError::AuthRequired + )); +} + +#[test] +fn managed_agy_terminal_detects_cursor_request_across_reads() { + let mut tail = Vec::new(); + + assert!(!terminal_requested_cursor_position( + &mut tail, + b"ready\x1b[" + )); + assert!(terminal_requested_cursor_position(&mut tail, b"6n")); + assert!(!terminal_requested_cursor_position( + &mut tail, + b"plain output" + )); +} + // ── agy CLI process matching ─────────────────────────────────────── #[test] From 520e5f77e67e6d99f2f1884a8c3dc91914ee044f Mon Sep 17 00:00:00 2001 From: iiiMohammed <311510718+iiiMohammed@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:47:25 +0300 Subject: [PATCH 02/10] Harden managed Antigravity refresh --- rust/Cargo.toml | 1 + rust/src/providers/antigravity/mod.rs | 564 ++++++++++++++++++------ rust/src/providers/antigravity/tests.rs | 108 +++-- 3 files changed, 522 insertions(+), 151 deletions(-) diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7742a5f73d..c0b643ccd6 100755 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -89,6 +89,7 @@ windows = { version = "0.58", features = [ "Win32_Security_Credentials", "Win32_Security_Cryptography", "Win32_Media_Audio", + "Win32_NetworkManagement_IpHelper", "Win32_System_LibraryLoader", "Win32_System_JobObjects", "Win32_System_Threading", diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index dad3819628..0fcc491557 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -10,33 +10,61 @@ mod local_step_resolver; mod quota_summary; use async_trait::async_trait; +#[cfg(windows)] +use futures::{StreamExt, stream}; use regex_lite::Regex; use serde::Deserialize; +#[cfg(windows)] use std::io::{Read, Write}; #[cfg(windows)] +use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; +#[cfg(windows)] use std::os::windows::process::CommandExt; use std::path::PathBuf; use std::process::Command; use std::sync::{LazyLock, OnceLock}; -use std::time::{Duration, Instant}; +use std::time::Duration; +#[cfg(windows)] +use std::time::Instant; +#[cfg(windows)] +use windows::Win32::Foundation::{ERROR_INSUFFICIENT_BUFFER, HANDLE}; +#[cfg(windows)] +use windows::Win32::NetworkManagement::IpHelper::{ + GetExtendedTcpTable, MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID, TCP_TABLE_OWNER_PID_LISTENER, +}; +#[cfg(windows)] +use windows::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_BASIC_LIMIT_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, + JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject, +}; +#[cfg(windows)] +use windows::core::PCWSTR; use crate::core::{ FetchContext, NamedRateWindow, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; -const NOT_RUNNING_MESSAGE: &str = - "Antigravity language server not running. Start Google Antigravity and sign in, then retry."; const AGY_NOT_FOUND_MESSAGE: &str = "Antigravity is not running and the signed-in agy CLI was not found."; -const AGY_READY_TIMEOUT: Duration = Duration::from_secs(20); +#[cfg(windows)] +const AGY_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(25); +#[cfg(windows)] +const AGY_CLEANUP_RESERVE: Duration = Duration::from_secs(2); +#[cfg(windows)] +const AGY_PROBE_TIMEOUT: Duration = Duration::from_millis(750); +#[cfg(windows)] const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); +#[cfg(any(windows, test))] +const AGY_MAX_CURSOR_REPLIES: usize = 32; const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; /// Serialize task-owned `agy` launches so concurrent app surfaces never start /// multiple interactive CLI servers at the same time. +#[cfg(windows)] static MANAGED_AGY_FETCH: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); /// Antigravity provider @@ -82,15 +110,21 @@ fn is_agy_cli_command(command_line: &str) -> bool { CLI_PATH_RE.is_match(&lower) || AGY_RE.is_match(&lower) } -fn terminal_requested_cursor_position(tail: &mut Vec, chunk: &[u8]) -> bool { +#[cfg(any(windows, test))] +fn terminal_cursor_position_request_count(tail: &mut Vec, chunk: &[u8]) -> usize { tail.extend_from_slice(chunk); - let requested = tail.windows(4).any(|bytes| bytes == b"\x1b[6n"); + let requested = tail.windows(4).filter(|bytes| *bytes == b"\x1b[6n").count(); if tail.len() > 3 { tail.drain(..tail.len() - 3); } requested } +#[cfg(any(windows, test))] +fn terminal_cursor_reply_allowance(sent: usize, requested: usize) -> usize { + requested.min(AGY_MAX_CURSOR_REPLIES.saturating_sub(sent)) +} + impl AntigravityProvider { pub fn new() -> Self { Self { @@ -110,13 +144,16 @@ impl AntigravityProvider { } /// Detect running Antigravity language server and extract connection info - fn detect_process_info() -> Result { + fn detect_process_info() -> Result, ProviderError> { // Use PowerShell to get process command lines #[cfg(windows)] const CREATE_NO_WINDOW: u32 = 0x08000000; let mut cmd = Command::new("powershell.exe"); cmd.args([ + "-NoLogo", + "-NoProfile", + "-NonInteractive", "-ExecutionPolicy", "Bypass", "-Command", // Match the desktop IDE/app language server (language_server.exe / @@ -138,8 +175,7 @@ impl AntigravityProvider { } let stdout = String::from_utf8_lossy(&output.stdout); - Self::parse_process_info(&stdout) - .ok_or_else(|| ProviderError::NotInstalled(NOT_RUNNING_MESSAGE.to_string())) + Ok(Self::parse_process_info(&stdout)) } fn parse_process_info(stdout: &str) -> Option { @@ -240,8 +276,10 @@ impl AntigravityProvider { // equivalent of `lsof`), then the heuristic window above the extension port, then a // few known ports as a last resort. let mut candidates: Vec = Vec::new(); - if let Some(pid) = pid { - candidates.extend(Self::listening_ports_for_pid(pid)); + if let Some(pid) = pid + && let Ok(ports) = Self::listening_ports_for_pid(pid) + { + candidates.extend(ports); } if let Some(ep) = extension_port.filter(|&p| p > 0) { candidates.extend((0..20u16).map(|offset| ep.saturating_add(offset))); @@ -287,47 +325,107 @@ impl AntigravityProvider { } } - /// Enumerate the TCP ports a given PID is listening on (Windows `lsof` equivalent). - /// On Windows this uses `Get-NetTCPConnection`; it returns an empty list on any failure - /// so the caller deterministically falls back to the heuristic candidate ports. + /// Enumerate IPv4 TCP listener ports for a PID through the Windows IP Helper API. + /// This avoids starting PowerShell inside the managed readiness poll. #[cfg(windows)] - fn listening_ports_for_pid(pid: u32) -> Vec { - const CREATE_NO_WINDOW: u32 = 0x08000000; - - let mut cmd = Command::new("powershell.exe"); - cmd.args([ - "-ExecutionPolicy", - "Bypass", - "-Command", - &format!( - "Get-NetTCPConnection -OwningProcess {pid} -State Listen \ - -ErrorAction SilentlyContinue | Select-Object -ExpandProperty LocalPort" - ), - ]); - cmd.creation_flags(CREATE_NO_WINDOW); - - let Ok(output) = cmd.output() else { - return Vec::new(); + fn listening_ports_for_pid(pid: u32) -> Result, ProviderError> { + const AF_INET_FAMILY: u32 = 2; + const NO_ERROR: u32 = 0; + + let mut bytes = 0_u32; + // SAFETY: the first call supplies no destination buffer and only asks Windows + // for the required byte count. + let query = unsafe { + GetExtendedTcpTable( + None, + &mut bytes, + false, + AF_INET_FAMILY, + TCP_TABLE_OWNER_PID_LISTENER, + 0, + ) }; - if !output.status.success() { - return Vec::new(); + if query != ERROR_INSUFFICIENT_BUFFER.0 && query != NO_ERROR { + return Err(ProviderError::Other(format!( + "Failed to size the Windows TCP listener table (error {query})" + ))); + } + if bytes < u32::try_from(std::mem::size_of::()).unwrap_or(u32::MAX) { + return Ok(Vec::new()); } - let stdout = String::from_utf8_lossy(&output.stdout); - let mut ports: Vec = stdout - .lines() - .filter_map(|l| l.trim().parse::().ok()) + let mut buffer = Vec::new(); + let mut loaded = false; + // The table can grow between the sizing call and the read. Retry with + // the updated size instead of failing a refresh on that benign race. + for _ in 0..3 { + // A u32 allocation supplies the alignment required by the all-DWORD MIB rows. + let words = (bytes as usize).div_ceil(std::mem::size_of::()); + buffer.resize(words, 0_u32); + // SAFETY: `buffer` is writable for at least `bytes` bytes and remains alive + // while the returned table is inspected. + let result = unsafe { + GetExtendedTcpTable( + Some(buffer.as_mut_ptr().cast()), + &mut bytes, + false, + AF_INET_FAMILY, + TCP_TABLE_OWNER_PID_LISTENER, + 0, + ) + }; + if result == NO_ERROR { + loaded = true; + break; + } + if result != ERROR_INSUFFICIENT_BUFFER.0 { + return Err(ProviderError::Other(format!( + "Failed to read the Windows TCP listener table (error {result})" + ))); + } + } + if !loaded { + return Err(ProviderError::Other( + "Windows TCP listener table kept changing during the query".to_string(), + )); + } + + let table = buffer.as_ptr().cast::(); + // SAFETY: Windows initialized the header on the successful call above. + let count = unsafe { (*table).dwNumEntries as usize }; + let rows_offset = std::mem::offset_of!(MIB_TCPTABLE_OWNER_PID, table); + let available = (bytes as usize).saturating_sub(rows_offset); + let max_rows = available / std::mem::size_of::(); + if count > max_rows { + return Err(ProviderError::Parse( + "Windows returned an invalid TCP listener table".to_string(), + )); + } + // SAFETY: `count` was bounded by the returned buffer size. Windows lays out + // the fixed-size owner-PID rows consecutively after the table header. + let rows = unsafe { + std::slice::from_raw_parts( + std::ptr::addr_of!((*table).table).cast::(), + count, + ) + }; + let mut ports: Vec = rows + .iter() + .filter(|row| row.dwOwningPid == pid) + .filter_map(|row| u16::try_from(row.dwLocalPort).ok()) + .map(u16::from_be) + .filter(|port| *port != 0) .collect(); ports.sort_unstable(); ports.dedup(); - ports + Ok(ports) } /// Non-Windows platforms have no `Get-NetTCPConnection`; return an empty list by design so /// the caller falls back to the heuristic candidate ports. #[cfg(not(windows))] - fn listening_ports_for_pid(_pid: u32) -> Vec { - Vec::new() + fn listening_ports_for_pid(_pid: u32) -> Result, ProviderError> { + Ok(Vec::new()) } /// Fetch user status from Antigravity API. @@ -346,11 +444,21 @@ impl AntigravityProvider { usage } - async fn fetch_user_status(&self) -> Result { - let process_info = Self::detect_process_info()?; + async fn fetch_user_status(&self) -> Result, ProviderError> { + let process_info = tokio::task::spawn_blocking(Self::detect_process_info) + .await + .map_err(|error| { + ProviderError::Other(format!( + "Failed to join the Antigravity process detector: {error}" + )) + })??; + let Some(process_info) = process_info else { + return Ok(None); + }; let api_port = Self::find_api_port(process_info.extension_port, process_info.pid).await?; self.fetch_user_status_at_port(&process_info, api_port) .await + .map(Some) } async fn fetch_user_status_at_port( @@ -440,13 +548,43 @@ impl AntigravityProvider { } /// Start a short-lived, headless `agy` session when neither the Antigravity - /// desktop app nor a user-owned CLI session is running. The child is kept - /// alive only for this fetch and is always reaped by `ManagedAgyProcess`. - async fn fetch_with_managed_agy(&self) -> Result { - let _launch_guard = MANAGED_AGY_FETCH.lock().await; + /// desktop app nor a user-owned CLI session is running. The deadline includes + /// launch serialization, the after-lock recheck, startup, probing and cleanup. + #[cfg(windows)] + async fn fetch_with_managed_agy(&self) -> Result { + let deadline = Instant::now() + AGY_ATTEMPT_TIMEOUT; + let lock_budget = deadline.saturating_duration_since(Instant::now()); + let _launch_guard = tokio::time::timeout(lock_budget, MANAGED_AGY_FETCH.lock()) + .await + .map_err(|_| { + ProviderError::Other( + "Timed out waiting for another managed agy refresh to finish".to_string(), + ) + })?; + + // A desktop app or user-owned CLI may have appeared while this request + // waited for the launch lock. Reuse it and never include it in our job. + let recheck_budget = deadline.saturating_duration_since(Instant::now()); + if recheck_budget.is_zero() { + return Err(Self::managed_agy_timeout()); + } + match tokio::time::timeout(recheck_budget, self.fetch_user_status()).await { + Ok(Ok(Some(usage))) => return Ok(ManagedAgyOutcome::Reused(usage)), + Ok(Ok(None)) => {} + Ok(Err(error)) => return Err(error), + Err(_) => return Err(Self::managed_agy_timeout()), + } - let binary = Self::locate_agy_binary() - .ok_or_else(|| ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.to_string()))?; + let Some(binary) = Self::locate_agy_binary() else { + return Ok(ManagedAgyOutcome::Missing); + }; + let probe_client = crate::core::credentialed_http_client_builder() + .no_proxy() + .timeout(AGY_PROBE_TIMEOUT) + .danger_accept_invalid_certs(true) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|error| ProviderError::Other(error.to_string()))?; let mut managed = ManagedAgyProcess::spawn(&binary)?; let pid = managed.id(); let process_info = ProcessInfo { @@ -456,68 +594,100 @@ impl AntigravityProvider { pid: Some(pid), source: ProcessSource::Cli, }; - let deadline = Instant::now() + AGY_READY_TIMEOUT; - let probe_client = crate::core::credentialed_http_client_builder() - .no_proxy() - .timeout(Duration::from_secs(2)) - .danger_accept_invalid_certs(true) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|error| ProviderError::Other(error.to_string()))?; - let mut last_error = None; - loop { - if let Some(status) = managed.try_wait()? { - return Err(ProviderError::NotInstalled(format!( - "agy exited before its local quota service was ready ({status}). Open Antigravity or run agy and sign in, then retry." - ))); - } + let result = async { + let work_deadline = deadline.checked_sub(AGY_CLEANUP_RESERVE).unwrap_or(deadline); + let mut last_error = None; + loop { + if let Some(status) = managed.try_wait()? { + return Err(ProviderError::NotInstalled(format!( + "agy exited before its local quota service was ready ({status}). Open Antigravity or run agy and sign in, then retry." + ))); + } - let ports = Self::listening_ports_for_pid(pid); - for port in ports { - if !Self::probe_api_port(&probe_client, port).await { - continue; + match Self::listening_ports_for_pid(pid) { + Ok(ports) => { + if let Some(port) = Self::first_ready_api_port(&probe_client, ports).await { + let remaining = work_deadline + .saturating_duration_since(Instant::now()); + if remaining.is_zero() { + break; + } + match tokio::time::timeout( + remaining, + self.fetch_user_status_at_port(&process_info, port), + ) + .await + { + Ok(Ok(usage)) => return Ok(ManagedAgyOutcome::Fetched(usage)), + Ok(Err(ProviderError::AuthRequired)) => { + return Err(ProviderError::AuthRequired); + } + Ok(Err(error)) => last_error = Some(error), + Err(_) => break, + } + } + } + Err(error) => last_error = Some(error), } - let remaining = deadline.saturating_duration_since(Instant::now()); + let remaining = work_deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() { break; } - match tokio::time::timeout( - remaining, - self.fetch_user_status_at_port(&process_info, port), - ) - .await - { - Ok(Ok(usage)) => return Ok(usage), - Ok(Err(ProviderError::AuthRequired)) => { - return Err(ProviderError::AuthRequired); - } - Ok(Err(error)) => last_error = Some(error), - Err(_) => break, - } + tokio::time::sleep(AGY_READY_POLL_INTERVAL.min(remaining)).await; } - if Instant::now() >= deadline { - break; + if let Some(error) = last_error { + tracing::debug!(%error, "managed agy quota service did not become ready"); } - tokio::time::sleep( - AGY_READY_POLL_INTERVAL.min(deadline.saturating_duration_since(Instant::now())), - ) - .await; + Err(Self::managed_agy_timeout()) } + .await; + + managed.shutdown().await; + result + } - if let Some(error) = last_error { - tracing::debug!(%error, "managed agy quota service did not become ready"); + #[cfg(windows)] + async fn first_ready_api_port(client: &reqwest::Client, ports: Vec) -> Option { + let mut probes = stream::iter( + ports + .into_iter() + .map(|port| async move { (port, Self::probe_api_port(client, port).await) }), + ) + .buffer_unordered(4); + while let Some((port, ready)) = probes.next().await { + if ready { + return Some(port); + } } - Err(ProviderError::Other( - "agy started but its quota service did not become ready. Open Antigravity or run agy and sign in, then retry." + None + } + + #[cfg(windows)] + fn managed_agy_timeout() -> ProviderError { + ProviderError::Other( + "agy started but its quota service did not become ready before the managed refresh deadline. Open Antigravity or run agy and sign in, then retry." .to_string(), - )) + ) } - fn is_not_running_error(error: &ProviderError) -> bool { - matches!(error, ProviderError::NotInstalled(message) if message == NOT_RUNNING_MESSAGE) + fn offline_usage_result() -> Option { + let count = local_sessions::offline_conversation_count(); + if count == 0 { + return None; + } + let noun = if count == 1 { + "conversation" + } else { + "conversations" + }; + let usage = UsageSnapshot::new(RateWindow::informational(format!( + "Offline · {count} {noun}" + ))) + .with_login_method("offline"); + Some(ProviderFetchResult::new(usage, "offline")) } fn locate_agy_binary() -> Option { @@ -779,38 +949,36 @@ impl Provider for AntigravityProvider { tracing::debug!("Fetching Antigravity usage via local probe"); match self.fetch_user_status().await { - Ok(usage) => Ok(ProviderFetchResult::new( + Ok(Some(usage)) => Ok(ProviderFetchResult::new( Self::with_cadence_labels(usage), "local", )), - Err(error) if Self::is_not_running_error(&error) => { - match self.fetch_with_managed_agy().await { - Ok(usage) => { - return Ok(ProviderFetchResult::new( - Self::with_cadence_labels(usage), - "cli", - )); - } - Err(ProviderError::NotInstalled(message)) - if message == AGY_NOT_FOUND_MESSAGE => {} - Err(error) => { - tracing::warn!(%error, "managed Antigravity CLI probe failed"); - return Err(error); + Ok(None) => { + #[cfg(windows)] + { + match self.fetch_with_managed_agy().await { + Ok(ManagedAgyOutcome::Reused(usage)) => { + return Ok(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "local", + )); + } + Ok(ManagedAgyOutcome::Fetched(usage)) => { + return Ok(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "cli", + )); + } + Ok(ManagedAgyOutcome::Missing) => {} + Err(error) => { + tracing::warn!(%error, "managed Antigravity CLI probe failed"); + return Err(error); + } } } - let count = local_sessions::offline_conversation_count(); - if count > 0 { - let noun = if count == 1 { - "conversation" - } else { - "conversations" - }; - let usage = UsageSnapshot::new(RateWindow::informational(format!( - "Offline · {count} {noun}" - ))) - .with_login_method("offline"); - return Ok(ProviderFetchResult::new(usage, "offline")); + if let Some(result) = Self::offline_usage_result() { + return Ok(result); } Err(ProviderError::NotInstalled( AGY_NOT_FOUND_MESSAGE.to_string(), @@ -857,17 +1025,31 @@ struct ProcessInfo { source: ProcessSource, } +#[cfg(windows)] +enum ManagedAgyOutcome { + /// A user-owned desktop or CLI process appeared after the launch lock. + Reused(UsageSnapshot), + /// Usage came from the short-lived process owned by this fetch. + Fetched(UsageSnapshot), + /// No configured `agy` executable exists, so offline history may be used. + Missing, +} + /// RAII owner for the exact `agy` process started by this provider. Dropping it /// cannot affect Antigravity or CLI processes that were already running. +#[cfg(windows)] struct ManagedAgyProcess { - child: Box, + child: Option>, pid: u32, + job: Option, master: Option>, drain_thread: Option>, } +#[cfg(windows)] impl ManagedAgyProcess { fn spawn(binary: &std::path::Path) -> Result { + let job = create_managed_agy_job()?; let pty_system = portable_pty::native_pty_system(); let pair = pty_system .openpty(portable_pty::PtySize { @@ -904,25 +1086,44 @@ impl ManagedAgyProcess { "Failed to determine the managed agy process id".to_string(), )); }; + let Some(process_handle) = child.as_raw_handle() else { + drop(child.kill()); + drop(child.wait()); + return Err(ProviderError::Other( + "Failed to access the managed agy process handle".to_string(), + )); + }; + if let Err(error) = assign_process_to_job(&job, process_handle) { + drop(child.kill()); + drop(child.wait()); + return Err(error); + } let drain_thread = std::thread::spawn(move || { // Drain without logging: terminal output can contain account data. // Windows ConPTY programs may request the cursor position and wait // for a terminal response before continuing initialization. let mut buffer = [0_u8; 4096]; let mut tail = Vec::with_capacity(3); + let mut cursor_replies = 0_usize; while let Ok(read) = reader.read(&mut buffer) { if read == 0 { break; } - if terminal_requested_cursor_position(&mut tail, &buffer[..read]) { + let requested = terminal_cursor_position_request_count(&mut tail, &buffer[..read]); + let allowed = terminal_cursor_reply_allowance(cursor_replies, requested); + for _ in 0..allowed { drop(writer.write_all(b"\x1b[1;1R")); + } + if allowed > 0 { drop(writer.flush()); + cursor_replies += allowed; } } }); Ok(Self { - child, + child: Some(child), pid, + job: Some(job), master: Some(pair.master), drain_thread: Some(drain_thread), }) @@ -933,23 +1134,134 @@ impl ManagedAgyProcess { } fn try_wait(&mut self) -> Result, ProviderError> { - self.child.try_wait().map_err(|error| { - ProviderError::Other(format!("Failed to inspect the agy CLI: {error}")) + self.child + .as_mut() + .expect("managed agy child is present until cleanup") + .try_wait() + .map_err(|error| { + ProviderError::Other(format!("Failed to inspect the agy CLI: {error}")) + }) + } + + async fn shutdown(mut self) { + let Some(resources) = self.take_resources() else { + return; + }; + let cleanup = tokio::task::spawn_blocking(move || resources.terminate_and_reap()); + // A stuck platform wait must not hold the async provider worker. The + // blocking cleanup task remains detached and still owns every handle. + drop(tokio::time::timeout(AGY_CLEANUP_RESERVE, cleanup).await); + } + + fn take_resources(&mut self) -> Option { + Some(ManagedAgyResources { + child: self.child.take()?, + job: Some( + self.job + .take() + .expect("managed agy job is present until cleanup"), + ), + master: self.master.take(), + drain_thread: self.drain_thread.take(), }) } } +#[cfg(windows)] impl Drop for ManagedAgyProcess { fn drop(&mut self) { - if self.child.try_wait().ok().flatten().is_none() { + let Some(mut resources) = self.take_resources() else { + return; + }; + resources.terminate(); + // Drop can run when an outer timeout cancels the fetch. Reaping and + // joining the terminal drain must therefore never block that worker. + drop( + std::thread::Builder::new() + .name("codexbar-agy-cleanup".to_string()) + .spawn(move || resources.reap()), + ); + } +} + +#[cfg(windows)] +struct ManagedAgyResources { + child: Box, + job: Option, + master: Option>, + drain_thread: Option>, +} + +#[cfg(windows)] +impl ManagedAgyResources { + fn terminate(&mut self) { + // SAFETY: this job is private to the single process launched above; + // user-owned Antigravity and agy processes were never assigned to it. + let terminated = self + .job + .as_ref() + .is_some_and(|job| unsafe { TerminateJobObject(win_handle(job), 1) }.is_ok()); + // KILL_ON_JOB_CLOSE is the second termination path if the explicit API + // fails. Close it before wait so a failure cannot strand the reaper. + drop(self.job.take()); + if !terminated { drop(self.child.kill()); } + } + + fn reap(mut self) { drop(self.child.wait()); drop(self.master.take()); if let Some(thread) = self.drain_thread.take() { drop(thread.join()); } } + + fn terminate_and_reap(mut self) { + self.terminate(); + self.reap(); + } +} + +#[cfg(windows)] +fn create_managed_agy_job() -> Result { + // SAFETY: a successful call transfers a unique job handle to this owner. + let raw = unsafe { CreateJobObjectW(None, PCWSTR::null()) } + .map_err(|error| ProviderError::Other(format!("Failed to create agy job: {error}")))?; + // SAFETY: `raw` is a unique valid handle returned by CreateJobObjectW. + let job = unsafe { OwnedHandle::from_raw_handle(raw.0) }; + let limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION { + BasicLimitInformation: JOBOBJECT_BASIC_LIMIT_INFORMATION { + LimitFlags: JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + ..Default::default() + }, + ..Default::default() + }; + let size = u32::try_from(std::mem::size_of_val(&limits)) + .map_err(|error| ProviderError::Other(format!("Invalid agy job limit size: {error}")))?; + // SAFETY: `job` is valid and `limits` is initialized for the requested class. + unsafe { + SetInformationJobObject( + win_handle(&job), + JobObjectExtendedLimitInformation, + (&limits as *const JOBOBJECT_EXTENDED_LIMIT_INFORMATION).cast(), + size, + ) + } + .map_err(|error| ProviderError::Other(format!("Failed to configure agy job: {error}")))?; + Ok(job) +} + +#[cfg(windows)] +fn assign_process_to_job(job: &OwnedHandle, process: RawHandle) -> Result<(), ProviderError> { + // SAFETY: both handles are valid and remain owned by their respective wrappers. + unsafe { AssignProcessToJobObject(win_handle(job), HANDLE(process)) } + .map_err(|error| ProviderError::Other(format!("Failed to contain agy process: {error}"))) +} + +#[cfg(windows)] +fn win_handle(value: &OwnedHandle) -> HANDLE { + HANDLE(value.as_raw_handle()) } // API Response types diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index bf2383ce72..c6ea8614b3 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -219,10 +219,11 @@ fn test_noisy_models_do_not_drive_summary_windows() { } #[test] -fn not_running_error_tells_user_how_to_start() { - let error = ProviderError::NotInstalled(NOT_RUNNING_MESSAGE.to_string()).to_string(); +fn missing_cli_error_explains_runtime_state() { + let error = ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.to_string()).to_string(); - assert!(error.contains("Start Google Antigravity and sign in")); + assert!(error.contains("not running")); + assert!(error.contains("agy CLI was not found")); } #[test] @@ -251,31 +252,88 @@ fn managed_agy_candidates_prefer_override_then_path_then_known_installs() { } #[test] -fn managed_agy_not_running_check_is_exact() { - assert!(AntigravityProvider::is_not_running_error( - &ProviderError::NotInstalled(NOT_RUNNING_MESSAGE.to_string()) - )); - assert!(!AntigravityProvider::is_not_running_error( - &ProviderError::NotInstalled("Failed to detect Antigravity process".to_string()) - )); - assert!(!AntigravityProvider::is_not_running_error( - &ProviderError::AuthRequired - )); +fn managed_agy_terminal_detects_cursor_request_across_reads() { + let mut tail = Vec::new(); + + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"ready\x1b["), + 0 + ); + assert_eq!(terminal_cursor_position_request_count(&mut tail, b"6n"), 1); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"plain output"), + 0 + ); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"\x1b[6nmore\x1b[6n"), + 2 + ); } #[test] -fn managed_agy_terminal_detects_cursor_request_across_reads() { - let mut tail = Vec::new(); +fn managed_agy_terminal_caps_cursor_replies() { + assert_eq!(terminal_cursor_reply_allowance(0, 2), 2); + assert_eq!(terminal_cursor_reply_allowance(31, 4), 1); + assert_eq!( + terminal_cursor_reply_allowance(AGY_MAX_CURSOR_REPLIES, 1), + 0 + ); +} - assert!(!terminal_requested_cursor_position( - &mut tail, - b"ready\x1b[" - )); - assert!(terminal_requested_cursor_position(&mut tail, b"6n")); - assert!(!terminal_requested_cursor_position( - &mut tail, - b"plain output" - )); +#[cfg(windows)] +#[test] +fn windows_listener_table_finds_current_process_port() { + let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) + .expect("bind a local IPv4 listener"); + let port = listener.local_addr().expect("listener address").port(); + + let ports = AntigravityProvider::listening_ports_for_pid(std::process::id()) + .expect("read the Windows TCP listener table"); + + assert!( + ports.contains(&port), + "listener table should contain {port}" + ); +} + +#[cfg(windows)] +#[test] +fn managed_agy_job_terminates_its_owned_process() { + use std::os::windows::io::AsRawHandle as _; + use std::os::windows::process::CommandExt as _; + + const CREATE_NO_WINDOW: u32 = 0x08000000; + let mut command = std::process::Command::new("powershell.exe"); + command + .args([ + "-NoLogo", + "-NoProfile", + "-Command", + "Start-Sleep -Seconds 30", + ]) + .creation_flags(CREATE_NO_WINDOW); + let mut child = command.spawn().expect("spawn an isolated test child"); + let job = create_managed_agy_job().expect("create a kill-on-close job"); + if let Err(error) = assign_process_to_job(&job, child.as_raw_handle()) { + drop(child.kill()); + drop(child.wait()); + panic!("assign the test child to its job: {error}"); + } + + // SAFETY: only the isolated test child was assigned to this private job. + unsafe { TerminateJobObject(win_handle(&job), 1) }.expect("terminate the private job"); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(3); + loop { + if child.try_wait().expect("inspect the test child").is_some() { + break; + } + if std::time::Instant::now() >= deadline { + drop(child.kill()); + drop(child.wait()); + panic!("job termination did not stop the test child"); + } + std::thread::sleep(std::time::Duration::from_millis(25)); + } } // ── agy CLI process matching ─────────────────────────────────────── @@ -449,7 +507,7 @@ fn not_installed_maps_to_local_runtime_offline() { // credential problem. assert_eq!( AntigravityProvider::new() - .error_state_kind(&ProviderError::NotInstalled(NOT_RUNNING_MESSAGE.into())), + .error_state_kind(&ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.into())), crate::core::ProviderStateKind::LocalRuntimeOffline ); } From ac5edc80c5a9e2c855d85afc8b4a913b30f071af Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 13:33:57 +0700 Subject: [PATCH 03/10] Restore Antigravity offline fallback on probe failure Sign-in failures (AuthRequired) keep surfacing, but non-auth local-probe and managed agy start failures now fall through to the offline conversation-history snapshot instead of replacing it with a hard error. Build the agy candidate-path test expectations with PathBuf::join so the test passes on non-Windows hosts, and note IPv6 listener discovery as an accepted follow-up. --- rust/src/providers/antigravity/mod.rs | 54 ++++++++++++++++---- rust/src/providers/antigravity/tests.rs | 67 ++++++++++++++++++++----- 2 files changed, 99 insertions(+), 22 deletions(-) diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 0fcc491557..4c615f3425 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -327,6 +327,11 @@ impl AntigravityProvider { /// Enumerate IPv4 TCP listener ports for a PID through the Windows IP Helper API. /// This avoids starting PowerShell inside the managed readiness poll. + /// + /// Known follow-up: only the AF_INET table is enumerated and candidates are + /// probed at `127.0.0.1`, so an IPv6-only loopback listener would be missed. + /// Accepted for now because the managed `agy` service is observed to bind + /// IPv4 on Windows; add AF_INET6 enumeration with `[::1]` probes later. #[cfg(windows)] fn listening_ports_for_pid(pid: u32) -> Result, ProviderError> { const AF_INET_FAMILY: u32 = 2; @@ -690,6 +695,27 @@ impl AntigravityProvider { Some(ProviderFetchResult::new(usage, "offline")) } + /// Resolve a failure to obtain live usage. + /// + /// A failed sign-in is actionable, so it always surfaces. Every other + /// failure means the runtime/CLI is unavailable or inconclusive, so an + /// available offline conversation-history snapshot is preferred over + /// discarding it for a transient error. + fn resolve_probe_failure( + error: ProviderError, + offline: Option, + ) -> Result { + if matches!(error, ProviderError::AuthRequired) { + return Err(error); + } + offline.ok_or(error) + } + + fn offline_or_unavailable() -> Result { + Self::offline_usage_result() + .ok_or_else(|| ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.to_string())) + } + fn locate_agy_binary() -> Option { let candidates = Self::agy_binary_candidates( std::env::var_os("ANTIGRAVITY_CLI_PATH").map(PathBuf::from), @@ -970,23 +996,31 @@ impl Provider for AntigravityProvider { )); } Ok(ManagedAgyOutcome::Missing) => {} + // A signed-out CLI is actionable, so surface it rather + // than hiding it behind offline history. Any other + // managed-start failure still leaves the runtime + // unavailable, so keep the offline-history fallback. Err(error) => { - tracing::warn!(%error, "managed Antigravity CLI probe failed"); - return Err(error); + if !matches!(error, ProviderError::AuthRequired) { + tracing::debug!(%error, "managed Antigravity CLI probe failed"); + } + return Self::resolve_probe_failure( + error, + Self::offline_usage_result(), + ); } } } - if let Some(result) = Self::offline_usage_result() { - return Ok(result); - } - Err(ProviderError::NotInstalled( - AGY_NOT_FOUND_MESSAGE.to_string(), - )) + Self::offline_or_unavailable() } Err(error) => { - tracing::warn!(%error, "Antigravity local probe failed"); - Err(error) + // The local probe is inconclusive (e.g. PowerShell unavailable); + // preserve offline history before surfacing the probe error. + if !matches!(error, ProviderError::AuthRequired) { + tracing::debug!(%error, "Antigravity local probe failed"); + } + Self::resolve_probe_failure(error, Self::offline_usage_result()) } } } diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index c6ea8614b3..08b6abf63b 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -228,26 +228,31 @@ fn missing_cli_error_explains_runtime_state() { #[test] fn managed_agy_candidates_prefer_override_then_path_then_known_installs() { + let explicit = PathBuf::from(r"D:\tools\agy.exe"); + let path_lookup = PathBuf::from(r"C:\path\agy.exe"); + let local_app_data = PathBuf::from(r"C:\Users\test\AppData\Local"); + let home = PathBuf::from(r"C:\Users\test"); + let candidates = AntigravityProvider::agy_binary_candidates( - Some(PathBuf::from(r"D:\tools\agy.exe")), - Some(PathBuf::from(r"C:\path\agy.exe")), - Some(PathBuf::from(r"C:\Users\test\AppData\Local")), - Some(PathBuf::from(r"C:\Users\test")), + Some(explicit.clone()), + Some(path_lookup.clone()), + Some(local_app_data.clone()), + Some(home.clone()), ); - assert_eq!(candidates[0], PathBuf::from(r"D:\tools\agy.exe")); - assert_eq!(candidates[1], PathBuf::from(r"C:\path\agy.exe")); + assert_eq!(candidates[0], explicit); + assert_eq!(candidates[1], path_lookup); + // Build expectations with `join` so the assertions match on every host: + // on Unix `\` is an ordinary character and `join` inserts `/`. assert_eq!( candidates[2], - PathBuf::from(r"C:\Users\test\AppData\Local\agy\bin\agy.exe") + local_app_data.join("agy").join("bin").join("agy.exe") ); assert_eq!( candidates[3], - PathBuf::from(r"C:\Users\test\.local\bin").join(if cfg!(windows) { - "agy.exe" - } else { - "agy" - }) + home.join(".local") + .join("bin") + .join(if cfg!(windows) { "agy.exe" } else { "agy" }) ); } @@ -523,3 +528,41 @@ fn probe_failure_maps_to_unknown() { crate::core::ProviderStateKind::Unknown ); } + +// ── Offline-history fallback on probe failure ────────────────────── + +fn offline_result() -> ProviderFetchResult { + ProviderFetchResult::new( + UsageSnapshot::new(RateWindow::new(0.0)).with_login_method("offline"), + "offline", + ) +} + +#[test] +fn auth_required_surfaces_instead_of_offline_history() { + let resolved = AntigravityProvider::resolve_probe_failure( + ProviderError::AuthRequired, + Some(offline_result()), + ); + assert!(matches!(resolved, Err(ProviderError::AuthRequired))); +} + +#[test] +fn non_auth_failure_prefers_offline_history() { + // A transient managed-start or local-probe failure must not discard the + // existing offline conversation-history snapshot. + let resolved = AntigravityProvider::resolve_probe_failure( + ProviderError::Other("agy readiness timeout".to_string()), + Some(offline_result()), + ); + assert!(resolved.is_ok()); +} + +#[test] +fn non_auth_failure_without_history_surfaces_error() { + let resolved = AntigravityProvider::resolve_probe_failure( + ProviderError::Other("agy readiness timeout".to_string()), + None, + ); + assert!(matches!(resolved, Err(ProviderError::Other(_)))); +} From 4a01b317ca57cfad62990403141d14984b3fb9c8 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:17:02 +0700 Subject: [PATCH 04/10] refactor(antigravity): isolate managed process lifecycle --- rust/src/lib.rs | 2 + rust/src/managed_process.rs | 648 ++++++++++++++++++++++++ rust/src/providers/antigravity/mod.rs | 452 +++-------------- rust/src/providers/antigravity/tests.rs | 99 +--- 4 files changed, 744 insertions(+), 457 deletions(-) create mode 100644 rust/src/managed_process.rs diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 3fda0076e7..06046a2e5c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -16,6 +16,8 @@ pub mod host; pub mod locale; pub mod logging; pub mod login; +#[cfg(windows)] +pub mod managed_process; pub mod notifications; pub mod providers; pub mod secure_file; diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs new file mode 100644 index 0000000000..adde854936 --- /dev/null +++ b/rust/src/managed_process.rs @@ -0,0 +1,648 @@ +//! Provider-neutral ownership of a short-lived, task-owned interactive child +//! process. +//! +//! A provider decides *whether* to launch a helper and supplies its +//! configuration; this module owns the *lifecycle* of the child it created: +//! PTY start, Windows Job Object containment, terminal drain, readiness +//! introspection, restart, shutdown and Drop cleanup. It never adopts or +//! terminates a process it did not create, so user-owned processes are always +//! isolated from cleanup. + +use std::ffi::OsString; +use std::io::{Read, Write}; +use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; +use std::path::PathBuf; +use std::time::Duration; + +use windows::Win32::Foundation::{ERROR_INSUFFICIENT_BUFFER, HANDLE}; +use windows::Win32::NetworkManagement::IpHelper::{ + GetExtendedTcpTable, MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID, TCP_TABLE_OWNER_PID_LISTENER, +}; +use windows::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_BASIC_LIMIT_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, + JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject, +}; +use windows::core::PCWSTR; + +/// Maximum number of terminal cursor-position replies sent to one child. +const MAX_CURSOR_REPLIES: usize = 32; + +/// Provider-supplied configuration for one managed child process. +/// +/// The provider owns this policy: which executable, arguments and environment +/// to use, and what readiness means for its protocol. The owner only enforces +/// lifecycle and containment. +#[derive(Debug, Clone)] +pub struct ManagedProcessConfig { + /// Executable to launch. + pub program: PathBuf, + /// Arguments passed to the executable. + pub args: Vec, + /// Additional environment variables for the child. + pub env: Vec<(OsString, OsString)>, + /// Working directory, when the provider wants to pin one. + pub cwd: Option, + /// PTY geometry. + pub pty_rows: u16, + /// PTY geometry. + pub pty_cols: u16, + /// Short label used only in diagnostics (e.g. `"agy"`). + pub label: String, +} + +/// Error raised by the managed-process owner. It stays provider-neutral; the +/// caller maps it into its own error surface. +#[derive(Debug, thiserror::Error)] +#[error("{0}")] +pub struct ManagedProcessError(String); + +impl ManagedProcessError { + fn new(message: impl Into) -> Self { + Self(message.into()) + } +} + +/// Result type for managed-process operations. +pub type ManagedProcessResult = Result; + +/// RAII owner for the exact child process it started. Dropping it cannot affect +/// any process that was already running. +pub struct ManagedProcess { + config: ManagedProcessConfig, + child: Option>, + pid: u32, + job: Option, + master: Option>, + drain_thread: Option>, +} + +impl ManagedProcess { + /// Start `config.program` in a PTY inside its own kill-on-close job. + pub fn spawn(config: &ManagedProcessConfig) -> ManagedProcessResult { + let job = create_managed_job(&config.label)?; + let pty_system = portable_pty::native_pty_system(); + let pair = pty_system + .openpty(portable_pty::PtySize { + rows: config.pty_rows, + cols: config.pty_cols, + pixel_width: 0, + pixel_height: 0, + }) + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to create a terminal for {}: {error}", + config.label + )) + })?; + let mut command = portable_pty::CommandBuilder::new(config.program.as_os_str()); + for arg in &config.args { + command.arg(arg); + } + if let Some(cwd) = &config.cwd { + command.cwd(cwd.as_os_str()); + } + for (key, value) in &config.env { + command.env(key, value); + } + + let reader = pair.master.try_clone_reader().map_err(|error| { + ManagedProcessError::new(format!( + "Failed to read the {} terminal: {error}", + config.label + )) + })?; + let writer = pair.master.take_writer().map_err(|error| { + ManagedProcessError::new(format!( + "Failed to open the {} terminal: {error}", + config.label + )) + })?; + let mut child = pair.slave.spawn_command(command).map_err(|error| { + ManagedProcessError::new(format!( + "Failed to launch the {} CLI: {error}", + config.label + )) + })?; + drop(pair.slave); + + let Some(pid) = child.process_id() else { + drop(child.kill()); + drop(child.wait()); + return Err(ManagedProcessError::new(format!( + "Failed to determine the managed {} process id", + config.label + ))); + }; + let Some(process_handle) = child.as_raw_handle() else { + drop(child.kill()); + drop(child.wait()); + return Err(ManagedProcessError::new(format!( + "Failed to access the managed {} process handle", + config.label + ))); + }; + if let Err(error) = assign_process_to_job(&job, process_handle, &config.label) { + drop(child.kill()); + drop(child.wait()); + return Err(error); + } + let drain_thread = spawn_drain_thread(reader, writer); + Ok(Self { + config: config.clone(), + child: Some(child), + pid, + job: Some(job), + master: Some(pair.master), + drain_thread: Some(drain_thread), + }) + } + + /// Process id of the owned child. + pub fn pid(&self) -> u32 { + self.pid + } + + /// Poll the owned child without blocking. + pub fn try_wait(&mut self) -> ManagedProcessResult> { + self.child + .as_mut() + .expect("managed child is present until cleanup") + .try_wait() + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to inspect the {} CLI: {error}", + self.config.label + )) + }) + } + + /// Candidate IPv4 loopback ports the owned child is currently listening on. + /// Providers use this to decide when the child's local service is ready. + pub fn listening_ports(&self) -> ManagedProcessResult> { + listening_ports_for_pid(self.pid) + } + + /// Terminate and reap the owned child, bounded by `cleanup_reserve`. + pub async fn shutdown(mut self, cleanup_reserve: Duration) { + let Some(resources) = self.take_resources() else { + return; + }; + let cleanup = tokio::task::spawn_blocking(move || resources.terminate_and_reap()); + // A stuck platform wait must not hold the async provider worker. The + // blocking cleanup task remains detached and still owns every handle. + drop(tokio::time::timeout(cleanup_reserve, cleanup).await); + } + + /// Replace the owned child with a fresh one from the same configuration. + /// The previous child is terminated and reaped before the new child starts, + /// so a restart cannot leak a process, job, or drain thread. + pub fn restart(&mut self) -> ManagedProcessResult<()> { + if let Some(resources) = self.take_resources() { + resources.terminate_and_reap(); + } + let replacement = Self::spawn(&self.config)?; + *self = replacement; + Ok(()) + } + + fn take_resources(&mut self) -> Option { + Some(ManagedProcessResources { + child: self.child.take()?, + job: Some( + self.job + .take() + .expect("managed job is present until cleanup"), + ), + master: self.master.take(), + drain_thread: self.drain_thread.take(), + }) + } +} + +impl Drop for ManagedProcess { + fn drop(&mut self) { + let Some(mut resources) = self.take_resources() else { + return; + }; + resources.terminate(); + // Drop can run when an outer timeout cancels the fetch. Reaping and + // joining the terminal drain must therefore never block that worker. + drop( + std::thread::Builder::new() + .name("codexbar-proc-cleanup".to_string()) + .spawn(move || resources.reap()), + ); + } +} + +struct ManagedProcessResources { + child: Box, + job: Option, + master: Option>, + drain_thread: Option>, +} + +impl ManagedProcessResources { + fn terminate(&mut self) { + // SAFETY: this job is private to the single process launched above; + // user-owned processes were never assigned to it. + let terminated = self + .job + .as_ref() + .is_some_and(|job| unsafe { TerminateJobObject(win_handle(job), 1) }.is_ok()); + // KILL_ON_JOB_CLOSE is the second termination path if the explicit API + // fails. Close it before wait so a failure cannot strand the reaper. + drop(self.job.take()); + if !terminated { + drop(self.child.kill()); + } + } + + fn reap(mut self) { + drop(self.child.wait()); + drop(self.master.take()); + if let Some(thread) = self.drain_thread.take() { + drop(thread.join()); + } + } + + fn terminate_and_reap(mut self) { + self.terminate(); + self.reap(); + } +} + +/// Drain PTY output without logging it: terminal output can contain account +/// data. Windows ConPTY programs may request the cursor position and wait for a +/// terminal response before continuing initialization, so answer a bounded +/// number of those requests. +fn spawn_drain_thread( + mut reader: Box, + mut writer: Box, +) -> std::thread::JoinHandle<()> { + std::thread::spawn(move || { + let mut buffer = [0_u8; 4096]; + let mut tail = Vec::with_capacity(3); + let mut cursor_replies = 0_usize; + while let Ok(read) = reader.read(&mut buffer) { + if read == 0 { + break; + } + let requested = terminal_cursor_position_request_count(&mut tail, &buffer[..read]); + let allowed = terminal_cursor_reply_allowance(cursor_replies, requested); + for _ in 0..allowed { + drop(writer.write_all(b"\x1b[1;1R")); + } + if allowed > 0 { + drop(writer.flush()); + cursor_replies += allowed; + } + } + }) +} + +fn terminal_cursor_position_request_count(tail: &mut Vec, chunk: &[u8]) -> usize { + tail.extend_from_slice(chunk); + let requested = tail.windows(4).filter(|bytes| *bytes == b"\x1b[6n").count(); + if tail.len() > 3 { + tail.drain(..tail.len() - 3); + } + requested +} + +fn terminal_cursor_reply_allowance(sent: usize, requested: usize) -> usize { + requested.min(MAX_CURSOR_REPLIES.saturating_sub(sent)) +} + +fn create_managed_job(label: &str) -> ManagedProcessResult { + // SAFETY: a successful call transfers a unique job handle to this owner. + let raw = unsafe { CreateJobObjectW(None, PCWSTR::null()) }.map_err(|error| { + ManagedProcessError::new(format!("Failed to create {label} job: {error}")) + })?; + // SAFETY: `raw` is a unique valid handle returned by CreateJobObjectW. + let job = unsafe { OwnedHandle::from_raw_handle(raw.0) }; + let limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION { + BasicLimitInformation: JOBOBJECT_BASIC_LIMIT_INFORMATION { + LimitFlags: JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + ..Default::default() + }, + ..Default::default() + }; + let size = u32::try_from(std::mem::size_of_val(&limits)).map_err(|error| { + ManagedProcessError::new(format!("Invalid {label} job limit size: {error}")) + })?; + // SAFETY: `job` is valid and `limits` is initialized for the requested class. + unsafe { + SetInformationJobObject( + win_handle(&job), + JobObjectExtendedLimitInformation, + (&limits as *const JOBOBJECT_EXTENDED_LIMIT_INFORMATION).cast(), + size, + ) + } + .map_err(|error| { + ManagedProcessError::new(format!("Failed to configure {label} job: {error}")) + })?; + Ok(job) +} + +fn assign_process_to_job( + job: &OwnedHandle, + process: RawHandle, + label: &str, +) -> ManagedProcessResult<()> { + // SAFETY: both handles are valid and remain owned by their respective wrappers. + unsafe { AssignProcessToJobObject(win_handle(job), HANDLE(process)) }.map_err(|error| { + ManagedProcessError::new(format!("Failed to contain {label} process: {error}")) + }) +} + +fn win_handle(value: &OwnedHandle) -> HANDLE { + HANDLE(value.as_raw_handle()) +} + +/// Enumerate IPv4 TCP listener ports for a PID through the Windows IP Helper API. +/// This avoids starting PowerShell inside a readiness poll. +/// +/// Known follow-up: only the AF_INET table is enumerated and candidates are +/// probed at `127.0.0.1`, so an IPv6-only loopback listener would be missed. +/// Accepted for now because the managed `agy` service is observed to bind IPv4 +/// on Windows; add AF_INET6 enumeration with `[::1]` probes later. +pub fn listening_ports_for_pid(pid: u32) -> ManagedProcessResult> { + const AF_INET_FAMILY: u32 = 2; + const NO_ERROR: u32 = 0; + + let mut bytes = 0_u32; + // SAFETY: the first call supplies no destination buffer and only asks Windows + // for the required byte count. + let query = unsafe { + GetExtendedTcpTable( + None, + &mut bytes, + false, + AF_INET_FAMILY, + TCP_TABLE_OWNER_PID_LISTENER, + 0, + ) + }; + if query != ERROR_INSUFFICIENT_BUFFER.0 && query != NO_ERROR { + return Err(ManagedProcessError::new(format!( + "Failed to size the Windows TCP listener table (error {query})" + ))); + } + if bytes < u32::try_from(std::mem::size_of::()).unwrap_or(u32::MAX) { + return Ok(Vec::new()); + } + + let mut buffer = Vec::new(); + let mut loaded = false; + // The table can grow between the sizing call and the read. Retry with + // the updated size instead of failing on that benign race. + for _ in 0..3 { + // A u32 allocation supplies the alignment required by the all-DWORD MIB rows. + let words = (bytes as usize).div_ceil(std::mem::size_of::()); + buffer.resize(words, 0_u32); + // SAFETY: `buffer` is writable for at least `bytes` bytes and remains alive + // while the returned table is inspected. + let result = unsafe { + GetExtendedTcpTable( + Some(buffer.as_mut_ptr().cast()), + &mut bytes, + false, + AF_INET_FAMILY, + TCP_TABLE_OWNER_PID_LISTENER, + 0, + ) + }; + if result == NO_ERROR { + loaded = true; + break; + } + if result != ERROR_INSUFFICIENT_BUFFER.0 { + return Err(ManagedProcessError::new(format!( + "Failed to read the Windows TCP listener table (error {result})" + ))); + } + } + if !loaded { + return Err(ManagedProcessError::new( + "Windows TCP listener table kept changing during the query".to_string(), + )); + } + + let table = buffer.as_ptr().cast::(); + // SAFETY: Windows initialized the header on the successful call above. + let count = unsafe { (*table).dwNumEntries as usize }; + let rows_offset = std::mem::offset_of!(MIB_TCPTABLE_OWNER_PID, table); + let available = (bytes as usize).saturating_sub(rows_offset); + let max_rows = available / std::mem::size_of::(); + if count > max_rows { + return Err(ManagedProcessError::new( + "Windows returned an invalid TCP listener table".to_string(), + )); + } + // SAFETY: `count` was bounded by the returned buffer size. Windows lays out + // the fixed-size owner-PID rows consecutively after the table header. + let rows = unsafe { + std::slice::from_raw_parts( + std::ptr::addr_of!((*table).table).cast::(), + count, + ) + }; + let mut ports: Vec = rows + .iter() + .filter(|row| row.dwOwningPid == pid) + .filter_map(|row| u16::try_from(row.dwLocalPort).ok()) + .map(u16::from_be) + .filter(|port| *port != 0) + .collect(); + ports.sort_unstable(); + ports.dedup(); + Ok(ports) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Instant; + + fn test_config() -> ManagedProcessConfig { + ManagedProcessConfig { + program: PathBuf::from("powershell.exe"), + args: vec![ + OsString::from("-NoLogo"), + OsString::from("-NoProfile"), + OsString::from("-NonInteractive"), + OsString::from("-Command"), + OsString::from("Start-Sleep -Seconds 30"), + ], + env: Vec::new(), + cwd: None, + pty_rows: 30, + pty_cols: 120, + label: "test".to_string(), + } + } + + fn start_test_process() -> ManagedProcess { + ManagedProcess::spawn(&test_config()).expect("start a managed test process") + } + + fn wait_for_exit(pid: u32) { + let deadline = Instant::now() + Duration::from_secs(5); + while process_is_alive(pid) { + assert!( + Instant::now() < deadline, + "managed child {pid} was not cleaned up" + ); + std::thread::sleep(Duration::from_millis(25)); + } + } + + 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, WaitForSingleObject, + }; + + // SAFETY: OpenProcess returns a handle owned by this function and closed below. + match unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 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, + } + } + + #[test] + fn terminal_detects_cursor_request_across_reads() { + let mut tail = Vec::new(); + + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"ready\x1b["), + 0 + ); + assert_eq!(terminal_cursor_position_request_count(&mut tail, b"6n"), 1); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"plain output"), + 0 + ); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"\x1b[6nmore\x1b[6n"), + 2 + ); + } + + #[test] + fn terminal_caps_cursor_replies() { + assert_eq!(terminal_cursor_reply_allowance(0, 2), 2); + assert_eq!(terminal_cursor_reply_allowance(31, 4), 1); + assert_eq!(terminal_cursor_reply_allowance(MAX_CURSOR_REPLIES, 1), 0); + } + + #[test] + fn listener_table_finds_current_process_port() { + let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) + .expect("bind a local IPv4 listener"); + let port = listener.local_addr().expect("listener address").port(); + + let ports = listening_ports_for_pid(std::process::id()) + .expect("read the Windows TCP listener table"); + + assert!( + ports.contains(&port), + "listener table should contain {port}" + ); + } + + #[test] + fn managed_job_terminates_its_owned_process() { + use std::os::windows::io::AsRawHandle as _; + use std::os::windows::process::CommandExt as _; + + const CREATE_NO_WINDOW: u32 = 0x08000000; + let mut command = std::process::Command::new("powershell.exe"); + command + .args([ + "-NoLogo", + "-NoProfile", + "-Command", + "Start-Sleep -Seconds 30", + ]) + .creation_flags(CREATE_NO_WINDOW); + let mut child = command.spawn().expect("spawn an isolated test child"); + let job = create_managed_job("test").expect("create a kill-on-close job"); + if let Err(error) = assign_process_to_job(&job, child.as_raw_handle(), "test") { + drop(child.kill()); + drop(child.wait()); + panic!("assign the test child to its job: {error}"); + } + + // SAFETY: only the isolated test child was assigned to this private job. + unsafe { TerminateJobObject(win_handle(&job), 1) }.expect("terminate the private job"); + let deadline = Instant::now() + Duration::from_secs(3); + loop { + if child.try_wait().expect("inspect the test child").is_some() { + break; + } + if Instant::now() >= deadline { + drop(child.kill()); + drop(child.wait()); + panic!("job termination did not stop the test child"); + } + std::thread::sleep(Duration::from_millis(25)); + } + } + + #[tokio::test] + async fn managed_process_shutdown_stops_owned_child() { + let mut process = start_test_process(); + let pid = process.pid(); + assert!(pid > 0, "managed process exposes its pid"); + assert!( + process.try_wait().expect("poll the child").is_none(), + "managed child is running before shutdown" + ); + + process.shutdown(Duration::from_secs(5)).await; + + assert!(!process_is_alive(pid), "shutdown stops the owned child"); + } + + #[test] + fn managed_process_drop_terminates_owned_child() { + let process = start_test_process(); + let pid = process.pid(); + assert!(pid > 0, "managed process exposes its pid"); + + drop(process); + + // Drop terminates synchronously and detaches reaping; wait for the OS + // to report the process gone so the test does not race the cleanup thread. + wait_for_exit(pid); + } + + #[test] + fn managed_process_restart_replaces_without_leaking() { + let mut process = start_test_process(); + let first = process.pid(); + + process.restart().expect("restart the managed child"); + let second = process.pid(); + + assert_ne!(first, second, "restart launches a fresh child"); + assert!(!process_is_alive(first), "restart reaps the previous child"); + assert!( + process.try_wait().expect("poll the replacement").is_none(), + "replacement is running after restart" + ); + + drop(process); + wait_for_exit(second); + } +} diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 4c615f3425..dcfc0fe662 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -15,9 +15,7 @@ use futures::{StreamExt, stream}; use regex_lite::Regex; use serde::Deserialize; #[cfg(windows)] -use std::io::{Read, Write}; -#[cfg(windows)] -use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; +use std::ffi::OsString; #[cfg(windows)] use std::os::windows::process::CommandExt; use std::path::PathBuf; @@ -26,20 +24,9 @@ use std::sync::{LazyLock, OnceLock}; use std::time::Duration; #[cfg(windows)] use std::time::Instant; + #[cfg(windows)] -use windows::Win32::Foundation::{ERROR_INSUFFICIENT_BUFFER, HANDLE}; -#[cfg(windows)] -use windows::Win32::NetworkManagement::IpHelper::{ - GetExtendedTcpTable, MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID, TCP_TABLE_OWNER_PID_LISTENER, -}; -#[cfg(windows)] -use windows::Win32::System::JobObjects::{ - AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, - JOBOBJECT_BASIC_LIMIT_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, - JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject, -}; -#[cfg(windows)] -use windows::core::PCWSTR; +use crate::managed_process::{ManagedProcess, ManagedProcessConfig, ManagedProcessError}; use crate::core::{ FetchContext, NamedRateWindow, Provider, ProviderError, ProviderFetchResult, ProviderId, @@ -56,8 +43,6 @@ const AGY_CLEANUP_RESERVE: Duration = Duration::from_secs(2); const AGY_PROBE_TIMEOUT: Duration = Duration::from_millis(750); #[cfg(windows)] const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); -#[cfg(any(windows, test))] -const AGY_MAX_CURSOR_REPLIES: usize = 32; const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; @@ -67,6 +52,15 @@ const QUOTA_SUMMARY_PATH: &str = #[cfg(windows)] static MANAGED_AGY_FETCH: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); +/// The managed-process owner stays provider-neutral; the provider only maps its +/// error surface into [`ProviderError`]. +#[cfg(windows)] +impl From for ProviderError { + fn from(error: ManagedProcessError) -> Self { + ProviderError::Other(error.to_string()) + } +} + /// Antigravity provider pub struct AntigravityProvider { metadata: ProviderMetadata, @@ -110,21 +104,6 @@ fn is_agy_cli_command(command_line: &str) -> bool { CLI_PATH_RE.is_match(&lower) || AGY_RE.is_match(&lower) } -#[cfg(any(windows, test))] -fn terminal_cursor_position_request_count(tail: &mut Vec, chunk: &[u8]) -> usize { - tail.extend_from_slice(chunk); - let requested = tail.windows(4).filter(|bytes| *bytes == b"\x1b[6n").count(); - if tail.len() > 3 { - tail.drain(..tail.len() - 3); - } - requested -} - -#[cfg(any(windows, test))] -fn terminal_cursor_reply_allowance(sent: usize, requested: usize) -> usize { - requested.min(AGY_MAX_CURSOR_REPLIES.saturating_sub(sent)) -} - impl AntigravityProvider { pub fn new() -> Self { Self { @@ -328,102 +307,12 @@ impl AntigravityProvider { /// Enumerate IPv4 TCP listener ports for a PID through the Windows IP Helper API. /// This avoids starting PowerShell inside the managed readiness poll. /// - /// Known follow-up: only the AF_INET table is enumerated and candidates are - /// probed at `127.0.0.1`, so an IPv6-only loopback listener would be missed. - /// Accepted for now because the managed `agy` service is observed to bind - /// IPv4 on Windows; add AF_INET6 enumeration with `[::1]` probes later. + /// The owner lives in the provider-neutral + /// [`crate::managed_process::listening_ports_for_pid`]; this binding only + /// maps its error into the provider surface. #[cfg(windows)] fn listening_ports_for_pid(pid: u32) -> Result, ProviderError> { - const AF_INET_FAMILY: u32 = 2; - const NO_ERROR: u32 = 0; - - let mut bytes = 0_u32; - // SAFETY: the first call supplies no destination buffer and only asks Windows - // for the required byte count. - let query = unsafe { - GetExtendedTcpTable( - None, - &mut bytes, - false, - AF_INET_FAMILY, - TCP_TABLE_OWNER_PID_LISTENER, - 0, - ) - }; - if query != ERROR_INSUFFICIENT_BUFFER.0 && query != NO_ERROR { - return Err(ProviderError::Other(format!( - "Failed to size the Windows TCP listener table (error {query})" - ))); - } - if bytes < u32::try_from(std::mem::size_of::()).unwrap_or(u32::MAX) { - return Ok(Vec::new()); - } - - let mut buffer = Vec::new(); - let mut loaded = false; - // The table can grow between the sizing call and the read. Retry with - // the updated size instead of failing a refresh on that benign race. - for _ in 0..3 { - // A u32 allocation supplies the alignment required by the all-DWORD MIB rows. - let words = (bytes as usize).div_ceil(std::mem::size_of::()); - buffer.resize(words, 0_u32); - // SAFETY: `buffer` is writable for at least `bytes` bytes and remains alive - // while the returned table is inspected. - let result = unsafe { - GetExtendedTcpTable( - Some(buffer.as_mut_ptr().cast()), - &mut bytes, - false, - AF_INET_FAMILY, - TCP_TABLE_OWNER_PID_LISTENER, - 0, - ) - }; - if result == NO_ERROR { - loaded = true; - break; - } - if result != ERROR_INSUFFICIENT_BUFFER.0 { - return Err(ProviderError::Other(format!( - "Failed to read the Windows TCP listener table (error {result})" - ))); - } - } - if !loaded { - return Err(ProviderError::Other( - "Windows TCP listener table kept changing during the query".to_string(), - )); - } - - let table = buffer.as_ptr().cast::(); - // SAFETY: Windows initialized the header on the successful call above. - let count = unsafe { (*table).dwNumEntries as usize }; - let rows_offset = std::mem::offset_of!(MIB_TCPTABLE_OWNER_PID, table); - let available = (bytes as usize).saturating_sub(rows_offset); - let max_rows = available / std::mem::size_of::(); - if count > max_rows { - return Err(ProviderError::Parse( - "Windows returned an invalid TCP listener table".to_string(), - )); - } - // SAFETY: `count` was bounded by the returned buffer size. Windows lays out - // the fixed-size owner-PID rows consecutively after the table header. - let rows = unsafe { - std::slice::from_raw_parts( - std::ptr::addr_of!((*table).table).cast::(), - count, - ) - }; - let mut ports: Vec = rows - .iter() - .filter(|row| row.dwOwningPid == pid) - .filter_map(|row| u16::try_from(row.dwLocalPort).ok()) - .map(u16::from_be) - .filter(|port| *port != 0) - .collect(); - ports.sort_unstable(); - ports.dedup(); - Ok(ports) + crate::managed_process::listening_ports_for_pid(pid).map_err(ProviderError::from) } /// Non-Windows platforms have no `Get-NetTCPConnection`; return an empty list by design so @@ -590,8 +479,20 @@ impl AntigravityProvider { .redirect(reqwest::redirect::Policy::none()) .build() .map_err(|error| ProviderError::Other(error.to_string()))?; - let mut managed = ManagedAgyProcess::spawn(&binary)?; - let pid = managed.id(); + let config = ManagedProcessConfig { + program: binary, + args: Vec::new(), + env: vec![ + (OsString::from("TERM"), OsString::from("xterm-256color")), + (OsString::from("COLORTERM"), OsString::from("truecolor")), + ], + cwd: dirs::home_dir().filter(|path| path.is_dir()), + pty_rows: 30, + pty_cols: 120, + label: "agy".to_string(), + }; + let mut managed = ManagedProcess::spawn(&config)?; + let pid = managed.pid(); let process_info = ProcessInfo { csrf_token: String::new(), extension_server_csrf_token: None, @@ -650,7 +551,7 @@ impl AntigravityProvider { } .await; - managed.shutdown().await; + managed.shutdown(AGY_CLEANUP_RESERVE).await; result } @@ -711,6 +612,37 @@ impl AntigravityProvider { offline.ok_or(error) } + /// Map a managed-lifecycle outcome onto provider policy. + /// + /// `Reused` means a user-owned runtime answered and the fetch stays local; + /// `Fetched` means the task-owned CLI answered; `Missing` is a policy no-op + /// so the caller can fall back to offline history; an error follows the + /// same offline-preferred resolution as the local probe. This is the policy + /// seam that the lifecycle owner deliberately does not own, so a fake + /// lifecycle can be driven through it in tests. + #[cfg(windows)] + fn resolve_managed_outcome( + outcome: Result, + ) -> Result, ProviderError> { + match outcome { + Ok(ManagedAgyOutcome::Reused(usage)) => Ok(Some(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "local", + ))), + Ok(ManagedAgyOutcome::Fetched(usage)) => Ok(Some(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "cli", + ))), + Ok(ManagedAgyOutcome::Missing) => Ok(None), + Err(error) => { + if !matches!(error, ProviderError::AuthRequired) { + tracing::debug!(%error, "managed Antigravity CLI probe failed"); + } + Self::resolve_probe_failure(error, Self::offline_usage_result()).map(Some) + } + } + } + fn offline_or_unavailable() -> Result { Self::offline_usage_result() .ok_or_else(|| ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.to_string())) @@ -982,33 +914,10 @@ impl Provider for AntigravityProvider { Ok(None) => { #[cfg(windows)] { - match self.fetch_with_managed_agy().await { - Ok(ManagedAgyOutcome::Reused(usage)) => { - return Ok(ProviderFetchResult::new( - Self::with_cadence_labels(usage), - "local", - )); - } - Ok(ManagedAgyOutcome::Fetched(usage)) => { - return Ok(ProviderFetchResult::new( - Self::with_cadence_labels(usage), - "cli", - )); - } - Ok(ManagedAgyOutcome::Missing) => {} - // A signed-out CLI is actionable, so surface it rather - // than hiding it behind offline history. Any other - // managed-start failure still leaves the runtime - // unavailable, so keep the offline-history fallback. - Err(error) => { - if !matches!(error, ProviderError::AuthRequired) { - tracing::debug!(%error, "managed Antigravity CLI probe failed"); - } - return Self::resolve_probe_failure( - error, - Self::offline_usage_result(), - ); - } + if let Some(result) = + Self::resolve_managed_outcome(self.fetch_with_managed_agy().await)? + { + return Ok(result); } } @@ -1069,235 +978,6 @@ enum ManagedAgyOutcome { Missing, } -/// RAII owner for the exact `agy` process started by this provider. Dropping it -/// cannot affect Antigravity or CLI processes that were already running. -#[cfg(windows)] -struct ManagedAgyProcess { - child: Option>, - pid: u32, - job: Option, - master: Option>, - drain_thread: Option>, -} - -#[cfg(windows)] -impl ManagedAgyProcess { - fn spawn(binary: &std::path::Path) -> Result { - let job = create_managed_agy_job()?; - let pty_system = portable_pty::native_pty_system(); - let pair = pty_system - .openpty(portable_pty::PtySize { - rows: 30, - cols: 120, - pixel_width: 0, - pixel_height: 0, - }) - .map_err(|error| { - ProviderError::Other(format!("Failed to create a terminal for agy: {error}")) - })?; - let mut command = portable_pty::CommandBuilder::new(binary.as_os_str()); - if let Some(home) = dirs::home_dir().filter(|path| path.is_dir()) { - command.cwd(home.as_os_str()); - } - command.env("TERM", "xterm-256color"); - command.env("COLORTERM", "truecolor"); - - let mut reader = pair.master.try_clone_reader().map_err(|error| { - ProviderError::Other(format!("Failed to read the agy terminal: {error}")) - })?; - let mut writer = pair.master.take_writer().map_err(|error| { - ProviderError::Other(format!("Failed to open the agy terminal: {error}")) - })?; - let mut child = pair.slave.spawn_command(command).map_err(|error| { - ProviderError::Other(format!("Failed to launch the agy CLI: {error}")) - })?; - drop(pair.slave); - - let Some(pid) = child.process_id() else { - drop(child.kill()); - drop(child.wait()); - return Err(ProviderError::Other( - "Failed to determine the managed agy process id".to_string(), - )); - }; - let Some(process_handle) = child.as_raw_handle() else { - drop(child.kill()); - drop(child.wait()); - return Err(ProviderError::Other( - "Failed to access the managed agy process handle".to_string(), - )); - }; - if let Err(error) = assign_process_to_job(&job, process_handle) { - drop(child.kill()); - drop(child.wait()); - return Err(error); - } - let drain_thread = std::thread::spawn(move || { - // Drain without logging: terminal output can contain account data. - // Windows ConPTY programs may request the cursor position and wait - // for a terminal response before continuing initialization. - let mut buffer = [0_u8; 4096]; - let mut tail = Vec::with_capacity(3); - let mut cursor_replies = 0_usize; - while let Ok(read) = reader.read(&mut buffer) { - if read == 0 { - break; - } - let requested = terminal_cursor_position_request_count(&mut tail, &buffer[..read]); - let allowed = terminal_cursor_reply_allowance(cursor_replies, requested); - for _ in 0..allowed { - drop(writer.write_all(b"\x1b[1;1R")); - } - if allowed > 0 { - drop(writer.flush()); - cursor_replies += allowed; - } - } - }); - Ok(Self { - child: Some(child), - pid, - job: Some(job), - master: Some(pair.master), - drain_thread: Some(drain_thread), - }) - } - - fn id(&self) -> u32 { - self.pid - } - - fn try_wait(&mut self) -> Result, ProviderError> { - self.child - .as_mut() - .expect("managed agy child is present until cleanup") - .try_wait() - .map_err(|error| { - ProviderError::Other(format!("Failed to inspect the agy CLI: {error}")) - }) - } - - async fn shutdown(mut self) { - let Some(resources) = self.take_resources() else { - return; - }; - let cleanup = tokio::task::spawn_blocking(move || resources.terminate_and_reap()); - // A stuck platform wait must not hold the async provider worker. The - // blocking cleanup task remains detached and still owns every handle. - drop(tokio::time::timeout(AGY_CLEANUP_RESERVE, cleanup).await); - } - - fn take_resources(&mut self) -> Option { - Some(ManagedAgyResources { - child: self.child.take()?, - job: Some( - self.job - .take() - .expect("managed agy job is present until cleanup"), - ), - master: self.master.take(), - drain_thread: self.drain_thread.take(), - }) - } -} - -#[cfg(windows)] -impl Drop for ManagedAgyProcess { - fn drop(&mut self) { - let Some(mut resources) = self.take_resources() else { - return; - }; - resources.terminate(); - // Drop can run when an outer timeout cancels the fetch. Reaping and - // joining the terminal drain must therefore never block that worker. - drop( - std::thread::Builder::new() - .name("codexbar-agy-cleanup".to_string()) - .spawn(move || resources.reap()), - ); - } -} - -#[cfg(windows)] -struct ManagedAgyResources { - child: Box, - job: Option, - master: Option>, - drain_thread: Option>, -} - -#[cfg(windows)] -impl ManagedAgyResources { - fn terminate(&mut self) { - // SAFETY: this job is private to the single process launched above; - // user-owned Antigravity and agy processes were never assigned to it. - let terminated = self - .job - .as_ref() - .is_some_and(|job| unsafe { TerminateJobObject(win_handle(job), 1) }.is_ok()); - // KILL_ON_JOB_CLOSE is the second termination path if the explicit API - // fails. Close it before wait so a failure cannot strand the reaper. - drop(self.job.take()); - if !terminated { - drop(self.child.kill()); - } - } - - fn reap(mut self) { - drop(self.child.wait()); - drop(self.master.take()); - if let Some(thread) = self.drain_thread.take() { - drop(thread.join()); - } - } - - fn terminate_and_reap(mut self) { - self.terminate(); - self.reap(); - } -} - -#[cfg(windows)] -fn create_managed_agy_job() -> Result { - // SAFETY: a successful call transfers a unique job handle to this owner. - let raw = unsafe { CreateJobObjectW(None, PCWSTR::null()) } - .map_err(|error| ProviderError::Other(format!("Failed to create agy job: {error}")))?; - // SAFETY: `raw` is a unique valid handle returned by CreateJobObjectW. - let job = unsafe { OwnedHandle::from_raw_handle(raw.0) }; - let limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION { - BasicLimitInformation: JOBOBJECT_BASIC_LIMIT_INFORMATION { - LimitFlags: JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, - ..Default::default() - }, - ..Default::default() - }; - let size = u32::try_from(std::mem::size_of_val(&limits)) - .map_err(|error| ProviderError::Other(format!("Invalid agy job limit size: {error}")))?; - // SAFETY: `job` is valid and `limits` is initialized for the requested class. - unsafe { - SetInformationJobObject( - win_handle(&job), - JobObjectExtendedLimitInformation, - (&limits as *const JOBOBJECT_EXTENDED_LIMIT_INFORMATION).cast(), - size, - ) - } - .map_err(|error| ProviderError::Other(format!("Failed to configure agy job: {error}")))?; - Ok(job) -} - -#[cfg(windows)] -fn assign_process_to_job(job: &OwnedHandle, process: RawHandle) -> Result<(), ProviderError> { - // SAFETY: both handles are valid and remain owned by their respective wrappers. - unsafe { AssignProcessToJobObject(win_handle(job), HANDLE(process)) } - .map_err(|error| ProviderError::Other(format!("Failed to contain agy process: {error}"))) -} - -#[cfg(windows)] -fn win_handle(value: &OwnedHandle) -> HANDLE { - HANDLE(value.as_raw_handle()) -} - // API Response types #[derive(Debug, Deserialize)] diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index 08b6abf63b..149120b8f7 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -256,89 +256,46 @@ fn managed_agy_candidates_prefer_override_then_path_then_known_installs() { ); } -#[test] -fn managed_agy_terminal_detects_cursor_request_across_reads() { - let mut tail = Vec::new(); +// ── Managed lifecycle policy (fake outcomes) ─────────────────────── +// +// The process lifecycle itself is covered by `crate::managed_process`; these +// exercise the provider-side policy that maps a lifecycle outcome onto a fetch +// result without spawning a real `agy`. - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"ready\x1b["), - 0 - ); - assert_eq!(terminal_cursor_position_request_count(&mut tail, b"6n"), 1); - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"plain output"), - 0 - ); - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"\x1b[6nmore\x1b[6n"), - 2 - ); +#[cfg(windows)] +#[test] +fn reused_user_runtime_stays_local() { + let usage = UsageSnapshot::new(RateWindow::new(10.0)); + let result = AntigravityProvider::resolve_managed_outcome(Ok(ManagedAgyOutcome::Reused(usage))) + .expect("reused outcome resolves") + .expect("reused outcome yields usage"); + assert_eq!(result.source_label, "local"); } +#[cfg(windows)] #[test] -fn managed_agy_terminal_caps_cursor_replies() { - assert_eq!(terminal_cursor_reply_allowance(0, 2), 2); - assert_eq!(terminal_cursor_reply_allowance(31, 4), 1); - assert_eq!( - terminal_cursor_reply_allowance(AGY_MAX_CURSOR_REPLIES, 1), - 0 - ); +fn owned_cli_fetch_reports_cli_source() { + let usage = UsageSnapshot::new(RateWindow::new(10.0)); + let result = + AntigravityProvider::resolve_managed_outcome(Ok(ManagedAgyOutcome::Fetched(usage))) + .expect("owned outcome resolves") + .expect("owned outcome yields usage"); + assert_eq!(result.source_label, "cli"); } #[cfg(windows)] #[test] -fn windows_listener_table_finds_current_process_port() { - let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) - .expect("bind a local IPv4 listener"); - let port = listener.local_addr().expect("listener address").port(); - - let ports = AntigravityProvider::listening_ports_for_pid(std::process::id()) - .expect("read the Windows TCP listener table"); - - assert!( - ports.contains(&port), - "listener table should contain {port}" - ); +fn missing_runtime_is_a_policy_no_op() { + let result = AntigravityProvider::resolve_managed_outcome(Ok(ManagedAgyOutcome::Missing)) + .expect("a missing runtime is not an error"); + assert!(result.is_none(), "missing runtime falls through to offline"); } #[cfg(windows)] #[test] -fn managed_agy_job_terminates_its_owned_process() { - use std::os::windows::io::AsRawHandle as _; - use std::os::windows::process::CommandExt as _; - - const CREATE_NO_WINDOW: u32 = 0x08000000; - let mut command = std::process::Command::new("powershell.exe"); - command - .args([ - "-NoLogo", - "-NoProfile", - "-Command", - "Start-Sleep -Seconds 30", - ]) - .creation_flags(CREATE_NO_WINDOW); - let mut child = command.spawn().expect("spawn an isolated test child"); - let job = create_managed_agy_job().expect("create a kill-on-close job"); - if let Err(error) = assign_process_to_job(&job, child.as_raw_handle()) { - drop(child.kill()); - drop(child.wait()); - panic!("assign the test child to its job: {error}"); - } - - // SAFETY: only the isolated test child was assigned to this private job. - unsafe { TerminateJobObject(win_handle(&job), 1) }.expect("terminate the private job"); - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(3); - loop { - if child.try_wait().expect("inspect the test child").is_some() { - break; - } - if std::time::Instant::now() >= deadline { - drop(child.kill()); - drop(child.wait()); - panic!("job termination did not stop the test child"); - } - std::thread::sleep(std::time::Duration::from_millis(25)); - } +fn managed_auth_required_surfaces_instead_of_offline() { + let result = AntigravityProvider::resolve_managed_outcome(Err(ProviderError::AuthRequired)); + assert!(matches!(result, Err(ProviderError::AuthRequired))); } // ── agy CLI process matching ─────────────────────────────────────── From 5583da211f195994a0e15b683a5a958342a57f4a Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:27:31 +0700 Subject: [PATCH 05/10] fix(antigravity): contain managed process at creation --- rust/Cargo.toml | 3 + rust/src/managed_process.rs | 557 +++++++++++++++++++++++++++++++----- 2 files changed, 487 insertions(+), 73 deletions(-) diff --git a/rust/Cargo.toml b/rust/Cargo.toml index c0b643ccd6..2f04469077 100755 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -90,8 +90,11 @@ windows = { version = "0.58", features = [ "Win32_Security_Cryptography", "Win32_Media_Audio", "Win32_NetworkManagement_IpHelper", + "Win32_Security", + "Win32_System_Console", "Win32_System_LibraryLoader", "Win32_System_JobObjects", + "Win32_System_Pipes", "Win32_System_Threading", "Win32_UI_WindowsAndMessaging", "Win32_Storage_FileSystem", diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index adde854936..6f4e33a64f 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -7,27 +7,58 @@ //! introspection, restart, shutdown and Drop cleanup. It never adopts or //! terminates a process it did not create, so user-owned processes are always //! isolated from cleanup. +//! +//! The child joins its kill-on-close job through the process-creation +//! `PROC_THREAD_ATTRIBUTE_JOB_LIST` attribute, so containment is atomic: no +//! thread of the child (or a descendant it starts) can run before the kernel +//! has bound the process tree to the job. The PTY is attached with the +//! `PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE` attribute in the same creation call, +//! preserving the previous ConPTY behavior. -use std::ffi::OsString; +use std::ffi::{OsStr, OsString}; +use std::fs::File; use std::io::{Read, Write}; +use std::os::windows::ffi::OsStrExt; use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::time::Duration; -use windows::Win32::Foundation::{ERROR_INSUFFICIENT_BUFFER, HANDLE}; +use windows::Win32::Foundation::{HANDLE, INVALID_HANDLE_VALUE}; use windows::Win32::NetworkManagement::IpHelper::{ GetExtendedTcpTable, MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID, TCP_TABLE_OWNER_PID_LISTENER, }; +use windows::Win32::System::Console::{ + ClosePseudoConsole, CreatePseudoConsole, HPCON, PSEUDOCONSOLE_INHERIT_CURSOR, +}; +#[cfg(test)] +use windows::Win32::System::JobObjects::AssignProcessToJobObject; use windows::Win32::System::JobObjects::{ - AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, - JOBOBJECT_BASIC_LIMIT_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, - JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject, + CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOBOBJECT_BASIC_LIMIT_INFORMATION, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, TerminateJobObject, }; -use windows::core::PCWSTR; +use windows::Win32::System::Pipes::CreatePipe; +use windows::Win32::System::Threading::{ + CREATE_UNICODE_ENVIRONMENT, CreateProcessW, DeleteProcThreadAttributeList, + EXTENDED_STARTUPINFO_PRESENT, GetExitCodeProcess, GetProcessId, INFINITE, + InitializeProcThreadAttributeList, LPPROC_THREAD_ATTRIBUTE_LIST, + PROC_THREAD_ATTRIBUTE_JOB_LIST, PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE, PROCESS_INFORMATION, + STARTF_USESTDHANDLES, STARTUPINFOEXW, STARTUPINFOW, TerminateProcess, + UpdateProcThreadAttribute, WaitForSingleObject, +}; +use windows::core::{PCWSTR, PWSTR}; /// Maximum number of terminal cursor-position replies sent to one child. const MAX_CURSOR_REPLIES: usize = 32; +/// `GetExitCodeProcess` value reported while the process is still running. +const STILL_ACTIVE: u32 = 259; + +/// ConPTY flags portable-pty applies to preserve interactive startup behavior. +/// The `windows` crate only names `PSEUDOCONSOLE_INHERIT_CURSOR`. +const PSEUDOCONSOLE_RESIZE_QUIRK: u32 = 0x2; +const PSEUDOCONSOLE_WIN32_INPUT_MODE: u32 = 0x4; + /// Provider-supplied configuration for one managed child process. /// /// The provider owns this policy: which executable, arguments and environment @@ -70,90 +101,48 @@ pub type ManagedProcessResult = Result; /// any process that was already running. pub struct ManagedProcess { config: ManagedProcessConfig, - child: Option>, + child: Option, pid: u32, job: Option, - master: Option>, + master: Option, drain_thread: Option>, } impl ManagedProcess { /// Start `config.program` in a PTY inside its own kill-on-close job. + /// + /// The job and the pseudoconsole are both supplied as process-creation + /// attributes, so the child is contained before any of its threads execute + /// and cannot spawn an uncontained descendant during startup. pub fn spawn(config: &ManagedProcessConfig) -> ManagedProcessResult { let job = create_managed_job(&config.label)?; - let pty_system = portable_pty::native_pty_system(); - let pair = pty_system - .openpty(portable_pty::PtySize { - rows: config.pty_rows, - cols: config.pty_cols, - pixel_width: 0, - pixel_height: 0, - }) - .map_err(|error| { - ManagedProcessError::new(format!( - "Failed to create a terminal for {}: {error}", - config.label - )) - })?; - let mut command = portable_pty::CommandBuilder::new(config.program.as_os_str()); - for arg in &config.args { - command.arg(arg); - } - if let Some(cwd) = &config.cwd { - command.cwd(cwd.as_os_str()); - } - for (key, value) in &config.env { - command.env(key, value); + let (mut pty, child) = spawn_pty_child(config, &job)?; + let pid = child.process_id(); + if pid == 0 { + // The job owns the child now; dropping it closes the job handle and + // its KILL_ON_JOB_CLOSE limit terminates the untracked process. + return Err(ManagedProcessError::new(format!( + "Failed to determine the managed {} process id", + config.label + ))); } - let reader = pair.master.try_clone_reader().map_err(|error| { + let reader = pty.try_clone_reader().map_err(|error| { ManagedProcessError::new(format!( "Failed to read the {} terminal: {error}", config.label )) })?; - let writer = pair.master.take_writer().map_err(|error| { - ManagedProcessError::new(format!( - "Failed to open the {} terminal: {error}", - config.label - )) - })?; - let mut child = pair.slave.spawn_command(command).map_err(|error| { - ManagedProcessError::new(format!( - "Failed to launch the {} CLI: {error}", - config.label - )) + let writer = pty.take_writer().ok_or_else(|| { + ManagedProcessError::new(format!("Failed to open the {} terminal", config.label)) })?; - drop(pair.slave); - - let Some(pid) = child.process_id() else { - drop(child.kill()); - drop(child.wait()); - return Err(ManagedProcessError::new(format!( - "Failed to determine the managed {} process id", - config.label - ))); - }; - let Some(process_handle) = child.as_raw_handle() else { - drop(child.kill()); - drop(child.wait()); - return Err(ManagedProcessError::new(format!( - "Failed to access the managed {} process handle", - config.label - ))); - }; - if let Err(error) = assign_process_to_job(&job, process_handle, &config.label) { - drop(child.kill()); - drop(child.wait()); - return Err(error); - } let drain_thread = spawn_drain_thread(reader, writer); Ok(Self { config: config.clone(), child: Some(child), pid, job: Some(job), - master: Some(pair.master), + master: Some(pty), drain_thread: Some(drain_thread), }) } @@ -169,6 +158,7 @@ impl ManagedProcess { .as_mut() .expect("managed child is present until cleanup") .try_wait() + .map(|status| status.map(portable_pty::ExitStatus::with_exit_code)) .map_err(|error| { ManagedProcessError::new(format!( "Failed to inspect the {} CLI: {error}", @@ -237,9 +227,9 @@ impl Drop for ManagedProcess { } struct ManagedProcessResources { - child: Box, + child: ManagedChild, job: Option, - master: Option>, + master: Option, drain_thread: Option>, } @@ -255,7 +245,7 @@ impl ManagedProcessResources { // fails. Close it before wait so a failure cannot strand the reaper. drop(self.job.take()); if !terminated { - drop(self.child.kill()); + self.child.kill(); } } @@ -273,6 +263,207 @@ impl ManagedProcessResources { } } +/// The exact process created for this owner. It holds the only process handle +/// that cleanup waits on; containment comes from the job bound at creation. +struct ManagedChild { + process: OwnedHandle, +} + +impl ManagedChild { + fn process_id(&self) -> u32 { + // SAFETY: the process handle remains owned and valid for `self`. + unsafe { GetProcessId(win_handle(&self.process)) } + } + + fn try_wait(&mut self) -> std::io::Result> { + let mut code = 0_u32; + // SAFETY: the process handle and output pointer are valid. + unsafe { GetExitCodeProcess(win_handle(&self.process), &mut code) } + .map_err(std::io::Error::other)?; + Ok((code != STILL_ACTIVE).then_some(code)) + } + + fn wait(&mut self) -> std::io::Result { + // SAFETY: the process handle remains owned and valid for `self`. + unsafe { + WaitForSingleObject(win_handle(&self.process), INFINITE); + } + let mut code = 0_u32; + // SAFETY: the process handle and output pointer are valid. + unsafe { GetExitCodeProcess(win_handle(&self.process), &mut code) } + .map_err(std::io::Error::other)?; + Ok(code) + } + + fn kill(&mut self) { + // SAFETY: the process handle remains owned and valid for `self`. + drop(unsafe { TerminateProcess(win_handle(&self.process), 1) }); + } +} + +/// The host side of a ConPTY. Dropping it closes the pseudoconsole, which makes +/// the drain reader observe EOF and lets the drain thread finish. +struct PseudoConsole { + con: HPCON, + input: Option, + output: OwnedHandle, +} + +impl PseudoConsole { + fn try_clone_reader(&self) -> std::io::Result> { + let handle = self.output.try_clone()?; + Ok(Box::new(File::from(handle))) + } + + fn take_writer(&mut self) -> Option> { + self.input + .take() + .map(|handle| Box::new(File::from(handle)) as Box) + } +} + +impl Drop for PseudoConsole { + fn drop(&mut self) { + if !self.con.is_invalid() { + // SAFETY: this pseudoconsole was created here and is closed once. + unsafe { ClosePseudoConsole(self.con) }; + } + } +} + +/// Create the PTY and the owned child in one process-creation call, binding the +/// child to `job` and attaching the pseudoconsole atomically. +fn spawn_pty_child( + config: &ManagedProcessConfig, + job: &OwnedHandle, +) -> ManagedProcessResult<(PseudoConsole, ManagedChild)> { + let cols = i16::try_from(config.pty_cols).map_err(|error| { + ManagedProcessError::new(format!("Invalid {} terminal width: {error}", config.label)) + })?; + let rows = i16::try_from(config.pty_rows).map_err(|error| { + ManagedProcessError::new(format!("Invalid {} terminal height: {error}", config.label)) + })?; + + let (input_read, input_write) = create_pipe(&config.label)?; + let (output_read, output_write) = create_pipe(&config.label)?; + // SAFETY: both handles are valid and the pseudoconsole duplicates them. + let con = unsafe { + CreatePseudoConsole( + windows::Win32::System::Console::COORD { X: cols, Y: rows }, + win_handle(&input_read), + win_handle(&output_write), + PSEUDOCONSOLE_INHERIT_CURSOR + | PSEUDOCONSOLE_RESIZE_QUIRK + | PSEUDOCONSOLE_WIN32_INPUT_MODE, + ) + } + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to create a terminal for {}: {error}", + config.label + )) + })?; + // The pseudoconsole owns duplicates of the child-side pipe ends. + drop(input_read); + drop(output_write); + let pty = PseudoConsole { + con, + input: Some(input_write), + output: output_read, + }; + + let mut cmdline = build_command_line(&config.program, &config.args)?; + let env_block = build_environment_block(&config.env)?; + let cwd = config + .cwd + .as_ref() + .map(|dir| encode_wide_nul(dir.as_os_str())) + .transpose()?; + + let mut attributes = Attributes::new(2)?; + let jobs = [win_handle(job)]; + // SAFETY: `jobs` and `con` stay alive through CreateProcessW. The job-list + // attribute binds the child before its first thread can run; the + // pseudoconsole attribute wires its stdio to the PTY above. + unsafe { + UpdateProcThreadAttribute( + attributes.ptr(), + 0, + PROC_THREAD_ATTRIBUTE_JOB_LIST as usize, + Some(jobs.as_ptr().cast()), + std::mem::size_of_val(&jobs), + None, + None, + ) + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to contain {} process: {error}", + config.label + )) + })?; + UpdateProcThreadAttribute( + attributes.ptr(), + 0, + PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE as usize, + Some(con.0 as *const core::ffi::c_void), + std::mem::size_of::(), + None, + None, + ) + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to attach the {} terminal: {error}", + config.label + )) + })?; + } + + let startup = STARTUPINFOEXW { + StartupInfo: STARTUPINFOW { + cb: u32::try_from(std::mem::size_of::()).map_err(|error| { + ManagedProcessError::new(format!("Invalid {} startup size: {error}", config.label)) + })?, + dwFlags: STARTF_USESTDHANDLES, + // The pseudoconsole owns the child's stdio; invalid handles stop the + // child from inheriting this process's redirected handles. + hStdInput: INVALID_HANDLE_VALUE, + hStdOutput: INVALID_HANDLE_VALUE, + hStdError: INVALID_HANDLE_VALUE, + ..Default::default() + }, + lpAttributeList: attributes.ptr(), + }; + let mut info = PROCESS_INFORMATION::default(); + // SAFETY: all buffers (command line, environment, cwd, attribute list, + // startup info) outlive this call; the kernel fills `info` on success. + unsafe { + CreateProcessW( + PCWSTR::null(), + PWSTR(cmdline.as_mut_ptr()), + None, + None, + false, + EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT, + Some(env_block.as_ptr().cast()), + cwd.as_ref() + .map_or(PCWSTR::null(), |value| PCWSTR(value.as_ptr())), + &startup.StartupInfo, + &mut info, + ) + } + .map_err(|error| { + ManagedProcessError::new(format!( + "Failed to launch the {} CLI: {error}", + config.label + )) + })?; + // SAFETY: CreateProcessW returned a valid thread handle closed exactly once. + drop(unsafe { OwnedHandle::from_raw_handle(info.hThread.0) }); + // SAFETY: CreateProcessW returned a valid process handle owned by the child. + let process = unsafe { OwnedHandle::from_raw_handle(info.hProcess.0) }; + Ok((pty, ManagedChild { process })) +} + /// Drain PTY output without logging it: terminal output can contain account /// data. Windows ConPTY programs may request the cursor position and wait for a /// terminal response before continuing initialization, so answer a bounded @@ -347,6 +538,8 @@ fn create_managed_job(label: &str) -> ManagedProcessResult { Ok(job) } +/// Post-creation assignment retained only for the job-termination unit test. +#[cfg(test)] fn assign_process_to_job( job: &OwnedHandle, process: RawHandle, @@ -358,6 +551,151 @@ fn assign_process_to_job( }) } +fn create_pipe(label: &str) -> ManagedProcessResult<(OwnedHandle, OwnedHandle)> { + let mut read = HANDLE::default(); + let mut write = HANDLE::default(); + // SAFETY: CreatePipe writes both read/write handles on success. + unsafe { CreatePipe(&mut read, &mut write, None, 0) }.map_err(|error| { + ManagedProcessError::new(format!("Failed to create a {label} terminal pipe: {error}")) + })?; + // SAFETY: CreatePipe returned two unique valid, non-inheritable handles. + let read = unsafe { OwnedHandle::from_raw_handle(read.0) }; + // SAFETY: CreatePipe returned two unique valid, non-inheritable handles. + let write = unsafe { OwnedHandle::from_raw_handle(write.0) }; + Ok((read, write)) +} + +/// Owned `PROC_THREAD_ATTRIBUTE_LIST` storage for one process-creation call. +struct Attributes(Vec); + +impl Attributes { + fn new(count: u32) -> ManagedProcessResult { + let mut bytes = 0_usize; + // SAFETY: the sizing call writes only to `bytes` and is expected to + // report insufficient buffer; the result is intentionally ignored. + drop(unsafe { + InitializeProcThreadAttributeList( + LPPROC_THREAD_ATTRIBUTE_LIST::default(), + count, + 0, + &mut bytes, + ) + }); + if bytes == 0 { + return Err(ManagedProcessError::new( + "Failed to size a process attribute list".to_string(), + )); + } + let mut value = Self(vec![0; bytes.div_ceil(std::mem::size_of::())]); + // SAFETY: the owned allocation is aligned for and at least `bytes` long. + unsafe { InitializeProcThreadAttributeList(value.ptr(), count, 0, &mut bytes) }.map_err( + |error| { + ManagedProcessError::new(format!( + "Failed to allocate a process attribute list: {error}" + )) + }, + )?; + Ok(value) + } + + fn ptr(&mut self) -> LPPROC_THREAD_ATTRIBUTE_LIST { + LPPROC_THREAD_ATTRIBUTE_LIST(self.0.as_mut_ptr().cast()) + } +} + +impl Drop for Attributes { + fn drop(&mut self) { + if !self.0.is_empty() { + // SAFETY: this list was initialized and has not yet been deleted. + unsafe { DeleteProcThreadAttributeList(self.ptr()) }; + } + } +} + +fn build_command_line(program: &Path, args: &[OsString]) -> ManagedProcessResult> { + let mut cmdline = Vec::new(); + append_quoted(program.as_os_str(), &mut cmdline)?; + for arg in args { + cmdline.push(b' ' as u16); + append_quoted(arg, &mut cmdline)?; + } + cmdline.push(0); + Ok(cmdline) +} + +fn build_environment_block(overrides: &[(OsString, OsString)]) -> ManagedProcessResult> { + let mut values: Vec<(OsString, OsString)> = std::env::vars_os().collect(); + for (key, value) in overrides { + values.retain(|(existing, _)| { + !existing + .to_string_lossy() + .eq_ignore_ascii_case(&key.to_string_lossy()) + }); + values.push((key.clone(), value.clone())); + } + // CreateProcessW expects Unicode environment blocks sorted case-insensitively. + values.sort_by_cached_key(|(key, _)| key.to_string_lossy().to_uppercase()); + let mut block = Vec::new(); + for (key, value) in values { + let mut entry = key; + entry.push("="); + entry.push(value); + block.extend(encode_wide(&entry)?); + block.push(0); + } + block.push(0); + Ok(block) +} + +fn encode_wide(value: &OsStr) -> ManagedProcessResult> { + let wide: Vec = value.encode_wide().collect(); + if wide.contains(&0) { + return Err(ManagedProcessError::new( + "Process argument contains an embedded NUL".to_string(), + )); + } + Ok(wide) +} + +fn encode_wide_nul(value: &OsStr) -> ManagedProcessResult> { + let mut wide = encode_wide(value)?; + wide.push(0); + Ok(wide) +} + +/// Quote one argument using the MSDN command-line rules. +fn append_quoted(value: &OsStr, output: &mut Vec) -> ManagedProcessResult<()> { + let value = encode_wide(value)?; + let needs_quotes = value.is_empty() + || value + .iter() + .any(|code| matches!(*code, 0x20 | 0x09 | 0x0a | 0x0b | 0x22)); + if !needs_quotes { + output.extend(value); + return Ok(()); + } + + output.push(b'"' as u16); + let mut backslashes = 0_usize; + for code in value { + if code == b'\\' as u16 { + backslashes += 1; + continue; + } + let trailing = if code == b'"' as u16 { + backslashes * 2 + 1 + } else { + backslashes + }; + output.extend(std::iter::repeat_n(b'\\' as u16, trailing)); + output.push(code); + backslashes = 0; + } + output.extend(std::iter::repeat_n(b'\\' as u16, backslashes * 2)); + output.push(b'"' as u16); + Ok(()) +} + fn win_handle(value: &OwnedHandle) -> HANDLE { HANDLE(value.as_raw_handle()) } @@ -386,7 +724,7 @@ pub fn listening_ports_for_pid(pid: u32) -> ManagedProcessResult> { 0, ) }; - if query != ERROR_INSUFFICIENT_BUFFER.0 && query != NO_ERROR { + if query != windows::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER.0 && query != NO_ERROR { return Err(ManagedProcessError::new(format!( "Failed to size the Windows TCP listener table (error {query})" ))); @@ -419,7 +757,7 @@ pub fn listening_ports_for_pid(pid: u32) -> ManagedProcessResult> { loaded = true; break; } - if result != ERROR_INSUFFICIENT_BUFFER.0 { + if result != windows::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER.0 { return Err(ManagedProcessError::new(format!( "Failed to read the Windows TCP listener table (error {result})" ))); @@ -645,4 +983,77 @@ mod tests { drop(process); wait_for_exit(second); } + + #[tokio::test] + async fn managed_process_contains_descendants_of_the_pty_child() { + let marker = std::env::temp_dir().join(format!( + "codexbar-managed-descendant-{}.txt", + std::process::id() + )); + drop(std::fs::remove_file(&marker)); + let script = format!( + "$psi = New-Object System.Diagnostics.ProcessStartInfo; \ + $psi.FileName = 'powershell.exe'; \ + $psi.Arguments = '-NoLogo -NoProfile -Command \"Start-Sleep -Seconds 120\"'; \ + $psi.UseShellExecute = $false; \ + $p = [System.Diagnostics.Process]::Start($psi); \ + Set-Content -LiteralPath '{}' -Value $p.Id; \ + Start-Sleep -Seconds 120", + marker.display() + ); + let config = ManagedProcessConfig { + program: PathBuf::from("powershell.exe"), + args: vec![ + OsString::from("-NoLogo"), + OsString::from("-NoProfile"), + OsString::from("-NonInteractive"), + OsString::from("-Command"), + OsString::from(script), + ], + env: Vec::new(), + cwd: None, + pty_rows: 30, + pty_cols: 120, + label: "test-descendant".to_string(), + }; + let mut process = ManagedProcess::spawn(&config).expect("start a managed test process"); + + let descendant = wait_for_descendant_pid(&marker); + assert!( + process_is_alive(descendant), + "the descendant of the PTY child should be running before shutdown" + ); + + process.shutdown(Duration::from_secs(10)).await; + + wait_for_exit_within(descendant, Duration::from_secs(10)); + drop(std::fs::remove_file(&marker)); + } + + fn wait_for_descendant_pid(marker: &Path) -> u32 { + let deadline = Instant::now() + Duration::from_secs(15); + loop { + if let Ok(contents) = std::fs::read_to_string(marker) + && let Ok(pid) = contents.trim().parse::() + { + return pid; + } + assert!( + Instant::now() < deadline, + "the PTY child did not report a descendant pid" + ); + std::thread::sleep(Duration::from_millis(50)); + } + } + + fn wait_for_exit_within(pid: u32, budget: Duration) { + let deadline = Instant::now() + budget; + while process_is_alive(pid) { + assert!( + Instant::now() < deadline, + "descendant {pid} survived the managed job cleanup" + ); + std::thread::sleep(Duration::from_millis(25)); + } + } } From 0db5ed891e0d99bb459a2136e0982eb809d545c3 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:46:24 +0700 Subject: [PATCH 06/10] fix(antigravity): distinguish exited status 259 --- rust/src/managed_process.rs | 59 ++++++++++++++++++++++++++++++------- 1 file changed, 48 insertions(+), 11 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index 6f4e33a64f..22bc8fb05a 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -23,7 +23,7 @@ use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; use std::path::{Path, PathBuf}; use std::time::Duration; -use windows::Win32::Foundation::{HANDLE, INVALID_HANDLE_VALUE}; +use windows::Win32::Foundation::{HANDLE, INVALID_HANDLE_VALUE, WAIT_OBJECT_0, WAIT_TIMEOUT}; use windows::Win32::NetworkManagement::IpHelper::{ GetExtendedTcpTable, MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID, TCP_TABLE_OWNER_PID_LISTENER, }; @@ -51,9 +51,6 @@ use windows::core::{PCWSTR, PWSTR}; /// Maximum number of terminal cursor-position replies sent to one child. const MAX_CURSOR_REPLIES: usize = 32; -/// `GetExitCodeProcess` value reported while the process is still running. -const STILL_ACTIVE: u32 = 259; - /// ConPTY flags portable-pty applies to preserve interactive startup behavior. /// The `windows` crate only names `PSEUDOCONSOLE_INHERIT_CURSOR`. const PSEUDOCONSOLE_RESIZE_QUIRK: u32 = 0x2; @@ -276,17 +273,29 @@ impl ManagedChild { } fn try_wait(&mut self) -> std::io::Result> { + // `STILL_ACTIVE` is also a legal exit code. Use the process handle's + // signaled state as the liveness source, then read the code only after + // Windows confirms that the process has exited. + // SAFETY: the process handle remains owned and valid for `self`. + let state = unsafe { WaitForSingleObject(win_handle(&self.process), 0) }; + if state == WAIT_TIMEOUT { + return Ok(None); + } + if state != WAIT_OBJECT_0 { + return Err(std::io::Error::last_os_error()); + } let mut code = 0_u32; // SAFETY: the process handle and output pointer are valid. unsafe { GetExitCodeProcess(win_handle(&self.process), &mut code) } .map_err(std::io::Error::other)?; - Ok((code != STILL_ACTIVE).then_some(code)) + Ok(Some(code)) } fn wait(&mut self) -> std::io::Result { // SAFETY: the process handle remains owned and valid for `self`. - unsafe { - WaitForSingleObject(win_handle(&self.process), INFINITE); + let state = unsafe { WaitForSingleObject(win_handle(&self.process), INFINITE) }; + if state != WAIT_OBJECT_0 { + return Err(std::io::Error::last_os_error()); } let mut code = 0_u32; // SAFETY: the process handle and output pointer are valid. @@ -566,7 +575,10 @@ fn create_pipe(label: &str) -> ManagedProcessResult<(OwnedHandle, OwnedHandle)> } /// Owned `PROC_THREAD_ATTRIBUTE_LIST` storage for one process-creation call. -struct Attributes(Vec); +struct Attributes { + storage: Vec, + initialized: bool, +} impl Attributes { fn new(count: u32) -> ManagedProcessResult { @@ -586,7 +598,10 @@ impl Attributes { "Failed to size a process attribute list".to_string(), )); } - let mut value = Self(vec![0; bytes.div_ceil(std::mem::size_of::())]); + let mut value = Self { + storage: vec![0; bytes.div_ceil(std::mem::size_of::())], + initialized: false, + }; // SAFETY: the owned allocation is aligned for and at least `bytes` long. unsafe { InitializeProcThreadAttributeList(value.ptr(), count, 0, &mut bytes) }.map_err( |error| { @@ -595,17 +610,18 @@ impl Attributes { )) }, )?; + value.initialized = true; Ok(value) } fn ptr(&mut self) -> LPPROC_THREAD_ATTRIBUTE_LIST { - LPPROC_THREAD_ATTRIBUTE_LIST(self.0.as_mut_ptr().cast()) + LPPROC_THREAD_ATTRIBUTE_LIST(self.storage.as_mut_ptr().cast()) } } impl Drop for Attributes { fn drop(&mut self) { - if !self.0.is_empty() { + if self.initialized { // SAFETY: this list was initialized and has not yet been deleted. unsafe { DeleteProcThreadAttributeList(self.ptr()) }; } @@ -984,6 +1000,27 @@ mod tests { wait_for_exit(second); } + #[test] + fn managed_process_reports_exit_code_259_as_exited() { + let mut config = test_config(); + config.program = PathBuf::from("cmd.exe"); + config.args = vec![OsString::from("/C"), OsString::from("exit /B 259")]; + let mut process = ManagedProcess::spawn(&config).expect("start an exit-code test child"); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(status) = process.try_wait().expect("poll the exit-code test child") { + assert_eq!(status.exit_code(), 259); + break; + } + assert!( + Instant::now() < deadline, + "the child with exit code 259 was reported as running" + ); + std::thread::sleep(Duration::from_millis(25)); + } + drop(process); + } + #[tokio::test] async fn managed_process_contains_descendants_of_the_pty_child() { let marker = std::env::temp_dir().join(format!( From d67749deacf4654945386112e12db222d318357f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 19:13:23 +0700 Subject: [PATCH 07/10] chore(antigravity): keep Windows test imports warning-free --- rust/src/managed_process.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index 22bc8fb05a..584b85dd29 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -19,7 +19,9 @@ use std::ffi::{OsStr, OsString}; use std::fs::File; use std::io::{Read, Write}; use std::os::windows::ffi::OsStrExt; -use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle}; +#[cfg(test)] +use std::os::windows::io::RawHandle; +use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; use std::path::{Path, PathBuf}; use std::time::Duration; @@ -1053,7 +1055,7 @@ mod tests { pty_cols: 120, label: "test-descendant".to_string(), }; - let mut process = ManagedProcess::spawn(&config).expect("start a managed test process"); + let process = ManagedProcess::spawn(&config).expect("start a managed test process"); let descendant = wait_for_descendant_pid(&marker); assert!( From b3660ec01b420106341b68e3c3013133cb30d59f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 19:26:16 +0700 Subject: [PATCH 08/10] fix(antigravity): keep job attribute storage alive --- rust/src/managed_process.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index 584b85dd29..fad84c490b 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -391,11 +391,12 @@ fn spawn_pty_child( .map(|dir| encode_wide_nul(dir.as_os_str())) .transpose()?; - let mut attributes = Attributes::new(2)?; let jobs = [win_handle(job)]; - // SAFETY: `jobs` and `con` stay alive through CreateProcessW. The job-list - // attribute binds the child before its first thread can run; the - // pseudoconsole attribute wires its stdio to the PTY above. + let mut attributes = Attributes::new(2)?; + // SAFETY: `jobs` is declared before `attributes`, so its backing storage + // remains valid through `DeleteProcThreadAttributeList` in `Attributes`'s + // destructor. The job-list attribute binds the child before its first + // thread can run; the pseudoconsole attribute wires its stdio to the PTY. unsafe { UpdateProcThreadAttribute( attributes.ptr(), From 394968c6ab8f51cd74f9052764f88dd7866b10c7 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:13:05 +0700 Subject: [PATCH 09/10] test(antigravity): wait with synchronize process handles --- rust/src/managed_process.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index fad84c490b..feea53efc3 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -860,11 +860,18 @@ mod tests { 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, WaitForSingleObject, + 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, false, pid) } { + 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; @@ -967,7 +974,6 @@ mod tests { ); process.shutdown(Duration::from_secs(5)).await; - assert!(!process_is_alive(pid), "shutdown stops the owned child"); } From e4d87692135f223121e6155ded19ae0347b3638c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:18:34 +0700 Subject: [PATCH 10/10] fix(antigravity): close PTY output before ConPTY teardown --- rust/src/managed_process.rs | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index feea53efc3..d6a9627914 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -317,12 +317,16 @@ impl ManagedChild { struct PseudoConsole { con: HPCON, input: Option, - output: OwnedHandle, + output: Option, } impl PseudoConsole { fn try_clone_reader(&self) -> std::io::Result> { - let handle = self.output.try_clone()?; + let handle = self + .output + .as_ref() + .expect("PTY output remains owned until cleanup") + .try_clone()?; Ok(Box::new(File::from(handle))) } @@ -335,6 +339,9 @@ impl PseudoConsole { impl Drop for PseudoConsole { fn drop(&mut self) { + // Close the parent-side output before asking ConPTY to close. Windows + // may wait in ClosePseudoConsole while an output pipe remains open. + drop(self.output.take()); if !self.con.is_invalid() { // SAFETY: this pseudoconsole was created here and is closed once. unsafe { ClosePseudoConsole(self.con) }; @@ -380,7 +387,7 @@ fn spawn_pty_child( let pty = PseudoConsole { con, input: Some(input_write), - output: output_read, + output: Some(output_read), }; let mut cmdline = build_command_line(&config.program, &config.args)?; @@ -909,6 +916,25 @@ mod tests { assert_eq!(terminal_cursor_reply_allowance(MAX_CURSOR_REPLIES, 1), 0); } + #[test] + fn failed_spawn_closes_pseudoconsole_promptly() { + let mut config = test_config(); + config.args.push(OsString::from("\0")); + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let failed = ManagedProcess::spawn(&config).is_err(); + sender + .send(failed) + .expect("failure result receiver remains available"); + }); + + assert!( + receiver + .recv_timeout(Duration::from_secs(3)) + .expect("post-ConPTY setup failure should return promptly") + ); + } + #[test] fn listener_table_finds_current_process_port() { let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0))