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

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

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");
+    })
+}