From 64995cc40c8b0208bda67fc4bf0d34226d5f5be5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Humberto=20Yusta=20G=C3=B3mez?= Date: Sun, 18 May 2025 12:38:47 +0200 Subject: [PATCH] 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 --- README.md | 12 ++--- crates/rmcp/Cargo.toml | 8 ++- crates/rmcp/src/lib.rs | 10 ++-- crates/rmcp/src/transport.rs | 2 +- crates/rmcp/src/transport/child_process.rs | 62 ++++++++++++++++------ crates/rmcp/tests/test_with_js.rs | 9 ++-- crates/rmcp/tests/test_with_python.rs | 10 ++-- docs/readme/README.zh-cn.md | 10 ++-- examples/clients/src/collection.rs | 14 +++-- examples/clients/src/everything_stdio.rs | 12 ++--- examples/clients/src/std_io.rs | 21 ++++---- examples/rig-integration/src/config/mcp.rs | 9 ++-- examples/simple-chat-client/src/config.rs | 13 ++--- 13 files changed, 115 insertions(+), 77 deletions(-) diff --git a/README.md b/README.md index 9912015..d238454 100644 --- a/README.md +++ b/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> { - 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 diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 37818ea..e9aa78b 100644 --- a/crates/rmcp/Cargo.toml +++ b/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", diff --git a/crates/rmcp/src/lib.rs b/crates/rmcp/src/lib.rs index 65e48fc..e4096b4 100644 --- a/crates/rmcp/src/lib.rs +++ b/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(); diff --git a/crates/rmcp/src/transport.rs b/crates/rmcp/src/transport.rs index b699e41..8f47501 100644 --- a/crates/rmcp/src/transport.rs +++ b/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")))] diff --git a/crates/rmcp/src/transport/child_process.rs b/crates/rmcp/src/transport/child_process.rs index 20ef6be..f2e3c66 100644 --- a/crates/rmcp/src/transport/child_process.rs +++ b/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, +) -> std::io::Result<(Box, (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, +} + +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 { - child - .kill_on_drop(true) + pub fn new(mut command: tokio::process::Command) -> std::io::Result { + 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 IntoTransport 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 + } +} diff --git a/crates/rmcp/tests/test_with_js.rs b/crates/rmcp/tests/test_with_js.rs index 57c11d9..50f4293 100644 --- a/crates/rmcp/tests/test_with_js.rs +++ b/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?; diff --git a/crates/rmcp/tests/test_with_python.rs b/crates/rmcp/tests/test_with_python.rs index 7e0fbd7..f8fa56e 100644 --- a/crates/rmcp/tests/test_with_python.rs +++ b/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?; diff --git a/docs/readme/README.zh-cn.md b/docs/readme/README.zh-cn.md index 6061385..cc730c1 100644 --- a/docs/readme/README.zh-cn.md +++ b/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?; // 或发送通知 diff --git a/examples/clients/src/collection.rs b/examples/clients/src/collection.rs index f474ea8..244d369 100644 --- a/examples/clients/src/collection.rs +++ b/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); } diff --git a/examples/clients/src/everything_stdio.rs b/examples/clients/src/everything_stdio.rs index 091e905..a2a7475 100644 --- a/examples/clients/src/everything_stdio.rs +++ b/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 diff --git a/examples/clients/src/std_io.rs b/examples/clients/src/std_io.rs index 9874189..0e04763 100644 --- a/examples/clients/src/std_io.rs +++ b/examples/clients/src/std_io.rs @@ -1,5 +1,9 @@ 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}; @@ -14,17 +18,14 @@ async fn main() -> Result<()> { .with(tracing_subscriber::fmt::layer()) .init(); let service = () - .serve(TokioChildProcess::new( - Command::new("uvx").arg("mcp-server-git"), - )?) + .serve(TokioChildProcess::new(Command::new("uvx").configure( + |cmd| { + cmd.arg("mcp-server-git"); + }, + ))?) .await?; - // or - // serve_client( - // (), - // TokioChildProcess::new(Command::new("uvx").arg("mcp-server-git"))?, - // ) - // .await?; + // or serve_client((), TokioChildProcess::new(cmd)?).await?; // Initialize let server_info = service.peer_info(); diff --git a/examples/rig-integration/src/config/mcp.rs b/examples/rig-integration/src/config/mcp.rs index 95f412a..17663ca 100644 --- a/examples/rig-integration/src/config/mcp.rs +++ b/examples/rig-integration/src/config/mcp.rs @@ -1,6 +1,6 @@ use std::{collections::HashMap, process::Stdio}; -use rmcp::{RoleClient, ServiceExt, service::RunningService}; +use rmcp::{RoleClient, ServiceExt, service::RunningService, transport::ConfigureCommandExt}; use serde::{Deserialize, Serialize}; use crate::mcp_adaptor::McpManager; @@ -70,10 +70,9 @@ impl McpServerTransportConfig { envs, } => { let transport = rmcp::transport::TokioChildProcess::new( - tokio::process::Command::new(command) - .args(args) - .envs(envs) - .stderr(Stdio::null()), + tokio::process::Command::new(command).configure(|cmd| { + cmd.args(args).envs(envs).stderr(Stdio::null()); + }), )?; ().serve(transport).await? } diff --git a/examples/simple-chat-client/src/config.rs b/examples/simple-chat-client/src/config.rs index 356fd49..fc15989 100644 --- a/examples/simple-chat-client/src/config.rs +++ b/examples/simple-chat-client/src/config.rs @@ -1,7 +1,7 @@ use std::{collections::HashMap, path::Path, process::Stdio}; use anyhow::Result; -use rmcp::{RoleClient, ServiceExt, service::RunningService}; +use rmcp::{RoleClient, ServiceExt, service::RunningService, transport::ConfigureCommandExt}; use serde::{Deserialize, Serialize}; #[derive(Debug, Serialize, Deserialize)] @@ -55,11 +55,12 @@ impl McpServerTransportConfig { envs, } => { let transport = rmcp::transport::child_process::TokioChildProcess::new( - tokio::process::Command::new(command) - .args(args) - .envs(envs) - .stderr(Stdio::inherit()) - .stdout(Stdio::inherit()), + tokio::process::Command::new(command).configure(|cmd| { + cmd.args(args) + .envs(envs) + .stderr(Stdio::inherit()) + .stdout(Stdio::inherit()); + }), )?; ().serve(transport).await? }