fix: cleanup zombie processes for child process client (#156)

* feat(client): cleanup of zombie processes in child process client

Using process wrap library, Process Group for Unix and Job Object for Windows

* fix: install extra dep based on feature, update examples take owned command

Install process-wrap if transport-child-process feature is enabled. Update examples to take ownership
of the command instead of mutable reference

* fix: update other examples, comments and readme to take ownership of command

Updated more examples, comments and readme to take command instead of mutable ref.
Also small fix in cargo toml of example to use local path.

* refactor: add configure command ext

Added configure command ext to tokio process command so that
you can use .configure() to use inline commands for mcp stdio client.
Added warning if start kill process fails

Author: humbertoyusta

README.md

@@ -22,17 +22,17 @@ rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "mai
 
 ### Quick start
 
-Start a client in one line:
+Start a client:
 
 ```rust, ignore
-use rmcp::{ServiceExt, transport::TokioChildProcess};
+use rmcp::{ServiceExt, transport::{TokioChildProcess, ConfigureCommandExt}};
 use tokio::process::Command;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let client = ().serve(
-        TokioChildProcess::new(Command::new("npx").arg("-y").arg("@modelcontextprotocol/server-everything"))?
-    ).await?;
+    let client = ().serve(TokioChildProcess::new(Command::new("npx").configure(|cmd| {
+        cmd.arg("-y").arg("@modelcontextprotocol/server-everything");
+    }))?).await?;
     Ok(())
 }
 ```
@@ -92,7 +92,7 @@ let server = service.serve(transport).await?;
 Once the server is initialized, you can send requests or notifications:
 
 ```rust, ignore
-// request 
+// request
 let roots = server.list_roots().await?;
 
 // or send notification

crates/rmcp/Cargo.toml

@@ -45,6 +45,8 @@ url = { version = "2.4", optional = true }
 # For tower compatibility
 tower-service = { version = "0.3", optional = true }
 
+# for child process transport
+process-wrap = { version = "8.2", features = ["tokio1"], optional = true}
 
 # for ws transport
 # tokio-tungstenite ={ version = "0.26", optional = true }
@@ -89,7 +91,11 @@ transport-streamable-http-client = [
 
 transport-async-rw = ["tokio/io-util", "tokio-util/codec"]
 transport-io = ["transport-async-rw", "tokio/io-std"]
-transport-child-process = ["transport-async-rw", "tokio/process"]
+transport-child-process = [
+    "transport-async-rw",
+    "tokio/process",
+    "dep:process-wrap",
+]
 transport-sse-server = [
     "transport-async-rw",
     "axum",

crates/rmcp/src/lib.rs

@@ -58,15 +58,13 @@
 //!
 //! ```rust
 //! use anyhow::Result;
-//! use rmcp::{model::CallToolRequestParam, service::ServiceExt, transport::TokioChildProcess};
+//! use rmcp::{model::CallToolRequestParam, service::ServiceExt, transport::{TokioChildProcess, ConfigureCommandExt}};
 //! use tokio::process::Command;
 //!
 //! async fn client() -> Result<()> {
-//!     let service = ()
-//!       .serve(TokioChildProcess::new(
-//!             Command::new("uvx").arg("mcp-server-git"),
-//!         )?)
-//!         .await?;
+//!     let service = ().serve(TokioChildProcess::new(Command::new("uvx").configure(|cmd| {
+//!         cmd.arg("mcp-server-git");
+//!     }))?).await?;
 //!
 //!     // Initialize
 //!     let server_info = service.peer_info();

crates/rmcp/src/transport.rs

@@ -84,7 +84,7 @@ pub use worker::WorkerTransport;
 pub mod child_process;
 #[cfg(feature = "transport-child-process")]
 #[cfg_attr(docsrs, doc(cfg(feature = "transport-child-process")))]
-pub use child_process::TokioChildProcess;
+pub use child_process::{ConfigureCommandExt, TokioChildProcess};
 
 #[cfg(feature = "transport-io")]
 #[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]

crates/rmcp/src/transport/child_process.rs

@@ -1,3 +1,4 @@
+use process_wrap::tokio::{TokioChildWrapper, TokioCommandWrap};
 use tokio::{
     io::AsyncRead,
     process::{ChildStdin, ChildStdout},
@@ -7,29 +8,41 @@ use super::{IntoTransport, Transport};
 use crate::service::ServiceRole;
 
 pub(crate) fn child_process(
-    mut child: tokio::process::Child,
-) -> std::io::Result<(tokio::process::Child, (ChildStdout, ChildStdin))> {
-    if child.stdin.is_none() {
-        return Err(std::io::Error::other("std in was taken"));
-    }
-    if child.stdout.is_none() {
-        return Err(std::io::Error::other("std out was taken"));
-    }
-    let child_stdin = child.stdin.take().expect("already checked");
-    let child_stdout = child.stdout.take().expect("already checked");
+    mut child: Box<dyn TokioChildWrapper>,
+) -> std::io::Result<(Box<dyn TokioChildWrapper>, (ChildStdout, ChildStdin))> {
+    let child_stdin = match child.inner_mut().stdin().take() {
+        Some(stdin) => stdin,
+        None => return Err(std::io::Error::other("std in was taken")),
+    };
+    let child_stdout = match child.inner_mut().stdout().take() {
+        Some(stdout) => stdout,
+        None => return Err(std::io::Error::other("std out was taken")),
+    };
     Ok((child, (child_stdout, child_stdin)))
 }
 
 pub struct TokioChildProcess {
-    child: tokio::process::Child,
+    child: ChildWithCleanup,
     child_stdin: ChildStdin,
     child_stdout: ChildStdout,
 }
 
+pub struct ChildWithCleanup {
+    inner: Box<dyn TokioChildWrapper>,
+}
+
+impl Drop for ChildWithCleanup {
+    fn drop(&mut self) {
+        if let Err(e) = self.inner.start_kill() {
+            tracing::warn!("Failed to kill child process: {e}");
+        }
+    }
+}
+
 // we hold the child process with stdout, for it's easier to implement AsyncRead
 pin_project_lite::pin_project! {
     pub struct TokioChildProcessOut {
-        child: tokio::process::Child,
+        child: ChildWithCleanup,
         #[pin]
         child_stdout: ChildStdout,
     }
@@ -46,14 +59,18 @@ impl AsyncRead for TokioChildProcessOut {
 }
 
 impl TokioChildProcess {
-    pub fn new(child: &mut tokio::process::Command) -> std::io::Result<Self> {
-        child
-            .kill_on_drop(true)
+    pub fn new(mut command: tokio::process::Command) -> std::io::Result<Self> {
+        command
             .stdin(std::process::Stdio::piped())
             .stdout(std::process::Stdio::piped());
-        let (child, (child_stdout, child_stdin)) = child_process(child.spawn()?)?;
+        let mut command_wrap = TokioCommandWrap::from(command);
+        #[cfg(unix)]
+        command_wrap.wrap(process_wrap::tokio::ProcessGroup::leader());
+        #[cfg(windows)]
+        command_wrap.wrap(process_wrap::tokio::JobObject);
+        let (child, (child_stdout, child_stdin)) = child_process(command_wrap.spawn()?)?;
         Ok(Self {
-            child,
+            child: ChildWithCleanup { inner: child },
             child_stdin,
             child_stdout,
         })
@@ -82,3 +99,14 @@ impl<R: ServiceRole> IntoTransport<R, std::io::Error, ()> for TokioChildProcess
         )
     }
 }
+
+pub trait ConfigureCommandExt {
+    fn configure(self, f: impl FnOnce(&mut Self)) -> Self;
+}
+
+impl ConfigureCommandExt for tokio::process::Command {
+    fn configure(mut self, f: impl FnOnce(&mut Self)) -> Self {
+        f(&mut self);
+        self
+    }
+}

crates/rmcp/tests/test_with_js.rs

@@ -2,7 +2,7 @@ use rmcp::{
     ServiceExt,
     service::QuitReason,
     transport::{
-        SseServer, StreamableHttpClientTransport, TokioChildProcess,
+        ConfigureCommandExt, SseServer, StreamableHttpClientTransport, TokioChildProcess,
         streamable_http_server::axum::StreamableHttpServer,
     },
 };
@@ -59,9 +59,10 @@ async fn test_with_js_server() -> anyhow::Result<()> {
         .spawn()?
         .wait()
         .await?;
-    let transport = TokioChildProcess::new(
-        tokio::process::Command::new("node").arg("tests/test_with_js/server.js"),
-    )?;
+    let transport =
+        TokioChildProcess::new(tokio::process::Command::new("node").configure(|cmd| {
+            cmd.arg("tests/test_with_js/server.js");
+        }))?;
 
     let client = ().serve(transport).await?;
     let resources = client.list_all_resources().await?;

crates/rmcp/tests/test_with_python.rs

@@ -1,6 +1,6 @@
 use rmcp::{
     ServiceExt,
-    transport::{SseServer, TokioChildProcess},
+    transport::{ConfigureCommandExt, SseServer, TokioChildProcess},
 };
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 mod common;
@@ -54,11 +54,9 @@ async fn test_with_python_server() -> anyhow::Result<()> {
         .spawn()?
         .wait()
         .await?;
-    let transport = TokioChildProcess::new(
-        tokio::process::Command::new("uv")
-            .arg("run")
-            .arg("tests/test_with_python/server.py"),
-    )?;
+    let transport = TokioChildProcess::new(tokio::process::Command::new("uv").configure(|cmd| {
+        cmd.arg("run").arg("tests/test_with_python/server.py");
+    }))?;
 
     let client = ().serve(transport).await?;
     let resources = client.list_all_resources().await?;

docs/readme/README.zh-cn.md

@@ -17,12 +17,12 @@ rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "mai
 ### 快速上手
 一行代码启动客户端：
 ```rust
-use rmcp::{ServiceExt, transport::TokioChildProcess};
+use rmcp::{ServiceExt, transport::{TokioChildProcess, ConfigureCommandExt}};
 use tokio::process::Command;
 
-let client = ().serve(
-    TokioChildProcess::new(Command::new("npx").arg("-y").arg("@modelcontextprotocol/server-everything"))?
-).await?;
+let client = ().serve(TokioChildProcess::new(Command::new("npx").configure(|cmd| {
+    cmd.arg("-y").arg("@modelcontextprotocol/server-everything");
+}))?).await?;
 ```
 
 #### 1. 构建传输层
@@ -63,7 +63,7 @@ let server = service.serve(transport).await?;
 一旦服务初始化完成，你可以发送请求或通知：
 
 ```rust, ignore
-// 请求 
+// 请求
 let roots = server.list_roots().await?;
 
 // 或发送通知

examples/clients/src/collection.rs

@@ -1,7 +1,11 @@
 use std::collections::HashMap;
 
 use anyhow::Result;
-use rmcp::{model::CallToolRequestParam, service::ServiceExt, transport::TokioChildProcess};
+use rmcp::{
+    model::CallToolRequestParam,
+    service::ServiceExt,
+    transport::{ConfigureCommandExt, TokioChildProcess},
+};
 use tokio::process::Command;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
@@ -20,9 +24,11 @@ async fn main() -> Result<()> {
     for idx in 0..10 {
         let service = ()
             .into_dyn()
-            .serve(TokioChildProcess::new(
-                Command::new("uvx").arg("mcp-server-git"),
-            )?)
+            .serve(TokioChildProcess::new(Command::new("uvx").configure(
+                |cmd| {
+                    cmd.arg("mcp-client-git");
+                },
+            ))?)
             .await?;
         client_list.insert(idx, service);
     }

examples/clients/src/everything_stdio.rs

@@ -3,7 +3,7 @@ use rmcp::{
     ServiceExt,
     model::{CallToolRequestParam, GetPromptRequestParam, ReadResourceRequestParam},
     object,
-    transport::TokioChildProcess,
+    transport::{ConfigureCommandExt, TokioChildProcess},
 };
 use tokio::process::Command;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
@@ -21,11 +21,11 @@ async fn main() -> Result<()> {
 
     // Start server
     let service = ()
-        .serve(TokioChildProcess::new(
-            Command::new("npx")
-                .arg("-y")
-                .arg("@modelcontextprotocol/server-everything"),
-        )?)
+        .serve(TokioChildProcess::new(Command::new("npx").configure(
+            |cmd| {
+                cmd.arg("-y").arg("@modelcontextprotocol/server-everything");
+            },
+        ))?)
         .await?;
 
     // Initialize