feat(transport): add builder & expose child stderr (#305)
This commit is contained in:
parent
31155a440c
commit
69712e321a
3 changed files with 122 additions and 32 deletions
|
|
@ -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(
|
||||||
|
|
|
||||||
|
|
@ -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(())
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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"""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue