fix(mcp): keep stdio extensions alive across worker exits (#10364)

Signed-off-by: iroiro147 <sarthak.singh@juspay.in>

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<ActionRequiredManager>,
     extension_manager: Weak<ExtensionManager>,
 ) -> ExtensionResult<McpClient> {
-    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())
     })?;

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<io::Result<(TokioChildProcess, Option<ChildStderr>)>>,
+}
+
+#[cfg(target_os = "linux")]
+fn long_lived_spawn_sender() -> io::Result<mpsc::Sender<LongLivedSpawnRequest>> {
+    static SENDER: OnceLock<io::Result<mpsc::Sender<LongLivedSpawnRequest>>> = OnceLock::new();
+
+    match SENDER.get_or_init(|| {
+        let (sender, receiver) = mpsc::channel::<LongLivedSpawnRequest>();
+        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<ChildStderr>)> {
+    #[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()
+    }
+}

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::<u32>().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<char> {
+    let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
+    let (_, after_name) = stat.rsplit_once(") ")?;
+    after_name.chars().next()
+}