fix: cleanup MCP processes when CLI closes (#2469)

Co-authored-by: Alice Hau <ahau@squareup.com>

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"

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"] }

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(())
+        }
+    }
 }

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;
 

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<Box<dyn Future<Output = ()> + 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<Box<dyn Future<Output = ()> + Send>> {
+    Box::pin(async move {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    })
+}

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]

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<TransportMessage>,
+    receiver: Option<mpsc::Receiver<TransportMessage>>,
     pending_requests: Arc<PendingRequests>,
-    _process: Child, // we store the process to keep it alive
+    process: Child, // we store the process to keep it alive
     error_sender: mpsc::Sender<Error>,
-    stdin: ChildStdin,
-    stdout: ChildStdout,
-    stderr: ChildStderr,
+    stdin: Option<ChildStdin>,
+    stdout: Option<ChildStdout>,
+    stderr: Option<ChildStderr>,
+}
+
+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(())
     }
 }