diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 2f495912f..1b8198476 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -9,13 +9,10 @@ use rmcp::service::{ClientInitializeError, ServiceError}; use rmcp::transport::streamable_http_client::{ StreamableHttpClientTransportConfig, StreamableHttpError, }; -use rmcp::transport::{ - ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport, TokioChildProcess, -}; +use rmcp::transport::{ConfigureCommandExt, DynamicTransportError, StreamableHttpClientTransport}; use std::collections::HashMap; use std::path::PathBuf; use std::pin::Pin; -use std::process::Stdio; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Weak}; use std::task::{Context, Poll}; @@ -47,7 +44,7 @@ use crate::config::search_path::SearchPaths; use crate::config::{get_all_extensions, Config}; use crate::oauth::{oauth_flow, GooseCredentialStore, StaticOAuthClientConfig}; use crate::prompt_template; -use crate::subprocess::configure_subprocess; +use crate::subprocess::spawn_long_lived_mcp_subprocess; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, ErrorCode, ErrorData, GetPromptResult, MetaObject, Prompt, Resource, ResourceContents, ServerInfo, ServerNotification, Tool, @@ -433,8 +430,6 @@ async fn child_process_client( action_required: Arc, extension_manager: Weak, ) -> ExtensionResult { - configure_subprocess(&mut command); - if let Ok(path) = SearchPaths::builder().path() { command.env("PATH", path); } @@ -449,9 +444,7 @@ async fn child_process_client( ); } - let (transport, mut stderr) = TokioChildProcess::builder(command) - .stderr(Stdio::piped()) - .spawn()?; + let (transport, mut stderr) = spawn_long_lived_mcp_subprocess(command).await?; let mut stderr = stderr.take().ok_or_else(|| { ExtensionError::SetupError("failed to attach child process stderr".to_owned()) })?; diff --git a/crates/goose/src/subprocess.rs b/crates/goose/src/subprocess.rs index 1849e55c6..cf5b56ca5 100644 --- a/crates/goose/src/subprocess.rs +++ b/crates/goose/src/subprocess.rs @@ -1,3 +1,8 @@ +use rmcp::transport::TokioChildProcess; +use std::io; +#[cfg(target_os = "linux")] +use std::sync::{mpsc, OnceLock}; +use tokio::process::ChildStderr; use tokio::process::Command; #[cfg(windows)] @@ -60,13 +65,80 @@ impl SubprocessExt for std::process::Command { } } -#[allow(unused_variables)] -pub fn configure_subprocess(command: &mut Command) { +fn configure_common_subprocess(command: &mut Command) { // Isolate subprocess into its own process group so it does not receive // SIGINT when the user presses Ctrl+C in the terminal. #[cfg(unix)] command.process_group(0); - #[cfg(target_os = "linux")] - configure_parent_death_signal(command); command.set_no_window(); } + +#[allow(unused_variables)] +pub fn configure_subprocess(command: &mut Command) { + configure_common_subprocess(command); + #[cfg(target_os = "linux")] + configure_parent_death_signal(command); +} + +#[cfg(target_os = "linux")] +struct LongLivedSpawnRequest { + command: Command, + runtime: tokio::runtime::Handle, + response: tokio::sync::oneshot::Sender)>>, +} + +#[cfg(target_os = "linux")] +fn long_lived_spawn_sender() -> io::Result> { + static SENDER: OnceLock>> = OnceLock::new(); + + match SENDER.get_or_init(|| { + let (sender, receiver) = mpsc::channel::(); + std::thread::Builder::new() + .name("goose-extension-spawner".to_owned()) + .spawn(move || { + while let Ok(mut request) = receiver.recv() { + let _runtime_guard = request.runtime.enter(); + configure_subprocess(&mut request.command); + let result = TokioChildProcess::builder(request.command) + .stderr(std::process::Stdio::piped()) + .spawn(); + let _ = request.response.send(result); + } + }) + .map(|_| sender) + }) { + Ok(sender) => Ok(sender.clone()), + Err(error) => Err(io::Error::new(error.kind(), error.to_string())), + } +} + +/// Spawn a long-lived MCP subprocess without tying Linux parent-death cleanup +/// to the Tokio worker that happened to request it. +pub async fn spawn_long_lived_mcp_subprocess( + command: Command, +) -> io::Result<(TokioChildProcess, Option)> { + #[cfg(target_os = "linux")] + { + let runtime = tokio::runtime::Handle::try_current().map_err(io::Error::other)?; + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + long_lived_spawn_sender()? + .send(LongLivedSpawnRequest { + command, + runtime, + response: response_tx, + }) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "extension spawner exited"))?; + response_rx + .await + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "extension spawner exited"))? + } + + #[cfg(not(target_os = "linux"))] + { + let mut command = command; + configure_subprocess(&mut command); + TokioChildProcess::builder(command) + .stderr(std::process::Stdio::piped()) + .spawn() + } +} diff --git a/crates/goose/tests/subprocess_cleanup.rs b/crates/goose/tests/subprocess_cleanup.rs index ae58d421a..d731ab94a 100644 --- a/crates/goose/tests/subprocess_cleanup.rs +++ b/crates/goose/tests/subprocess_cleanup.rs @@ -1,15 +1,30 @@ #![cfg(target_os = "linux")] -use goose::subprocess::configure_subprocess; +use goose::subprocess::{configure_subprocess, spawn_long_lived_mcp_subprocess}; use std::io::{BufRead, BufReader, Write}; use std::path::PathBuf; -use std::process::{Command, Stdio}; +use std::process::{Child, Command, Stdio}; +use std::sync::mpsc; use std::time::{Duration, Instant}; const HELPER_ENV: &str = "GOOSE_SUBPROCESS_PARENT_DEATH_HELPER"; +const THREAD_HELPER_ENV: &str = "GOOSE_SUBPROCESS_THREAD_DEATH_HELPER"; + +struct HelperProcess(Child); + +impl Drop for HelperProcess { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} #[ctor::ctor] unsafe fn maybe_run_helper() { + if std::env::var_os(THREAD_HELPER_ENV).is_some() { + run_thread_death_helper(); + } + if std::env::var_os(HELPER_ENV).is_none() { return; } @@ -41,6 +56,49 @@ unsafe fn maybe_run_helper() { } } +fn run_thread_death_helper() { + let (tx, rx) = mpsc::channel(); + let spawn_thread = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + + let pid = runtime.block_on(async { + let mut command = tokio::process::Command::new("sleep"); + command.arg("30"); + command.stdin(Stdio::null()); + command.stdout(Stdio::null()); + command.stderr(Stdio::null()); + let (child, _) = spawn_long_lived_mcp_subprocess(command) + .await + .expect("spawn child"); + let pid = child.id().expect("child pid"); + std::mem::forget(child); + pid + }); + + tx.send(pid).expect("send child pid"); + }); + + spawn_thread.join().expect("spawn thread"); + let child_pid = rx.recv().expect("child pid"); + std::thread::sleep(Duration::from_millis(500)); + + if !process_is_running(child_pid) { + eprintln!("child process {child_pid} exited after spawning thread exit"); + unsafe { + libc::_exit(1); + } + } + + println!("{child_pid}"); + std::io::stdout().flush().expect("flush pid"); + loop { + std::thread::park(); + } +} + #[test] fn child_process_exits_when_parent_process_dies() { let current_exe = std::env::current_exe().expect("current test binary"); @@ -74,6 +132,71 @@ fn child_process_exits_when_parent_process_dies() { } } +#[test] +fn long_lived_child_process_survives_spawning_thread_exit() { + let current_exe = std::env::current_exe().expect("current test binary"); + let mut helper = HelperProcess( + Command::new(current_exe) + .env(THREAD_HELPER_ENV, "1") + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn helper"), + ); + + let stdout = helper.0.stdout.take().expect("helper stdout"); + let (pid_tx, pid_rx) = mpsc::channel(); + std::thread::spawn(move || { + let pid = BufReader::new(stdout) + .lines() + .next() + .ok_or("helper exited without reporting a child pid") + .and_then(|line| line.map_err(|_| "failed to read helper child pid")) + .and_then(|line| line.parse::().map_err(|_| "invalid helper child pid")); + let _ = pid_tx.send(pid); + }); + let child_pid = pid_rx + .recv_timeout(Duration::from_secs(5)) + .expect("timed out waiting for helper child pid") + .expect("helper child pid"); + + assert!( + process_is_running(child_pid), + "child process {child_pid} exited after spawning thread exit" + ); + + unsafe { + libc::kill(helper.0.id() as libc::pid_t, libc::SIGKILL); + } + let status = helper.0.wait().expect("wait for helper"); + assert!( + !status.success(), + "helper should have been killed: {status}" + ); + + let deadline = Instant::now() + Duration::from_secs(5); + while process_is_running(child_pid) && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(100)); + } + assert!( + !process_is_running(child_pid), + "child process {child_pid} survived parent process death" + ); +} + fn process_exists(pid: u32) -> bool { PathBuf::from(format!("/proc/{pid}")).exists() } + +fn process_is_running(pid: u32) -> bool { + match process_state(pid) { + Some('Z') | None => false, + Some(_) => true, + } +} + +fn process_state(pid: u32) -> Option { + let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?; + let (_, after_name) = stat.rsplit_once(") ")?; + after_name.chars().next() +}