feat(transport): add builder & expose child stderr (#305)

This commit is contained in:
Brendan Graham 2025-07-04 18:14:43 -07:00 • committed by GitHub
commit 69712e321a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 122 additions and 32 deletions

View file

@ -1,24 +1,37 @@
use std::process::Stdio;
use process_wrap::tokio::{TokioChildWrapper, TokioCommandWrap}; use process_wrap::tokio::{TokioChildWrapper, TokioCommandWrap};
use tokio::{ use tokio::{
io::AsyncRead, io::AsyncRead,
process::{ChildStdin, ChildStdout}, process::{ChildStderr, ChildStdin, ChildStdout},
}; };
use super::{IntoTransport, Transport}; use super::{IntoTransport, Transport};
use crate::service::ServiceRole; use crate::service::ServiceRole;
pub(crate) fn child_process( /// The parts of a child process.
mut child: Box<dyn TokioChildWrapper>, type ChildProcessParts = (
) -> std::io::Result<(Box<dyn TokioChildWrapper>, (ChildStdout, ChildStdin))> { Box<dyn TokioChildWrapper>,
ChildStdout,
ChildStdin,
Option<ChildStderr>,
);
/// 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<dyn TokioChildWrapper>) -> std::io::Result<ChildProcessParts> {
let child_stdin = match child.inner_mut().stdin().take() { let child_stdin = match child.inner_mut().stdin().take() {
Some(stdin) => stdin, 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() { let child_stdout = match child.inner_mut().stdout().take() {
Some(stdout) => stdout, 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 { pub struct TokioChildProcess {
@ -66,31 +79,15 @@ impl AsyncRead for TokioChildProcessOut {
} }
impl TokioChildProcess { impl TokioChildProcess {
/// Create a new Tokio child process with the given command. /// Convenience: spawn with default `piped` stdio
///
/// # 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);
/// ```
///
pub fn new(command: impl Into<TokioCommandWrap>) -> std::io::Result<Self> { pub fn new(command: impl Into<TokioCommandWrap>) -> std::io::Result<Self> {
let mut command_wrap = command.into(); let (proc, _ignored) = TokioChildProcessBuilder::new(command).spawn()?;
command_wrap Ok(proc)
.command_mut() }
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped()); /// Builder entry-point allowing fine-grained stdio control.
let (child, (child_stdout, child_stdin)) = child_process(command_wrap.spawn()?)?; pub fn builder(command: impl Into<TokioCommandWrap>) -> TokioChildProcessBuilder {
Ok(Self { TokioChildProcessBuilder::new(command)
child: ChildWithCleanup { inner: child },
child_stdin,
child_stdout,
})
} }
/// Get the process ID of the child process. /// Get the process ID of the child process.
@ -98,6 +95,7 @@ impl TokioChildProcess {
self.child.inner.id() self.child.inner.id()
} }
/// Split this helper into a reader (stdout) and writer (stdin).
pub fn split(self) -> (TokioChildProcessOut, ChildStdin) { pub fn split(self) -> (TokioChildProcessOut, ChildStdin) {
let TokioChildProcess { let TokioChildProcess {
child, 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<TokioCommandWrap>) -> 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<Stdio>) -> Self {
self.stdin = io.into();
self
}
/// Override the child stdout configuration.
pub fn stdout(mut self, io: impl Into<Stdio>) -> Self {
self.stdout = io.into();
self
}
/// Override the child stderr configuration.
pub fn stderr(mut self, io: impl Into<Stdio>) -> 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<ChildStderr>)> {
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<R: ServiceRole> IntoTransport<R, std::io::Error, ()> for TokioChildProcess { impl<R: ServiceRole> IntoTransport<R, std::io::Error, ()> for TokioChildProcess {
fn into_transport(self) -> impl Transport<R, Error = std::io::Error> + 'static { fn into_transport(self) -> impl Transport<R, Error = std::io::Error> + 'static {
IntoTransport::<R, std::io::Error, super::async_rw::TransportAdapterAsyncRW>::into_transport( IntoTransport::<R, std::io::Error, super::async_rw::TransportAdapterAsyncRW>::into_transport(

View file

@ -1,9 +1,11 @@
use std::process::Stdio;
use axum::Router; use axum::Router;
use rmcp::{ use rmcp::{
ServiceExt, ServiceExt,
transport::{ConfigureCommandExt, SseServer, TokioChildProcess, sse_server::SseServerConfig}, transport::{ConfigureCommandExt, SseServer, TokioChildProcess, sse_server::SseServerConfig},
}; };
use tokio::time::timeout; use tokio::{io::AsyncReadExt, time::timeout};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
mod common; mod common;
@ -119,3 +121,35 @@ async fn test_with_python_server() -> anyhow::Result<()> {
client.cancel().await?; client.cancel().await?;
Ok(()) 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(())
}

View file

@ -1,7 +1,12 @@
from fastmcp import FastMCP from fastmcp import FastMCP
import sys
mcp = FastMCP("Demo") mcp = FastMCP("Demo")
print("server starting up...", file=sys.stderr)
@mcp.tool() @mcp.tool()
def add(a: int, b: int) -> int: def add(a: int, b: int) -> int:
"""Add two numbers""" """Add two numbers"""