modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

refactor: Transport trait and worker transport, and streamable http client with those new features. (#167)

Browse source 
* refactor: a transport trait and streamable http client

* test: test with js streamable http server

* refactor: separate sse stream connection

* fix(test): streamable http client test with js

* fix(test): wait longer time for js server startup
This commit is contained in:
4t145 2025-05-17 23:53:38 +08:00 • committed by GitHub
commit 9a771fb157
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2761 additions and 1315 deletions
Download patch file Download diff file Expand all files Collapse all files

140
crates/rmcp/src/service.rs
View file

@ -1,4 +1,4 @@
use futures::future::BoxFuture;
use futures::{FutureExt, future::BoxFuture};
use thiserror::Error;
use crate::{
@ -9,20 +9,26 @@ use crate::{
JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta, NumberOrString, ProgressToken,
RequestId, ServerJsonRpcMessage,
},
transport::IntoTransport,
transport::{IntoTransport, Transport},
};
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
mod client;
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use client::*;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
mod server;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use server::*;
#[cfg(feature = "tower")]
#[cfg_attr(docsrs, doc(cfg(feature = "tower")))]
mod tower;
use tokio_util::sync::{CancellationToken, DropGuard};
#[cfg(feature = "tower")]
#[cfg_attr(docsrs, doc(cfg(feature = "tower")))]
pub use tower::*;
use tracing::instrument;
#[derive(Error, Debug)]
@ -459,10 +465,11 @@ impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
#[derive(Debug)]
pub enum QuitReason {
Cancelled,
Closed,
JoinError(tokio::task::JoinError),
}
/// Request execution context
@ -523,7 +530,6 @@ where
T: IntoTransport<R, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
use futures::{SinkExt, StreamExt};
const SINK_PROXY_BUFFER_SIZE: usize = 64;
let (sink_proxy_tx, mut sink_proxy_rx) =
tokio::sync::mpsc::channel::<TxJsonRpcMessage<R>>(SINK_PROXY_BUFFER_SIZE);
@ -535,7 +541,8 @@ where
}
service.set_peer(peer.clone());
let mut local_responder_pool = HashMap::new();
let mut local_responder_pool =
HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::new();
let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
let shared_service = Arc::new(service);
// for return
@ -546,29 +553,42 @@ where
let serve_loop_ct = ct.child_token();
let peer_return: Peer<R> = peer.clone();
let handle = tokio::spawn(async move {
let (mut sink, mut stream) = transport.into_transport();
let mut sink = std::pin::pin!(sink);
let mut stream = std::pin::pin!(stream);
let mut transport = transport.into_transport();
let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
let mut send_task_set = tokio::task::JoinSet::<SendTaskResult<E>>::new();
#[derive(Debug)]
enum Event<P, R, T> {
ProxyMessage(P),
PeerMessage(R),
ToSink(T),
enum SendTaskResult<E> {
Request {
id: RequestId,
result: Result<(), E>,
},
Notification {
responder: Responder<Result<(), ServiceError>>,
cancellation_param: Option<CancelledNotificationParam>,
result: Result<(), E>,
},
}
#[derive(Debug)]
enum Event<R: ServiceRole, E> {
ProxyMessage(PeerSinkMessage<R>),
PeerMessage(RxJsonRpcMessage<R>),
ToSink(TxJsonRpcMessage<R>),
SendTaskResult(SendTaskResult<E>),
}
let quit_reason = loop {
let evt = if let Some(m) = batch_messages.pop_front() {
Event::PeerMessage(m)
} else {
tokio::select! {
m = sink_proxy_rx.recv() => {
m = sink_proxy_rx.recv(), if !sink_proxy_rx.is_closed() => {
if let Some(m) = m {
Event::ToSink(m)
} else {
continue
}
}
m = stream.next() => {
m = transport.receive() => {
if let Some(m) = m {
Event::PeerMessage(m)
} else {
@ -577,13 +597,28 @@ where
break QuitReason::Closed
}
}
m = peer_rx.recv() => {
m = peer_rx.recv(), if !peer_rx.is_closed() => {
if let Some(m) = m {
Event::ProxyMessage(m)
} else {
continue
}
}
m = send_task_set.join_next(), if !send_task_set.is_empty() => {
let Some(result) = m else {
continue
};
match result {
Err(e) => {
// join error, which is serious, we should quit.
tracing::error!(%e, "send request task encounter a tokio join error");
break QuitReason::JoinError(e)
}
Ok(result) => {
Event::SendTaskResult(result)
}
}
}
_ = serve_loop_ct.cancelled() => {
tracing::info!("task cancelled");
break QuitReason::Cancelled
@ -593,6 +628,34 @@ where
tracing::trace!(?evt, "new event");
match evt {
Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
if let Err(e) = result {
if let Some(responder) = local_responder_pool.remove(&id) {
let _ = responder
.send(Err(ServiceError::Transport(std::io::Error::other(e))));
}
}
}
Event::SendTaskResult(SendTaskResult::Notification {
responder,
result,
cancellation_param,
}) => {
let response = if let Err(e) = result {
Err(ServiceError::Transport(std::io::Error::other(e)))
} else {
Ok(())
};
let _ = responder.send(response);
if let Some(param) = cancellation_param {
if let Some(responder) = local_responder_pool.remove(&param.request_id) {
tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
let _response_result = responder.send(Err(ServiceError::Cancelled {
reason: param.reason.clone(),
}));
}
}
}
// response and error
Event::ToSink(m) => {
if let Some(id) = match &m {
@ -603,10 +666,13 @@ where
if let Some(ct) = local_ct_pool.remove(id) {
ct.cancel();
}
let send_result = sink.send(m).await;
if let Err(error) = send_result {
tracing::error!(%error, "fail to response message");
}
let send = transport.send(m);
tokio::spawn(async move {
let send_result = send.await;
if let Err(error) = send_result {
tracing::error!(%error, "fail to response message");
}
});
}
}
Event::ProxyMessage(PeerSinkMessage::Request {
@ -615,14 +681,11 @@ where
responder,
}) => {
local_responder_pool.insert(id.clone(), responder);
let send_result = sink
.send(JsonRpcMessage::request(request, id.clone()))
.await;
if let Err(e) = send_result {
if let Some(responder) = local_responder_pool.remove(&id) {
let _ = responder
.send(Err(ServiceError::Transport(std::io::Error::other(e))));
}
let send = transport.send(JsonRpcMessage::request(request, id.clone()));
{
let id = id.clone();
send_task_set
.spawn(send.map(move |r| SendTaskResult::Request { id, result: r }));
}
}
Event::ProxyMessage(PeerSinkMessage::Notification {
@ -638,21 +701,12 @@ where
}
Err(notification) => notification,
};
let send_result = sink.send(JsonRpcMessage::notification(notification)).await;
let response = if let Err(e) = send_result {
Err(ServiceError::Transport(std::io::Error::other(e)))
} else {
Ok(())
};
let _ = responder.send(response);
if let Some(param) = cancellation_param {
if let Some(responder) = local_responder_pool.remove(&param.request_id) {
tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
let _response_result = responder.send(Err(ServiceError::Cancelled {
reason: param.reason.clone(),
}));
}
}
let send = transport.send(JsonRpcMessage::notification(notification));
send_task_set.spawn(send.map(move |result| SendTaskResult::Notification {
responder,
cancellation_param,
result,
}));
}
Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest {
id, request, ..
@ -749,7 +803,7 @@ where
}
}
};
let sink_close_result = sink.close().await;
let sink_close_result = transport.close().await;
if let Err(e) = sink_close_result {
tracing::error!(%e, "fail to close sink");
}