diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 8de6797989..b4445c0b63 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -434,6 +434,7 @@ dependencies = [ "getrandom 0.2.17", "http", "indexmap", + "libc", "libloading", "native-tls", "parking_lot", @@ -454,6 +455,7 @@ dependencies = [ "tracing", "ureq", "uuid", + "windows-sys 0.61.2", "zip", ] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 0f18a9b159..beb026e86b 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -70,6 +70,18 @@ tokio-tungstenite = { version = "0.24", default-features = false, features = ["c [target.'cfg(windows)'.dependencies] zip = { version = "2", default-features = false, features = ["deflate"], optional = true } +# Windows Job Object primitives backing `process_tree`'s CLI process-tree +# containment (`JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE`). +windows-sys = { version = "0.61", default-features = false, features = [ + "Win32_Foundation", + "Win32_System_JobObjects", + "Win32_System_Threading", +] } + +[target.'cfg(unix)'.dependencies] +# Process-group primitives backing `process_tree`'s CLI process-tree +# containment (`killpg`). +libc = "0.2" [dev-dependencies] rusqlite = { version = "0.35", features = ["bundled"] } diff --git a/rust/src/errors.rs b/rust/src/errors.rs index 70f4c14ff1..ab6a6bed2b 100644 --- a/rust/src/errors.rs +++ b/rust/src/errors.rs @@ -399,11 +399,11 @@ fn capture_backtrace() -> Option> { /// Aggregate of errors collected during [`crate::Client::stop`]. /// /// `Client::stop` performs cooperative shutdown across every active -/// session before killing the CLI child process. Errors from any -/// per-session `session.destroy` RPC and from the terminal child-kill -/// step are collected here rather than short-circuiting on the first -/// failure, so callers see the full picture of what went wrong during -/// teardown. +/// session before tearing down the CLI's process tree. Errors from any +/// per-session `session.destroy` RPC, runtime shutdown, process-tree +/// termination, and the terminal child reap are collected here rather +/// than short-circuiting on the first failure, so callers see the full +/// picture of what went wrong during teardown. /// /// Implements [`std::error::Error`] and forwards to `Display` for the /// first error, with a count suffix when there are more. @@ -411,8 +411,9 @@ fn capture_backtrace() -> Option> { pub struct StopErrors(pub(crate) Vec); impl StopErrors { - /// Borrow the collected errors as a slice, in the order they - /// occurred (per-session destroys first, then child-kill last). + /// Borrow the collected errors as a slice, in the order they occurred + /// (per-session destroys first, then runtime shutdown, process-tree + /// termination, and the final child reap). pub fn errors(&self) -> &[Error] { &self.0 } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 4a9f73ca4c..a7c3b1b898 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -31,6 +31,10 @@ pub mod hooks; mod jsonrpc; /// Permission-policy helpers that produce a [`handler::PermissionHandler`]. pub mod permission; +/// Cross-platform ownership of the spawned CLI process tree (Unix process +/// group / Windows Job Object) used by `Client::stop`, `force_stop`, and +/// `Drop` to reach descendants the CLI spawns, not only the CLI itself. +pub(crate) mod process_tree; /// BYOK bearer-token provider callbacks. pub mod provider_token; mod provider_token_dispatch; @@ -1017,6 +1021,11 @@ impl std::fmt::Debug for Client { struct ClientInner { child: parking_lot::Mutex>, + /// Containment for `child` and any descendants it spawns. `None` when + /// there is no child (streams-only transports) or when attaching + /// containment failed and teardown fell back to root-only (see + /// [`process_tree::attach`]). + process_tree: parking_lot::Mutex>, #[cfg(feature = "bundled-in-process")] /// In-process FFI runtime host, set only for [`Transport::InProcess`]. /// Closing it tears down the native runtime connection. @@ -1057,6 +1066,49 @@ struct ClientInner { startup_timings: OnceLock, } +struct SpawnedProcessGuard { + child: Option, + tree: Option, +} + +impl SpawnedProcessGuard { + fn new(child: Child, tree: Option) -> Self { + Self { + child: Some(child), + tree, + } + } + + fn child_mut(&mut self) -> &mut Child { + self.child.as_mut().expect("spawned process has a child") + } + + fn pid(&self) -> Option { + self.child.as_ref().and_then(Child::id) + } + + fn transfer_to(mut self, client: &Client) { + *client.inner.child.lock() = self.child.take(); + *client.inner.process_tree.lock() = self.tree.take(); + } +} + +impl Drop for SpawnedProcessGuard { + fn drop(&mut self) { + let pid = self.pid(); + if let Some(tree) = self.tree.take() + && let Err(error) = tree.terminate() + { + error!(pid = ?pid, %error, "failed to terminate CLI process tree during startup cleanup"); + } + if let Some(child) = self.child.as_mut() + && let Err(error) = child.start_kill() + { + error!(pid = ?pid, %error, "failed to kill CLI process during startup cleanup"); + } + } +} + impl Client { /// Start a CLI server process with the given options. /// @@ -1250,6 +1302,7 @@ impl Client { reader, writer, None, + None, working_directory, options.on_list_models, session_fs_config.is_some(), @@ -1264,7 +1317,7 @@ impl Client { port, connection_token: _, } => { - let (mut child, actual_port, spawn_elapsed, port_wait_elapsed) = + let (mut process, actual_port, spawn_elapsed, port_wait_elapsed) = Self::spawn_tcp(&program, &options, &working_directory, port).await?; timings.process_spawn_ms = Some(StartupTimings::millis(spawn_elapsed)); timings.port_wait_ms = Some(StartupTimings::millis(port_wait_elapsed)); @@ -1276,11 +1329,11 @@ impl Client { "Client::start TCP connect complete" ); let (reader, writer) = tokio::io::split(stream); - Self::drain_stderr(&mut child); - Self::from_transport( + Self::drain_stderr(process.child_mut()); + Self::from_spawned_transport( reader, writer, - Some(child), + process, working_directory, options.on_list_models, session_fs_config.is_some(), @@ -1292,7 +1345,7 @@ impl Client { )? } Transport::Stdio => { - let (mut child, spawn_elapsed) = + let (mut child, tree, spawn_elapsed) = Self::spawn_stdio(&program, &options, &working_directory)?; timings.process_spawn_ms = Some(StartupTimings::millis(spawn_elapsed)); let stdin = child.stdin.take().expect("stdin is piped"); @@ -1302,6 +1355,7 @@ impl Client { stdout, stdin, Some(child), + tree, working_directory, options.on_list_models, session_fs_config.is_some(), @@ -1359,6 +1413,7 @@ impl Client { reader, writer, None, + None, working_directory, options.on_list_models, session_fs_config.is_some(), @@ -1472,6 +1527,38 @@ impl Client { Ok(client) } + #[allow(clippy::too_many_arguments)] + fn from_spawned_transport( + reader: impl AsyncRead + Unpin + Send + 'static, + writer: impl AsyncWrite + Unpin + Send + 'static, + process: SpawnedProcessGuard, + cwd: PathBuf, + on_list_models: Option>, + session_fs_configured: bool, + session_fs_sqlite_declared: bool, + on_get_trace_context: Option>, + on_github_telemetry: Option, + effective_connection_token: Option, + mode: ClientMode, + ) -> Result { + let client = Self::from_transport( + reader, + writer, + None, + None, + cwd, + on_list_models, + session_fs_configured, + session_fs_sqlite_declared, + on_get_trace_context, + on_github_telemetry, + effective_connection_token, + mode, + )?; + process.transfer_to(&client); + Ok(client) + } + /// Create a Client from raw async streams (no child process). /// /// Useful for testing or connecting to a server over a custom transport. @@ -1484,6 +1571,7 @@ impl Client { reader, writer, None, + None, cwd, None, false, @@ -1513,6 +1601,7 @@ impl Client { reader, writer, None, + None, cwd, None, false, @@ -1538,6 +1627,7 @@ impl Client { reader, writer, None, + None, cwd, None, false, @@ -1563,6 +1653,7 @@ impl Client { reader, writer, None, + None, cwd, None, false, @@ -1589,6 +1680,7 @@ impl Client { reader: impl AsyncRead + Unpin + Send + 'static, writer: impl AsyncWrite + Unpin + Send + 'static, child: Option, + tree: Option, cwd: PathBuf, on_list_models: Option>, session_fs_configured: bool, @@ -1615,6 +1707,7 @@ impl Client { let client = Self { inner: Arc::new(ClientInner { child: parking_lot::Mutex::new(child), + process_tree: parking_lot::Mutex::new(tree), #[cfg(feature = "bundled-in-process")] ffi_host: parking_lot::Mutex::new(None), rpc, @@ -1691,6 +1784,7 @@ impl Client { fn build_command(program: &Path, options: &ClientOptions, working_directory: &Path) -> Command { let mut command = Command::new(program); command.kill_on_drop(true); + process_tree::configure(&mut command); for arg in &options.prefix_args { command.arg(arg); } @@ -1814,7 +1908,7 @@ impl Client { program: &Path, options: &ClientOptions, working_directory: &Path, - ) -> Result<(Child, Duration)> { + ) -> Result<(Child, Option, Duration)> { info!(cwd = ?working_directory, program = %program.display(), "spawning copilot CLI (stdio)"); let mut command = Self::build_command(program, options, working_directory); command @@ -1828,11 +1922,12 @@ impl Client { let spawn_start = Instant::now(); let child = command.spawn()?; let spawn_elapsed = spawn_start.elapsed(); + let tree = process_tree::attach(&child); debug!( elapsed_ms = spawn_elapsed.as_millis(), "Client::spawn_stdio subprocess spawned" ); - Ok((child, spawn_elapsed)) + Ok((child, tree, spawn_elapsed)) } async fn spawn_tcp( @@ -1840,7 +1935,7 @@ impl Client { options: &ClientOptions, working_directory: &Path, port: u16, - ) -> Result<(Child, u16, Duration, Duration)> { + ) -> Result<(SpawnedProcessGuard, u16, Duration, Duration)> { info!(cwd = ?working_directory, program = %program.display(), port = %port, "spawning copilot CLI (tcp)"); let mut command = Self::build_command(program, options, working_directory); command @@ -1852,13 +1947,15 @@ impl Client { .args(&options.extra_args) .stdin(Stdio::null()); let spawn_start = Instant::now(); - let mut child = command.spawn()?; + let child = command.spawn()?; let spawn_elapsed = spawn_start.elapsed(); + let tree = process_tree::attach(&child); + let mut process = SpawnedProcessGuard::new(child, tree); debug!( elapsed_ms = spawn_elapsed.as_millis(), "Client::spawn_tcp subprocess spawned" ); - let stdout = child.stdout.take().expect("stdout is piped"); + let stdout = process.child_mut().stdout.take().expect("stdout is piped"); let (port_tx, port_rx) = oneshot::channel::(); let span = tracing::error_span!("copilot_cli_port_scan"); @@ -1899,7 +1996,7 @@ impl Client { "Client::spawn_tcp TCP port wait complete" ); info!(port = %actual_port, "CLI server listening"); - Ok((child, actual_port, spawn_elapsed, port_wait_elapsed)) + Ok((process, actual_port, spawn_elapsed, port_wait_elapsed)) } fn drain_stderr(child: &mut Child) { @@ -2439,15 +2536,17 @@ impl Client { /// Cooperatively shut down the client and the CLI child process. /// /// Walks every still-registered session and sends `session.destroy` - /// for each one, asks SDK-owned runtimes to shut down, then kills the - /// CLI child. Errors from per-session destroys, runtime shutdown, and - /// the final child-kill are collected into - /// [`StopErrors`] rather than short-circuiting on the first failure - /// — so callers see the full picture of teardown. + /// for each one, asks SDK-owned runtimes to shut down, then terminates + /// the CLI and, when process-tree containment was attached successfully, + /// its descendants, then reaps the CLI child. Errors from per-session + /// destroys, runtime shutdown, process-tree termination, and the final + /// child reap are collected into [`StopErrors`] rather than + /// short-circuiting on the first failure, so callers see the full + /// picture of teardown. /// /// If you have already called [`Session::disconnect`] on every /// session this client created, the per-session destroy step is a - /// no-op (the router map is empty); only the child-kill remains. + /// no-op (the router map is empty); only process-tree teardown remains. /// /// [`Session::disconnect`]: crate::session::Session::disconnect /// @@ -2455,11 +2554,12 @@ impl Client { /// /// **Cancel-unsafe but recoverable.** The body sequentially destroys /// every registered session (each via [`Client::call`](Self::call), - /// individually cancel-safe) before killing the child. Cancelling - /// `stop()` mid-loop leaves some sessions still in the router map - /// and the child still running. Recovery: call [`force_stop`](Self::force_stop) - /// (sync, kills the child unconditionally and clears router state) - /// or call `stop()` again with a fresh future. The documented + /// individually cancel-safe) before tearing down the process tree. + /// Cancelling `stop()` mid-loop leaves some sessions still in the + /// router map and the tree still running. Recovery: call + /// [`force_stop`](Self::force_stop) (sync, terminates the tree + /// unconditionally and clears router state) or call `stop()` again + /// with a fresh future. The documented /// `tokio::time::timeout(..., client.stop())` pattern in the example /// below uses `force_stop` as the fallback for exactly this case. pub async fn stop(&self) -> std::result::Result<(), StopErrors> { @@ -2531,8 +2631,19 @@ impl Client { } let child = self.inner.child.lock().take(); + let tree = self.inner.process_tree.lock().take(); *self.inner.state.lock() = ConnectionState::Disconnected; *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new()); + // Reach every descendant the CLI spawned, not only the CLI itself, + // before reaping the root below. `runtime.shutdown` above already + // asked the CLI to clean up cooperatively; this is the backstop + // that still runs even when that cooperative cleanup left + // something behind. + if let Some(tree) = &tree + && let Err(e) = tree.terminate() + { + errors.push(e.into()); + } if let Some(mut child) = child { match child.try_wait() { Ok(Some(_status)) => {} @@ -2571,12 +2682,13 @@ impl Client { /// Forcibly stop the CLI process without waiting for it to exit. /// - /// Synchronous fallback when [`stop`](Self::stop) is unsuitable — for + /// Synchronous fallback when [`stop`](Self::stop) is unsuitable, for /// example when the awaiting tokio runtime is shutting down or the - /// process is wedged on I/O. Sends a kill signal without awaiting - /// reaper completion and immediately drops all per-session router - /// state so dependent tasks observe a closed channel rather than a - /// hang. + /// process is wedged on I/O. Terminates the CLI and, when process-tree + /// containment was attached successfully, its descendants. It sends a + /// kill signal to the CLI child without awaiting reaper completion and + /// immediately drops all per-session router state so dependent tasks + /// observe a closed channel rather than a hang. /// /// # Cancel safety /// @@ -2601,6 +2713,11 @@ impl Client { pub fn force_stop(&self) { let pid = self.pid(); info!(pid = ?pid, "force-stopping CLI process"); + if let Some(tree) = self.inner.process_tree.lock().take() + && let Err(e) = tree.terminate() + { + error!(pid = ?pid, error = %e, "failed to terminate CLI process tree"); + } if let Some(mut child) = self.inner.child.lock().take() && let Err(e) = child.start_kill() { @@ -2662,8 +2779,13 @@ impl Client { impl Drop for ClientInner { fn drop(&mut self) { + let pid = self.child.lock().as_ref().and_then(Child::id); + if let Some(tree) = self.process_tree.lock().take() + && let Err(e) = tree.terminate() + { + error!(pid = ?pid, error = %e, "failed to terminate CLI process tree on drop"); + } if let Some(ref mut child) = *self.child.lock() { - let pid = child.id(); if let Err(e) = child.start_kill() { error!(pid = ?pid, error = %e, "failed to kill CLI process on drop"); } else { @@ -3323,6 +3445,7 @@ mod tests { client_read, client_write, Some(child), + None, temp.path().to_path_buf(), None, false, @@ -3340,6 +3463,393 @@ mod tests { assert_test_child_killed(&survived).await; } + #[cfg(any(unix, windows))] + #[tokio::test] + async fn dropping_last_client_terminates_process_tree_and_releases_descendant_lock() { + let temp = tempfile::tempdir().unwrap(); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let start_path = temp.path().join("start"); + let mut command = root_with_lock_holding_descendant( + temp.path(), + &lock_path, + &descendant_ready, + &start_path, + ); + let child = command.spawn().unwrap(); + let tree = crate::process_tree::attach(&child); + + let (client_write, server_read) = tokio::io::duplex(64); + let (server_write, client_read) = tokio::io::duplex(64); + drop(server_read); + drop(server_write); + let client = Client::from_transport( + client_read, + client_write, + Some(child), + tree, + temp.path().to_path_buf(), + None, + false, + false, + None, + None, + None, + ClientMode::default(), + ) + .unwrap(); + + std::fs::write(&start_path, b"go").unwrap(); + wait_for_test_path(&descendant_ready).await; + let descendant_pid: u32 = std::fs::read_to_string(&descendant_ready) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!(crate::process_tree::process_alive(descendant_pid)); + + drop(client); + + wait_until_test( + || !crate::process_tree::process_alive(descendant_pid), + "descendant survived dropping the last Client", + ) + .await; + wait_until_test( + || descendant_lock_is_free(&lock_path), + "descendant's lock was never released after dropping the last Client", + ) + .await; + } + + #[cfg(unix)] + #[tokio::test] + async fn failed_tcp_startup_terminates_process_tree_and_releases_descendant_lock() { + let temp = tempfile::tempdir().unwrap(); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let this_test_binary = std::env::current_exe().unwrap(); + let unavailable_port = { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).unwrap(); + listener.local_addr().unwrap().port() + }; + const HELPER_FILTER: &str = "process_tree::tests::lock_holder_helper_entrypoint"; + let options = ClientOptions::new() + .with_program(CliProgram::Path(PathBuf::from("sh"))) + .with_prefix_args([ + OsString::from("-c"), + OsString::from( + "\"$HELPER_BIN\" \"$HELPER_FILTER\" --exact --nocapture & \ + while [ ! -f \"$PROCESS_TREE_TEST_READY_PATH\" ]; do sleep 0.01; done; \ + echo \"listening on port $REPORTED_PORT\"; wait", + ), + ]) + .with_cwd(temp.path()) + .with_env([ + ("HELPER_BIN", this_test_binary.into_os_string()), + ("HELPER_FILTER", OsString::from(HELPER_FILTER)), + ( + "PROCESS_TREE_TEST_LOCK_PATH", + lock_path.clone().into_os_string(), + ), + ( + "PROCESS_TREE_TEST_READY_PATH", + descendant_ready.clone().into_os_string(), + ), + ( + "REPORTED_PORT", + OsString::from(unavailable_port.to_string()), + ), + ]) + .with_transport(Transport::Tcp { + port: 0, + connection_token: None, + }); + + let error = Client::start(options).await.unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Io)); + wait_for_test_path(&descendant_ready).await; + let descendant_pid: u32 = std::fs::read_to_string(&descendant_ready) + .unwrap() + .trim() + .parse() + .unwrap(); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(15); + let (descendant_terminated, lock_released) = loop { + let descendant_terminated = !crate::process_tree::process_alive(descendant_pid); + let lock_released = descendant_lock_is_free(&lock_path); + if (descendant_terminated && lock_released) || tokio::time::Instant::now() >= deadline { + break (descendant_terminated, lock_released); + } + tokio::time::sleep(Duration::from_millis(10)).await; + }; + if !descendant_terminated { + // SAFETY: the pid came from the test-only child process. + unsafe { + libc::kill(descendant_pid as i32, libc::SIGKILL); + } + } + + assert!( + descendant_terminated, + "descendant survived failed Client::start; lock_released={lock_released}" + ); + assert!( + lock_released, + "descendant's lock remained held after failed Client::start" + ); + } + + #[cfg(any(unix, windows))] + #[tokio::test] + async fn force_stop_terminates_process_tree_and_releases_descendant_lock() { + let temp = tempfile::tempdir().unwrap(); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let start_path = temp.path().join("start"); + let mut command = root_with_lock_holding_descendant( + temp.path(), + &lock_path, + &descendant_ready, + &start_path, + ); + let child = command.spawn().unwrap(); + let root_pid = child.id().unwrap(); + // Note: whether attachment succeeds is an environment detail, not + // part of the contract under test — `force_stop` must reach the + // descendant when it does, and degrade to root-only teardown when + // it doesn't. This test exercises the succeeding case. + let tree = crate::process_tree::attach(&child); + + let (client_write, server_read) = tokio::io::duplex(64); + let (server_write, client_read) = tokio::io::duplex(64); + drop(server_read); + drop(server_write); + let client = Client::from_transport( + client_read, + client_write, + Some(child), + tree, + temp.path().to_path_buf(), + None, + false, + false, + None, + None, + None, + ClientMode::default(), + ) + .unwrap(); + + std::fs::write(&start_path, b"go").unwrap(); + wait_for_test_path(&descendant_ready).await; + let descendant_pid: u32 = std::fs::read_to_string(&descendant_ready) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!(crate::process_tree::process_alive(root_pid)); + assert!(crate::process_tree::process_alive(descendant_pid)); + + client.force_stop(); + + wait_until_test( + || !crate::process_tree::process_alive(descendant_pid), + "descendant survived Client::force_stop", + ) + .await; + wait_until_test( + || descendant_lock_is_free(&lock_path), + "descendant's lock was never released after Client::force_stop", + ) + .await; + } + + #[cfg(any(unix, windows))] + #[tokio::test] + async fn stop_awaits_root_reap_and_terminates_process_tree() { + let temp = tempfile::tempdir().unwrap(); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let start_path = temp.path().join("start"); + let mut command = root_with_lock_holding_descendant( + temp.path(), + &lock_path, + &descendant_ready, + &start_path, + ); + let child = command.spawn().unwrap(); + let root_pid = child.id().unwrap(); + // Note: whether attachment succeeds is an environment detail; see + // the matching comment in `force_stop_terminates_process_tree_and_ + // releases_descendant_lock` above. + let tree = crate::process_tree::attach(&child); + + let (client_write, server_read) = tokio::io::duplex(64); + let (server_write, client_read) = tokio::io::duplex(64); + drop(server_read); + drop(server_write); + let client = Client::from_transport( + client_read, + client_write, + Some(child), + tree, + temp.path().to_path_buf(), + None, + false, + false, + None, + None, + None, + ClientMode::default(), + ) + .unwrap(); + + std::fs::write(&start_path, b"go").unwrap(); + wait_for_test_path(&descendant_ready).await; + let descendant_pid: u32 = std::fs::read_to_string(&descendant_ready) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!(crate::process_tree::process_alive(root_pid)); + assert!(crate::process_tree::process_alive(descendant_pid)); + + // `stop()` returning at all — rather than hanging — is itself part + // of the proof: it must reap the root before returning. + let _ = client.stop().await; + + assert!( + !crate::process_tree::process_alive(root_pid), + "root must be reaped by the time stop() returns" + ); + wait_until_test( + || !crate::process_tree::process_alive(descendant_pid), + "descendant survived Client::stop", + ) + .await; + wait_until_test( + || descendant_lock_is_free(&lock_path), + "descendant's lock was never released after Client::stop", + ) + .await; + } + + /// Root command that waits for `start_path` to appear, then re-execs + /// this same test binary's `process_tree::tests::lock_holder_helper_entrypoint` + /// as a descendant and blocks so the root itself stays alive until + /// killed. Mirrors the identical helper in `process_tree`'s own tests; + /// duplicated here (rather than shared) because it builds the command + /// through `Client::build_command` to exercise the exact same spawn + /// path `Client::start` uses, which `process_tree`'s own tests + /// deliberately do not depend on. + #[cfg(any(unix, windows))] + fn root_with_lock_holding_descendant( + temp: &Path, + lock_path: &Path, + descendant_ready: &Path, + start_path: &Path, + ) -> Command { + let this_test_binary = std::env::current_exe().unwrap(); + const HELPER_FILTER: &str = "process_tree::tests::lock_holder_helper_entrypoint"; + #[cfg(unix)] + let mut command = { + let mut command = + Client::build_command(Path::new("sh"), &ClientOptions::default(), temp); + command.args([ + "-c", + "while [ ! -f \"$START\" ]; do sleep 0.05; done; \ + \"$HELPER_BIN\" \"$HELPER_FILTER\" --exact --nocapture & wait", + ]); + command + }; + #[cfg(windows)] + let mut command = { + let mut command = + Client::build_command(Path::new("powershell.exe"), &ClientOptions::default(), temp); + command.args([ + "-NoLogo", + "-NoProfile", + "-NonInteractive", + "-Command", + "while (-not (Test-Path $env:START)) { Start-Sleep -Milliseconds 50 }; \ + Start-Process -FilePath $env:HELPER_BIN -ArgumentList @( \ + $env:HELPER_FILTER,'--exact','--nocapture') -NoNewWindow -Wait", + ]); + command + }; + command + .env("HELPER_FILTER", HELPER_FILTER) + .env("HELPER_BIN", &this_test_binary) + .env("PROCESS_TREE_TEST_LOCK_PATH", lock_path) + .env("PROCESS_TREE_TEST_READY_PATH", descendant_ready) + .env("START", start_path) + // See the matching comment in `process_tree`'s own test helper: + // the re-exec'd descendant inherits stdio by default, which + // would otherwise keep this test's output pipe open for as + // long as that descendant is alive. + .stdout(Stdio::null()) + .stderr(Stdio::null()); + command + } + + #[cfg(any(unix, windows))] + async fn wait_for_test_path(path: &Path) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(15); + while !path.exists() { + assert!( + tokio::time::Instant::now() < deadline, + "expected file was never created: {}", + path.display() + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + #[cfg(any(unix, windows))] + async fn wait_until_test(mut predicate: impl FnMut() -> bool, message: &str) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(15); + while !predicate() { + assert!(tokio::time::Instant::now() < deadline, "{message}"); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + /// Attempts to acquire the same lock the descendant holds. Succeeds + /// only once the descendant process has actually released it. + #[cfg(any(unix, windows))] + fn descendant_lock_is_free(path: &Path) -> bool { + #[cfg(unix)] + { + use std::os::fd::AsRawFd; + + let Ok(file) = std::fs::OpenOptions::new().write(true).open(path) else { + return false; + }; + // SAFETY: `file` owns a valid fd for the duration of this call. + let acquired = + unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0; + if acquired { + // SAFETY: `file` still owns the fd we just locked. + unsafe { + libc::flock(file.as_raw_fd(), libc::LOCK_UN); + } + } + acquired + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + + std::fs::OpenOptions::new() + .write(true) + .share_mode(0) + .open(path) + .is_ok() + } + } + #[cfg(any(unix, windows))] #[tokio::test] async fn spawned_child_is_killed_when_dropped() { @@ -3411,6 +3921,7 @@ mod tests { Client { inner: Arc::new(ClientInner { child: parking_lot::Mutex::new(None), + process_tree: parking_lot::Mutex::new(None), #[cfg(feature = "bundled-in-process")] ffi_host: parking_lot::Mutex::new(None), rpc: { diff --git a/rust/src/process_tree.rs b/rust/src/process_tree.rs new file mode 100644 index 0000000000..95a18b1600 --- /dev/null +++ b/rust/src/process_tree.rs @@ -0,0 +1,591 @@ +//! Cross-platform ownership of the CLI process tree. +//! +//! [`Client`](crate::Client) spawns the CLI as a direct child, and the CLI +//! may itself spawn descendants (MCP servers, shell tools, subagents). +//! Without containment, [`Client::stop`](crate::Client::stop), +//! [`Client::force_stop`](crate::Client::force_stop), and `Drop` can only +//! reach the root: killing it leaves any descendant it spawned running +//! and holding whatever it held (files, sockets, locks). +//! +//! [`ProcessTree`] gives those three teardown paths a single primitive that +//! reaches the whole tree instead of only the root: a dedicated process +//! group on Unix, a Job Object carrying `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` +//! on Windows. +//! +//! # What this does and does not guarantee +//! +//! Both primitives terminate the tree once [`ProcessTree::terminate`] +//! actually runs. They differ in what happens if it never runs — the SDK's +//! own process crashes, is killed, or loses power before any teardown code +//! executes: +//! +//! - **Windows**: `JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE` is enforced by the +//! kernel when the last handle to the Job Object closes, including on an +//! unclean exit of this process. The tree dies even if no Rust code runs. +//! - **Unix**: a process group has no equivalent kernel-enforced cleanup. +//! [`ProcessTree::terminate`] must actually execute for the group to die; +//! an unclean exit of this process leaves the group running, and a +//! descendant that calls `setsid()` escapes the group entirely. This +//! asymmetry is inherent to the two platforms' process models, not an +//! implementation gap here. +//! +//! Attaching the Windows Job Object needs a small window after `spawn()` to +//! call `AssignProcessToJobObject` on the live process handle; a child that +//! spawns its own descendants faster than that call can still escape +//! containment. Closing that window requires creating the process suspended +//! and resuming its primary thread only after assignment, which trades this +//! narrow race for a more invasive spawn path. This module accepts the race +//! and assigns immediately after spawn instead. + +use std::io; + +use tokio::process::{Child, Command}; +use tracing::warn; + +/// Configure `command` so its eventual child is contained in a dedicated +/// process tree from the moment it is spawned. +/// +/// Call before `command.spawn()`. Cheap and infallible on both platforms — +/// it only sets flags on the `Command`, it never touches a live process. +/// [`attach`] performs the (Windows-only) work that needs a live process +/// handle. +pub(crate) fn configure(command: &mut Command) { + platform::configure(command); +} + +/// Attach the platform containment primitive to a freshly spawned `child`. +/// +/// Call immediately after `spawn()`, before the child has had a chance to +/// create descendants (see the module-level Windows caveat). Returns `None` +/// (after logging a warning) rather than an error when attachment fails on a +/// platform where failure is recoverable — a `Client` whose containment +/// setup failed should still start normally and fall back to root-only +/// teardown rather than turning a containment failure into a hard start +/// failure. Unix attachment cannot practically fail once `spawn()` has +/// already succeeded (the process group was established at fork, before +/// `exec`), so it never needs this fallback. +pub(crate) fn attach(child: &Child) -> Option { + match platform::attach(child) { + Ok(tree) => Some(ProcessTree(tree)), + Err(error) => { + warn!( + pid = ?child.id(), + %error, + "failed to attach CLI process to a containment tree; \ + falling back to root-process-only teardown" + ); + None + } + } +} + +/// Owns the platform primitive that contains one spawned root process and +/// its descendants. +pub(crate) struct ProcessTree(platform::Tree); + +impl ProcessTree { + /// Signal every process still alive in the tree to exit immediately. + /// + /// Idempotent: safe to call after the tree has already exited. + pub(crate) fn terminate(&self) -> io::Result<()> { + self.0.terminate() + } + + /// `true` once no process remains in the tree. + #[cfg(test)] + pub(crate) fn is_empty(&self) -> io::Result { + self.0.is_empty() + } +} + +/// `true` if a process with `pid` is still alive. Test-only introspection — +/// shared by this module's own tests and by the process-tree tests in +/// `lib.rs` so both check liveness the same way. +#[cfg(test)] +pub(crate) fn process_alive(pid: u32) -> bool { + platform::process_alive(pid) +} + +#[cfg(unix)] +mod platform { + use std::io; + + use tokio::process::{Child, Command}; + + pub(super) struct Tree { + /// The dedicated process group id, equal to the root child's pid. + pgid: i32, + } + + pub(super) fn configure(command: &mut Command) { + // Put the eventual child in its own new process group (pgid == + // its own pid), established by the kernel at fork, before `exec` + // runs. Anything the child spawns inherits this group unless it + // explicitly changes its own, so containment exists before we ever + // get a chance to call `attach`. + command.process_group(0); + } + + pub(super) fn attach(child: &Child) -> io::Result { + let pid = child.id().ok_or_else(|| { + io::Error::new( + io::ErrorKind::NotFound, + "CLI process exited before its process group could be recorded", + ) + })?; + Ok(Tree { pgid: pid as i32 }) + } + + impl Tree { + pub(super) fn terminate(&self) -> io::Result<()> { + // SAFETY: `killpg` only ever signals processes sharing + // `self.pgid`; no pointers are involved. + if unsafe { libc::killpg(self.pgid, libc::SIGKILL) } == 0 { + return Ok(()); + } + match io::Error::last_os_error() { + // No process left in the group — already terminated. + error if error.raw_os_error() == Some(libc::ESRCH) => Ok(()), + error => Err(error), + } + } + + #[cfg(test)] + pub(super) fn is_empty(&self) -> io::Result { + // Signal 0 only probes for existence. Delivering a signal to a + // process group requires no parent/child relationship, only + // that at least one process in the group is still alive, so + // this reads liveness without trying to reap anything — the + // root child stays owned solely by Tokio's `Child`. + // SAFETY: as above. + if unsafe { libc::killpg(self.pgid, 0) } == 0 { + return Ok(false); + } + match io::Error::last_os_error() { + error if error.raw_os_error() == Some(libc::ESRCH) => Ok(true), + error => Err(error), + } + } + } + + #[cfg(test)] + pub(super) fn process_alive(pid: u32) -> bool { + #[cfg(target_os = "linux")] + if let Ok(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) + && stat + .rsplit_once(") ") + .and_then(|(_, fields)| fields.chars().next()) + .is_some_and(|state| matches!(state, 'Z' | 'X')) + { + return false; + } + // SAFETY: signal 0 only probes process existence. + (unsafe { libc::kill(pid as i32, 0) }) == 0 + } +} + +#[cfg(windows)] +mod platform { + use std::{io, ptr}; + + use tokio::process::{Child, Command}; + use windows_sys::Win32::Foundation::{CloseHandle, HANDLE}; + use windows_sys::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, TerminateJobObject, + }; + #[cfg(test)] + use windows_sys::Win32::System::JobObjects::{ + JOBOBJECT_BASIC_ACCOUNTING_INFORMATION, JobObjectBasicAccountingInformation, + QueryInformationJobObject, + }; + + pub(super) fn configure(_command: &mut Command) { + // Nothing to set on the `Command` itself; the Job Object is + // created and assigned to the live process in `attach`, after + // `spawn()` returns a process handle. + } + + struct OwnedHandle(HANDLE); + + // SAFETY: Win32 handles may be used and closed from any thread. + unsafe impl Send for OwnedHandle {} + unsafe impl Sync for OwnedHandle {} + + impl Drop for OwnedHandle { + fn drop(&mut self) { + // SAFETY: `self.0` is a valid handle owned uniquely by this value. + unsafe { + CloseHandle(self.0); + } + } + } + + pub(super) struct Tree { + job: OwnedHandle, + } + + pub(super) fn attach(child: &Child) -> io::Result { + // SAFETY: null attributes and name create a private, + // non-inheritable Job Object owned solely by this process. + let raw_job = unsafe { CreateJobObjectW(ptr::null(), ptr::null()) }; + if raw_job.is_null() { + return Err(io::Error::last_os_error()); + } + let job = OwnedHandle(raw_job); + + let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + // SAFETY: `job.0` is a live Job Object handle, and `limits` has the + // exact layout `JobObjectExtendedLimitInformation` requires. + if unsafe { + SetInformationJobObject( + job.0, + JobObjectExtendedLimitInformation, + ptr::from_ref(&limits).cast(), + size_of::() as u32, + ) + } == 0 + { + return Err(io::Error::last_os_error()); + } + + // Assign as soon as we have a live process handle — see the + // module-level doc for the residual assign-race this accepts. + let process = child.raw_handle().ok_or_else(|| { + io::Error::new( + io::ErrorKind::NotFound, + "CLI process exited before it could be assigned to a Job Object", + ) + })?; + // SAFETY: `job.0` is live, and `process` is the handle Tokio owns + // for this still-running child. + if unsafe { AssignProcessToJobObject(job.0, process.cast()) } == 0 { + return Err(io::Error::last_os_error()); + } + + Ok(Tree { job }) + } + + impl Tree { + pub(super) fn terminate(&self) -> io::Result<()> { + // SAFETY: `self.job` is a live Job Object handle owned by this value. + if unsafe { TerminateJobObject(self.job.0, 1) } != 0 { + Ok(()) + } else { + Err(io::Error::last_os_error()) + } + } + + #[cfg(test)] + pub(super) fn is_empty(&self) -> io::Result { + let mut accounting = JOBOBJECT_BASIC_ACCOUNTING_INFORMATION::default(); + // SAFETY: `accounting` has the exact layout the query requires. + if unsafe { + QueryInformationJobObject( + self.job.0, + JobObjectBasicAccountingInformation, + ptr::from_mut(&mut accounting).cast(), + size_of::() as u32, + ptr::null_mut(), + ) + } == 0 + { + return Err(io::Error::last_os_error()); + } + Ok(accounting.ActiveProcesses == 0) + } + } + + #[cfg(test)] + pub(super) fn process_alive(pid: u32) -> bool { + use windows_sys::Win32::Foundation::WAIT_TIMEOUT; + use windows_sys::Win32::System::Threading::{ + OpenProcess, PROCESS_SYNCHRONIZE, WaitForSingleObject, + }; + + // SAFETY: the process handle is closed before returning. + unsafe { + let process = OpenProcess(PROCESS_SYNCHRONIZE, 0, pid); + if process.is_null() { + return false; + } + let exists = WaitForSingleObject(process, 0) == WAIT_TIMEOUT; + CloseHandle(process); + exists + } + } +} + +#[cfg(test)] +mod tests { + use std::path::Path; + use std::time::Duration; + + use tempfile::tempdir; + + use super::*; + + const TEST_TIMEOUT: Duration = Duration::from_secs(15); + + /// Not a real test — a re-exec entry point. The containment tests spawn + /// this same compiled test binary as a subprocess, filtered by exact + /// name to run only this "test", so its own OS process becomes the + /// descendant the tests kill. It reads its lock/ready paths from + /// environment variables (set on the child before spawn) rather than + /// CLI args, since the harness controls what CLI args a filtered run + /// receives. Run normally (no env vars set), it is a harmless no-op. + #[test] + fn lock_holder_helper_entrypoint() { + let Ok(lock_path) = std::env::var("PROCESS_TREE_TEST_LOCK_PATH") else { + return; + }; + let ready_path = std::env::var("PROCESS_TREE_TEST_READY_PATH").expect("ready path env var"); + + let _lock = acquire_exclusive_lock(Path::new(&lock_path)).expect("acquire exclusive lock"); + std::fs::write(&ready_path, std::process::id().to_string()).expect("write ready file"); + + // The containment tests kill this process (root-only, or the whole + // tree); there is no voluntary shutdown path that reaches past here. + std::thread::sleep(Duration::from_secs(120)); + } + + #[cfg(unix)] + fn acquire_exclusive_lock(path: &Path) -> std::io::Result { + use std::os::fd::AsRawFd; + + let file = std::fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(path)?; + // SAFETY: `file` owns a valid fd for the duration of this call. + if unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX) } != 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(file) + } + + #[cfg(windows)] + fn acquire_exclusive_lock(path: &Path) -> std::io::Result { + use std::os::windows::fs::OpenOptionsExt; + + // `share_mode(0)` denies every other open on this path — read, + // write, and delete — until this handle closes. Windows closes it + // automatically, and only then, when the process exits by any + // means, so the lock is always released exactly when the process + // is actually gone. + std::fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .share_mode(0) + .open(path) + } + + /// Root command that waits for `start_path` to appear, then re-execs + /// this test binary's [`lock_holder_helper_entrypoint`] as a descendant + /// and blocks so the root itself stays alive until killed. + /// + /// The start-file wait exists only so tests can call [`attach`] before + /// the descendant is created — deterministically avoiding the + /// assign-race the module docs describe, which real callers accept but + /// a test must not be flaky about. + fn root_with_lock_holding_descendant( + temp: &Path, + lock_path: &Path, + descendant_ready: &Path, + start_path: &Path, + ) -> Command { + let this_test_binary = std::env::current_exe().expect("locate current test binary"); + // libtest's `--exact` filter matches the fully qualified test path, + // not the bare function name. + const HELPER_FILTER: &str = "process_tree::tests::lock_holder_helper_entrypoint"; + #[cfg(unix)] + let mut command = { + let mut command = Command::new("sh"); + command.args([ + "-c", + "while [ ! -f \"$START\" ]; do sleep 0.05; done; \ + \"$HELPER_BIN\" \"$HELPER_FILTER\" --exact --nocapture & wait", + ]); + command + }; + #[cfg(windows)] + let mut command = { + let mut command = Command::new("powershell.exe"); + command.args([ + "-NoLogo", + "-NoProfile", + "-NonInteractive", + "-Command", + "while (-not (Test-Path $env:START)) { Start-Sleep -Milliseconds 50 }; \ + Start-Process -FilePath $env:HELPER_BIN -ArgumentList @( \ + $env:HELPER_FILTER,'--exact','--nocapture') -NoNewWindow -Wait", + ]); + command + }; + configure(&mut command); + command + .env("HELPER_FILTER", HELPER_FILTER) + .current_dir(temp) + .env("HELPER_BIN", &this_test_binary) + .env("PROCESS_TREE_TEST_LOCK_PATH", lock_path) + .env("PROCESS_TREE_TEST_READY_PATH", descendant_ready) + .env("START", start_path) + // The re-exec'd descendant inherits stdio by default; nulling it + // here keeps this test's own output pipe from staying open for + // as long as that descendant (and, transitively, its own nested + // test harness output) is alive. + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + command + } + + async fn wait_for_file(path: &Path) { + let deadline = tokio::time::Instant::now() + TEST_TIMEOUT; + while !path.exists() { + assert!( + tokio::time::Instant::now() < deadline, + "expected file was never created: {}", + path.display() + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + async fn wait_until(mut predicate: impl FnMut() -> bool, message: &str) { + let deadline = tokio::time::Instant::now() + TEST_TIMEOUT; + while !predicate() { + assert!(tokio::time::Instant::now() < deadline, "{message}"); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + /// Attempts to acquire the same lock the descendant holds. Succeeds + /// only once the descendant process (or whatever kernel action ended + /// it) has released it — proving the resource, not merely the pid, is + /// gone. + fn lock_is_free(path: &Path) -> bool { + #[cfg(unix)] + { + use std::os::fd::AsRawFd; + + let Ok(file) = std::fs::OpenOptions::new().write(true).open(path) else { + return false; + }; + // SAFETY: `file` owns a valid fd for the duration of this call. + let acquired = + unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0; + if acquired { + // SAFETY: `file` still owns the fd we just locked. + unsafe { + libc::flock(file.as_raw_fd(), libc::LOCK_UN); + } + } + acquired + } + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + + std::fs::OpenOptions::new() + .write(true) + .share_mode(0) + .open(path) + .is_ok() + } + } + + #[tokio::test] + async fn terminate_kills_descendant_and_releases_its_lock() { + let temp = tempdir().expect("create temp directory"); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let start_path = temp.path().join("start"); + + let mut command = root_with_lock_holding_descendant( + temp.path(), + &lock_path, + &descendant_ready, + &start_path, + ); + let mut root = command.spawn().expect("spawn root process"); + let root_pid = root.id().expect("root pid"); + let tree = attach(&root).expect("attach process tree"); + std::fs::write(&start_path, b"go").expect("signal root to spawn its descendant"); + + wait_for_file(&descendant_ready).await; + let descendant_pid: u32 = std::fs::read_to_string(&descendant_ready) + .expect("read descendant pid") + .trim() + .parse() + .expect("parse descendant pid"); + + assert!(process_alive(root_pid), "root must be alive before stop"); + assert!( + process_alive(descendant_pid), + "descendant must be alive before stop" + ); + assert!( + !lock_is_free(&lock_path), + "descendant must be holding its lock before stop" + ); + + tree.terminate().expect("terminate process tree"); + tokio::time::timeout(TEST_TIMEOUT, root.wait()) + .await + .expect("root did not exit within the timeout after process-tree termination") + .expect("reap root process"); + + wait_until( + || !process_alive(descendant_pid), + "descendant survived process-tree termination", + ) + .await; + assert!( + !process_alive(root_pid), + "root survived its own termination" + ); + wait_until( + || lock_is_free(&lock_path), + "descendant's lock was never released after process-tree termination", + ) + .await; + assert!( + tree.is_empty().expect("query tree emptiness"), + "tree must report empty once every process in it has exited" + ); + } + + #[tokio::test] + async fn is_empty_is_false_while_descendant_holds_the_tree_open() { + let temp = tempdir().expect("create temp directory"); + let lock_path = temp.path().join("resource.lock"); + let descendant_ready = temp.path().join("descendant-ready"); + let start_path = temp.path().join("start"); + + let mut command = root_with_lock_holding_descendant( + temp.path(), + &lock_path, + &descendant_ready, + &start_path, + ); + let mut root = command.spawn().expect("spawn root process"); + let tree = attach(&root).expect("attach process tree"); + std::fs::write(&start_path, b"go").expect("signal root to spawn its descendant"); + + wait_for_file(&descendant_ready).await; + assert!( + !tree.is_empty().expect("query tree emptiness"), + "tree must report non-empty while the descendant is still alive" + ); + + tree.terminate().expect("terminate process tree"); + tokio::time::timeout(TEST_TIMEOUT, root.wait()) + .await + .expect("root did not exit within the timeout after process-tree termination") + .expect("reap root process"); + } +}