From 69712e321a3946176058740ebfa0f92f641557fa Mon Sep 17 00:00:00 2001 From: Brendan Graham Date: Fri, 4 Jul 2025 18:14:43 -0700 Subject: [PATCH] feat(transport): add builder & expose child stderr (#305) --- crates/rmcp/src/transport/child_process.rs | 113 ++++++++++++++----- crates/rmcp/tests/test_with_python.rs | 36 +++++- crates/rmcp/tests/test_with_python/server.py | 5 + 3 files changed, 122 insertions(+), 32 deletions(-) diff --git a/crates/rmcp/src/transport/child_process.rs b/crates/rmcp/src/transport/child_process.rs index 391a6ce..2e7c034 100644 --- a/crates/rmcp/src/transport/child_process.rs +++ b/crates/rmcp/src/transport/child_process.rs @@ -1,24 +1,37 @@ +use std::process::Stdio; + use process_wrap::tokio::{TokioChildWrapper, TokioCommandWrap}; use tokio::{ io::AsyncRead, - process::{ChildStdin, ChildStdout}, + process::{ChildStderr, ChildStdin, ChildStdout}, }; use super::{IntoTransport, Transport}; use crate::service::ServiceRole; -pub(crate) fn child_process( - mut child: Box, -) -> std::io::Result<(Box, (ChildStdout, ChildStdin))> { +/// The parts of a child process. +type ChildProcessParts = ( + Box, + ChildStdout, + ChildStdin, + Option, +); + +/// Extract the stdio handles from a spawned child. +/// Returns `(child, stdout, stdin, stderr)` where `stderr` is `Some` only +/// if the process was spawned with `Stdio::piped()`. +#[inline] +fn child_process(mut child: Box) -> std::io::Result { let child_stdin = match child.inner_mut().stdin().take() { Some(stdin) => stdin, - None => return Err(std::io::Error::other("std in was taken")), + None => return Err(std::io::Error::other("stdin was already taken")), }; let child_stdout = match child.inner_mut().stdout().take() { Some(stdout) => stdout, - None => return Err(std::io::Error::other("std out was taken")), + None => return Err(std::io::Error::other("stdout was already taken")), }; - Ok((child, (child_stdout, child_stdin))) + let child_stderr = child.inner_mut().stderr().take(); + Ok((child, child_stdout, child_stdin, child_stderr)) } pub struct TokioChildProcess { @@ -66,31 +79,15 @@ impl AsyncRead for TokioChildProcessOut { } impl TokioChildProcess { - /// Create a new Tokio child process with the given command. - /// - /// # Manage the child process - /// You can also check these issue and pr for more information on how to manage the child process: - /// - [#156](https://github.com/modelcontextprotocol/rust-sdk/pull/156) - /// - [#253](https://github.com/modelcontextprotocol/rust-sdk/issues/253) - /// ```rust,ignore - /// #[cfg(unix)] - /// command_wrap.wrap(process_wrap::tokio::ProcessGroup::leader()); - /// #[cfg(windows)] - /// command_wrap.wrap(process_wrap::tokio::JobObject); - /// ``` - /// + /// Convenience: spawn with default `piped` stdio pub fn new(command: impl Into) -> std::io::Result { - let mut command_wrap = command.into(); - command_wrap - .command_mut() - .stdin(std::process::Stdio::piped()) - .stdout(std::process::Stdio::piped()); - let (child, (child_stdout, child_stdin)) = child_process(command_wrap.spawn()?)?; - Ok(Self { - child: ChildWithCleanup { inner: child }, - child_stdin, - child_stdout, - }) + let (proc, _ignored) = TokioChildProcessBuilder::new(command).spawn()?; + Ok(proc) + } + + /// Builder entry-point allowing fine-grained stdio control. + pub fn builder(command: impl Into) -> TokioChildProcessBuilder { + TokioChildProcessBuilder::new(command) } /// Get the process ID of the child process. @@ -98,6 +95,7 @@ impl TokioChildProcess { self.child.inner.id() } + /// Split this helper into a reader (stdout) and writer (stdin). pub fn split(self) -> (TokioChildProcessOut, ChildStdin) { let TokioChildProcess { child, @@ -114,6 +112,59 @@ impl TokioChildProcess { } } +/// Builder for `TokioChildProcess` allowing custom `Stdio` configuration. +pub struct TokioChildProcessBuilder { + cmd: TokioCommandWrap, + stdin: Stdio, + stdout: Stdio, + stderr: Stdio, +} + +impl TokioChildProcessBuilder { + fn new(cmd: impl Into) -> Self { + Self { + cmd: cmd.into(), + stdin: Stdio::piped(), + stdout: Stdio::piped(), + stderr: Stdio::inherit(), + } + } + + /// Override the child stdin configuration. + pub fn stdin(mut self, io: impl Into) -> Self { + self.stdin = io.into(); + self + } + /// Override the child stdout configuration. + pub fn stdout(mut self, io: impl Into) -> Self { + self.stdout = io.into(); + self + } + /// Override the child stderr configuration. + pub fn stderr(mut self, io: impl Into) -> Self { + self.stderr = io.into(); + self + } + + /// Spawn the child process. Returns the transport plus an optional captured stderr handle. + pub fn spawn(mut self) -> std::io::Result<(TokioChildProcess, Option)> { + self.cmd + .command_mut() + .stdin(self.stdin) + .stdout(self.stdout) + .stderr(self.stderr); + + let (child, stdout, stdin, stderr_opt) = child_process(self.cmd.spawn()?)?; + + let proc = TokioChildProcess { + child: ChildWithCleanup { inner: child }, + child_stdin: stdin, + child_stdout: stdout, + }; + Ok((proc, stderr_opt)) + } +} + impl IntoTransport for TokioChildProcess { fn into_transport(self) -> impl Transport + 'static { IntoTransport::::into_transport( diff --git a/crates/rmcp/tests/test_with_python.rs b/crates/rmcp/tests/test_with_python.rs index 9def172..8971a32 100644 --- a/crates/rmcp/tests/test_with_python.rs +++ b/crates/rmcp/tests/test_with_python.rs @@ -1,9 +1,11 @@ +use std::process::Stdio; + use axum::Router; use rmcp::{ ServiceExt, transport::{ConfigureCommandExt, SseServer, TokioChildProcess, sse_server::SseServerConfig}, }; -use tokio::time::timeout; +use tokio::{io::AsyncReadExt, time::timeout}; use tokio_util::sync::CancellationToken; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; mod common; @@ -119,3 +121,35 @@ async fn test_with_python_server() -> anyhow::Result<()> { client.cancel().await?; Ok(()) } + +#[tokio::test] +async fn test_with_python_server_stderr() -> anyhow::Result<()> { + init().await?; + + let (transport, stderr) = + TokioChildProcess::builder(tokio::process::Command::new("uv").configure(|cmd| { + cmd.arg("run") + .arg("server.py") + .current_dir("tests/test_with_python"); + })) + .stderr(Stdio::piped()) + .spawn()?; + + let mut stderr = stderr.expect("stderr must be piped"); + + let stderr_task = tokio::spawn(async move { + let mut buffer = String::new(); + stderr.read_to_string(&mut buffer).await?; + Ok::<_, std::io::Error>(buffer) + }); + + let client = ().serve(transport).await?; + let _ = client.list_all_resources().await?; + let _ = client.list_all_tools().await?; + client.cancel().await?; + + let stderr_output = stderr_task.await??; + assert!(stderr_output.contains("server starting up...")); + + Ok(()) +} diff --git a/crates/rmcp/tests/test_with_python/server.py b/crates/rmcp/tests/test_with_python/server.py index 985cc33..1c8ee69 100644 --- a/crates/rmcp/tests/test_with_python/server.py +++ b/crates/rmcp/tests/test_with_python/server.py @@ -1,7 +1,12 @@ from fastmcp import FastMCP +import sys + mcp = FastMCP("Demo") +print("server starting up...", file=sys.stderr) + + @mcp.tool() def add(a: int, b: int) -> int: """Add two numbers"""