From 85dd6375b5fe0fe42131ce2101e1c349456249c9 Mon Sep 17 00:00:00 2001 From: Alice Hau <110418948+ahau-square@users.noreply.github.com> Date: Thu, 8 May 2025 10:19:59 -0400 Subject: [PATCH] fix: cleanup MCP processes when CLI closes (#2469) Co-authored-by: Alice Hau --- Cargo.lock | 18 +++- crates/goose-cli/Cargo.toml | 1 + crates/goose-cli/src/commands/mcp.rs | 39 +++++++- crates/goose-cli/src/lib.rs | 1 + crates/goose-cli/src/signal.rs | 36 ++++++++ crates/mcp-client/Cargo.toml | 1 + crates/mcp-client/src/transport/stdio.rs | 110 ++++++++++++++++------- 7 files changed, 173 insertions(+), 33 deletions(-) create mode 100644 crates/goose-cli/src/signal.rs diff --git a/Cargo.lock b/Cargo.lock index fe08495d12..e6aeead22c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2485,6 +2485,7 @@ dependencies = [ "mcp-core", "mcp-server", "minijinja", + "nix 0.30.1", "once_cell", "rand", "regex", @@ -3488,9 +3489,9 @@ checksum = "03087c2bad5e1034e8cace5926dec053fb3790248370865f5117a7d0213354c8" [[package]] name = "libc" -version = "0.2.170" +version = "0.2.172" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "875b3680cb2f8f71bdcf9a30f38d48282f5d3c95cbf9b3fa57269bb5d5c06828" +checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" [[package]] name = "libdbus-sys" @@ -3700,6 +3701,7 @@ dependencies = [ "eventsource-client", "futures", "mcp-core", + "nix 0.30.1", "rand", "reqwest 0.11.27", "serde", @@ -3926,6 +3928,18 @@ dependencies = [ "libc", ] +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags 2.9.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" diff --git a/crates/goose-cli/Cargo.toml b/crates/goose-cli/Cargo.toml index cf9adb517d..43c5f2a4b6 100644 --- a/crates/goose-cli/Cargo.toml +++ b/crates/goose-cli/Cargo.toml @@ -53,6 +53,7 @@ async-trait = "0.1.86" base64 = "0.22.1" regex = "1.11.1" minijinja = "2.8.0" +nix = { version = "0.30.1", features = ["process", "signal"] } [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["wincred"] } diff --git a/crates/goose-cli/src/commands/mcp.rs b/crates/goose-cli/src/commands/mcp.rs index 791fe5adc8..168dab9169 100644 --- a/crates/goose-cli/src/commands/mcp.rs +++ b/crates/goose-cli/src/commands/mcp.rs @@ -7,6 +7,16 @@ use mcp_server::router::RouterService; use mcp_server::{BoundedService, ByteTransport, Server}; use tokio::io::{stdin, stdout}; +use std::sync::Arc; +use tokio::sync::Notify; + +#[cfg(unix)] +use nix::sys::signal::{kill, Signal}; +#[cfg(unix)] +use nix::unistd::getpgrp; +#[cfg(unix)] +use nix::unistd::Pid; + pub async fn run_server(name: &str) -> Result<()> { // Initialize logging crate::logging::setup_logging(Some(&format!("mcp-{name}")), None)?; @@ -26,10 +36,37 @@ pub async fn run_server(name: &str) -> Result<()> { _ => None, }; + // Create shutdown notification channel + let shutdown = Arc::new(Notify::new()); + let shutdown_clone = shutdown.clone(); + + // Spawn shutdown signal handler + tokio::spawn(async move { + crate::signal::shutdown_signal().await; + shutdown_clone.notify_one(); + }); + // Create and run the server let server = Server::new(router.unwrap_or_else(|| panic!("Unknown server requested {}", name))); let transport = ByteTransport::new(stdin(), stdout()); tracing::info!("Server initialized and ready to handle requests"); - Ok(server.run(transport).await?) + + tokio::select! { + result = server.run(transport) => { + Ok(result?) + } + _ = shutdown.notified() => { + // On Unix systems, kill the entire process group + #[cfg(unix)] + fn terminate_process_group() { + let pgid = getpgrp(); + kill(Pid::from_raw(-pgid.as_raw()), Signal::SIGTERM) + .expect("Failed to send SIGTERM to process group"); + } + terminate_process_group(); + + Ok(()) + } + } } diff --git a/crates/goose-cli/src/lib.rs b/crates/goose-cli/src/lib.rs index 3be1393551..2cf95d0671 100644 --- a/crates/goose-cli/src/lib.rs +++ b/crates/goose-cli/src/lib.rs @@ -6,6 +6,7 @@ pub mod logging; pub mod recipe; pub mod recipes; pub mod session; +pub mod signal; // Re-export commonly used types pub use session::Session; diff --git a/crates/goose-cli/src/signal.rs b/crates/goose-cli/src/signal.rs new file mode 100644 index 0000000000..bdf05dfa49 --- /dev/null +++ b/crates/goose-cli/src/signal.rs @@ -0,0 +1,36 @@ +use std::future::Future; +use std::pin::Pin; +use tokio::signal; + +#[cfg(unix)] +pub fn shutdown_signal() -> Pin + Send>> { + Box::pin(async move { + let ctrl_c = async { + signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } + }) +} + +#[cfg(not(unix))] +pub fn shutdown_signal() -> Pin + Send>> { + Box::pin(async move { + signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }) +} diff --git a/crates/mcp-client/Cargo.toml b/crates/mcp-client/Cargo.toml index d6699fd56d..914fd31da3 100644 --- a/crates/mcp-client/Cargo.toml +++ b/crates/mcp-client/Cargo.toml @@ -20,5 +20,6 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } tower = { version = "0.4", features = ["timeout", "util"] } tower-service = "0.3" rand = "0.8" +nix = { version = "0.30.1", features = ["process", "signal"] } [dev-dependencies] diff --git a/crates/mcp-client/src/transport/stdio.rs b/crates/mcp-client/src/transport/stdio.rs index 7980816bf0..5895e83e1e 100644 --- a/crates/mcp-client/src/transport/stdio.rs +++ b/crates/mcp-client/src/transport/stdio.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::sync::atomic::{AtomicI32, Ordering}; use std::sync::Arc; use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command}; @@ -7,31 +8,58 @@ use mcp_core::protocol::JsonRpcMessage; use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; use tokio::sync::{mpsc, Mutex}; +// Import nix crate components instead of libc +#[cfg(unix)] +use nix::sys::signal::{kill, Signal}; +#[cfg(unix)] +use nix::unistd::{getpgid, Pid}; + use super::{send_message, Error, PendingRequests, Transport, TransportHandle, TransportMessage}; +// Global to track process groups we've created +static PROCESS_GROUP: AtomicI32 = AtomicI32::new(-1); + /// A `StdioTransport` uses a child process's stdin/stdout as a communication channel. /// /// It uses channels for message passing and handles responses asynchronously through a background task. pub struct StdioActor { - receiver: mpsc::Receiver, + receiver: Option>, pending_requests: Arc, - _process: Child, // we store the process to keep it alive + process: Child, // we store the process to keep it alive error_sender: mpsc::Sender, - stdin: ChildStdin, - stdout: ChildStdout, - stderr: ChildStderr, + stdin: Option, + stdout: Option, + stderr: Option, +} + +impl Drop for StdioActor { + fn drop(&mut self) { + // Get the process group ID before attempting cleanup + #[cfg(unix)] + if let Some(pid) = self.process.id() { + if let Ok(pgid) = getpgid(Some(Pid::from_raw(pid as i32))) { + // Send SIGTERM to the entire process group + let _ = kill(Pid::from_raw(-pgid.as_raw()), Signal::SIGTERM); + // Give processes a moment to cleanup + std::thread::sleep(std::time::Duration::from_millis(100)); + // Force kill if still running + let _ = kill(Pid::from_raw(-pgid.as_raw()), Signal::SIGKILL); + } + } + } } impl StdioActor { pub async fn run(mut self) { use tokio::pin; - let incoming = Self::handle_incoming_messages(self.stdout, self.pending_requests.clone()); - let outgoing = Self::handle_outgoing_messages( - self.receiver, - self.stdin, - self.pending_requests.clone(), - ); + let stdout = self.stdout.take().expect("stdout should be available"); + let stdin = self.stdin.take().expect("stdin should be available"); + let receiver = self.receiver.take().expect("receiver should be available"); + + let incoming = Self::handle_incoming_messages(stdout, self.pending_requests.clone()); + let outgoing = + Self::handle_outgoing_messages(receiver, stdin, self.pending_requests.clone()); // take ownership of futures for tokio::select pin!(incoming); @@ -46,25 +74,27 @@ impl StdioActor { tracing::debug!("Stdout handler completed: {:?}", result); } // capture the status so we don't need to wait for a timeout - status = self._process.wait() => { + status = self.process.wait() => { tracing::debug!("Process exited with status: {:?}", status); } } // Then always try to read stderr before cleaning up let mut stderr_buffer = Vec::new(); - if let Ok(bytes) = self.stderr.read_to_end(&mut stderr_buffer).await { - let err_msg = if bytes > 0 { - String::from_utf8_lossy(&stderr_buffer).to_string() - } else { - "Process ended unexpectedly".to_string() - }; + if let Some(mut stderr) = self.stderr.take() { + if let Ok(bytes) = stderr.read_to_end(&mut stderr_buffer).await { + let err_msg = if bytes > 0 { + String::from_utf8_lossy(&stderr_buffer).to_string() + } else { + "Process ended unexpectedly".to_string() + }; - tracing::info!("Process stderr: {}", err_msg); - let _ = self - .error_sender - .send(Error::StdioProcessError(err_msg)) - .await; + tracing::info!("Process stderr: {}", err_msg); + let _ = self + .error_sender + .send(Error::StdioProcessError(err_msg)) + .await; + } } // Clean up regardless of which path we took @@ -213,9 +243,9 @@ impl StdioTransport { .stderr(std::process::Stdio::piped()) .kill_on_drop(true); - // Set process group only on Unix systems + // Set process group and ensure signal handling on Unix systems #[cfg(unix)] - command.process_group(0); // don't inherit signal handling from parent process + command.process_group(0); // Hide console window on Windows #[cfg(windows)] @@ -240,6 +270,15 @@ impl StdioTransport { .take() .ok_or_else(|| Error::StdioProcessError("Failed to get stderr".into()))?; + // Store the process group ID for cleanup + #[cfg(unix)] + if let Some(pid) = process.id() { + // Use nix instead of unsafe libc calls + if let Ok(pgid) = getpgid(Some(Pid::from_raw(pid as i32))) { + PROCESS_GROUP.store(pgid.as_raw(), Ordering::SeqCst); + } + } + Ok((process, stdin, stdout, stderr)) } } @@ -254,13 +293,13 @@ impl Transport for StdioTransport { let (error_tx, error_rx) = mpsc::channel(1); let actor = StdioActor { - receiver: message_rx, + receiver: Some(message_rx), pending_requests: Arc::new(PendingRequests::new()), - _process: process, + process, error_sender: error_tx, - stdin, - stdout, - stderr, + stdin: Some(stdin), + stdout: Some(stdout), + stderr: Some(stderr), }; tokio::spawn(actor.run()); @@ -273,6 +312,17 @@ impl Transport for StdioTransport { } async fn close(&self) -> Result<(), Error> { + // Attempt to clean up the process group on close + #[cfg(unix)] + if let Some(pgid) = PROCESS_GROUP.load(Ordering::SeqCst).checked_abs() { + // Use nix instead of unsafe libc calls + // Try SIGTERM first + let _ = kill(Pid::from_raw(-pgid), Signal::SIGTERM); + // Give processes a moment to cleanup + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + // Force kill if still running + let _ = kill(Pid::from_raw(-pgid), Signal::SIGKILL); + } Ok(()) } }