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

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
41 changed files with 2761 additions and 1315 deletions
Download patch file Download diff file Expand all files Collapse all files

7
.github/workflows/ci.yml vendored
View file

@ -157,15 +157,16 @@ jobs:
- uses: actions/checkout@v3
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
uses: dtolnay/rust-toolchain@nightly
- uses: Swatinem/rust-cache@v2
- name: Generate documentation
run: |
cargo doc --no-deps -p rmcp -p rmcp-macros
cargo +nightly doc --no-deps -p rmcp -p rmcp-macros --all-features
env:
RUSTDOCFLAGS: -Dwarnings
RUSTDOCFLAGS: --cfg docsrs -Dwarnings
RUSTFLAGS: --cfg docsrs
release:
name: Release crates

30
README.md
View file

@ -45,18 +45,24 @@ use tokio::io::{stdin, stdout};
let transport = (stdin(), stdout());
```
The transport type must implemented [`IntoTransport`](crate::transport::IntoTransport) trait, which allow split into a sink and a stream.
#### Transport
The transport type must implemented [`Transport`] trait, which allow it send message concurrently and receive message sequentially.
There are 3 pairs of standard transport types:
For client, the sink item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage) and stream item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage)
| transport | client | server |
|:-: |:-: |:-: |
| std IO | [`child_process::TokioChildProcess`] | [`io::stdio`] |
| streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::session::create_session`] |
| sse | [`sse_client::SseClientTransport`] | [`sse_server::SseServer`] |
For server, the sink item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage) and stream item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage)
#### [IntoTransport](`IntoTransport`) trait
[`IntoTransport`] is a helper trait that implicitly convert a type into a transport type.
##### These types is automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
1. The types that already implement both [`Sink`](futures::Sink) and [`Stream`](futures::Stream) trait.
2. A tuple of sink `Tx` and stream `Rx`: `(Tx, Rx)`.
3. The type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait.
4. A tuple of [`tokio::io::AsyncRead`] `R `and [`tokio::io::AsyncWrite`] `W`: `(R, W)`.
These types is automatically implemented [`IntoTransport`] trait
1. A type that already implement both [`futures::Sink`] and [`futures::Stream`] trait, or a tuple `(Tx, Rx)` where `Tx` is [`futures::Sink`] and `Rx` is [`futures::Stream`].
2. A type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait. or a tuple `(R, W)` where `R` is [`tokio::io::AsyncRead`] and `W` is [`tokio::io::AsyncWrite`].
3. A type that implement [Worker](`worker::Worker`) trait.
4. A type that implement [`Transport`] trait.
For example, you can see how we build a transport through TCP stream or http upgrade so easily. [examples](examples/README.md)
</details>
@ -194,13 +200,15 @@ See [examples](examples/README.md)
- `macros`: macros default
- `schemars`: implement `JsonSchema` for all model structs
### Transports
#### Transports
- `transport-io`: Server stdio transport
- `transport-sse-server`: Server SSE transport
- `transport-child-process`: Client stdio transport
- `transport-sse`: Client sse transport
- `transport-sse-client`: Client sse transport
- `transport-streamable-http-server` streamable http server transport
- `transport-streamable-client-server` streamable http server transport
## Related Resources

24
crates/rmcp/Cargo.toml
View file

@ -11,6 +11,7 @@ documentation = "https://docs.rs/rmcp"
[package.metadata.docs.rs]
all-features = true
rustdoc-args = ["--cfg", "docsrs"]
[dependencies]
serde = { version = "1.0", features = ["derive", "rc"] }
@ -38,7 +39,7 @@ reqwest = { version = "0.12", default-features = false, features = [
"json",
"stream",
], optional = true }
sse-stream = { version = "0.1.3", optional = true }
sse-stream = { version = "0.1.4", optional = true }
url = { version = "2.4", optional = true }
# For tower compatibility
@ -62,9 +63,26 @@ default = ["base64", "macros", "server"]
client = []
server = ["transport-async-rw", "dep:schemars"]
macros = ["dep:rmcp-macros", "dep:paste"]
__transport-sse = ["dep:reqwest", "dep:sse-stream", "dep:url"]
# reqwest http client
reqwest = ["dep:reqwest"]
# SSE client
client-side-sse = ["dep:sse-stream"]
__transport-sse = ["reqwest", "client-side-sse", "dep:url", "transport-worker"]
transport-worker = ["dep:tokio-stream"]
transport-sse = ["__transport-sse", "reqwest?/rustls-tls"]
transport-sse-tls-no-provider = ["__transport-sse", "reqwest?/rustls-tls-no-provider"]
# Streamable HTTP client
__transport-streamable-http-client = ["reqwest", "client-side-sse", "transport-worker", "dep:url"]
transport-streamable-http-client = ["__transport-streamable-http-client", "reqwest?/rustls-tls"]
transport-streamable-http-client-tls-no-provider = ["__transport-streamable-http-client", "reqwest?/rustls-tls-no-provider"]
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"]
@ -114,7 +132,7 @@ path = "tests/test_with_python.rs"
[[test]]
name = "test_with_js"
required-features = ["server", "client", "transport-sse-server", "transport-child-process", "transport-streamable-http-server"]
required-features = ["server", "client", "transport-sse-server", "transport-child-process", "transport-streamable-http-server", "transport-streamable-http-client"]
path = "tests/test_with_js.rs"
[[test]]

2
crates/rmcp/src/handler.rs
View file

@ -1,4 +1,6 @@
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub mod client;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub mod server;

1
crates/rmcp/src/handler/server/tool.rs
View file

@ -431,6 +431,7 @@ impl<S> ToolBox<S> {
}
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
#[macro_export]
macro_rules! tool_box {
(@pin_add $callee: ident, $attr: expr, $f: expr) => {

13
crates/rmcp/src/lib.rs
View file

@ -1,3 +1,5 @@
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(docsrs, allow(unused_attributes))]
//! The official Rust SDK for the Model Context Protocol (MCP).
//!
//! The MCP is a protocol that allows AI assistants to communicate with other
@ -93,16 +95,22 @@ pub use error::Error;
/// Basic data types in MCP specification
pub mod model;
#[cfg(any(feature = "client", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "client", feature = "server"))))]
pub mod service;
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use handler::client::ClientHandler;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use handler::server::ServerHandler;
#[cfg(any(feature = "client", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "client", feature = "server"))))]
pub use service::{Peer, Service, ServiceError, ServiceExt};
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use service::{RoleClient, serve_client};
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use service::{RoleServer, serve_server};
pub mod handler;
@ -110,12 +118,17 @@ pub mod transport;
// re-export
#[cfg(all(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(all(feature = "macros", feature = "server"))))]
pub use paste::paste;
#[cfg(all(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(all(feature = "macros", feature = "server"))))]
pub use rmcp_macros::tool;
#[cfg(all(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(all(feature = "macros", feature = "server"))))]
pub use schemars;
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
pub use serde;
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
pub use serde_json;

1
crates/rmcp/src/model.rs
View file

@ -36,6 +36,7 @@ pub fn object(value: serde_json::Value) -> JsonObject {
/// Use this macro just like [`serde_json::json!`]
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
#[macro_export]
macro_rules! object {
({$($tt:tt)*}) => {

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");
}

40
crates/rmcp/src/service/client.rs
View file

@ -1,4 +1,3 @@
use futures::{SinkExt, Stream, StreamExt};
use thiserror::Error;
use super::*;
@ -38,29 +37,29 @@ pub enum ClientError {
}
/// Helper function to get the next message from the stream
async fn expect_next_message<S>(
stream: &mut S,
async fn expect_next_message<T>(
transport: &mut T,
context: &str,
) -> Result<ServerJsonRpcMessage, ClientError>
where
S: Stream<Item = ServerJsonRpcMessage> + Unpin,
T: Transport<RoleClient>,
{
stream
.next()
transport
.receive()
.await
.ok_or_else(|| ClientError::ConnectionClosed(context.to_string()))
.map_err(|e| ClientError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))
}
/// Helper function to expect a response from the stream
async fn expect_response<S>(
stream: &mut S,
async fn expect_response<T>(
transport: &mut T,
context: &str,
) -> Result<(ServerResult, RequestId), ClientError>
where
S: Stream<Item = ServerJsonRpcMessage> + Unpin,
T: Transport<RoleClient>,
{
let msg = expect_next_message(stream, context).await?;
let msg = expect_next_message(transport, context).await?;
match msg {
ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => Ok((result, id)),
@ -123,9 +122,7 @@ where
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
let (sink, stream) = transport.into_transport();
let mut sink = Box::pin(sink);
let mut stream = Box::pin(stream);
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// Convert ClientError to std::io::Error, then to E
@ -143,13 +140,14 @@ where
params: service.get_info(),
extensions: Default::default(),
};
sink.send(ClientJsonRpcMessage::request(
ClientRequest::InitializeRequest(init_request),
id.clone(),
))
.await?;
transport
.send(ClientJsonRpcMessage::request(
ClientRequest::InitializeRequest(init_request),
id.clone(),
))
.await?;
let (response, response_id) = expect_response(&mut stream, "initialize response")
let (response, response_id) = expect_response(&mut transport, "initialize response")
.await
.map_err(handle_client_error)?;
@ -173,9 +171,9 @@ where
extensions: Default::default(),
}),
);
sink.send(notification).await?;
transport.send(notification).await?;
let (peer, peer_rx) = Peer::new(id_provider, initialize_result);
serve_inner(service, (sink, stream), peer, peer_rx, ct).await
serve_inner(service, transport, peer, peer_rx, ct).await
}
macro_rules! method {

51
crates/rmcp/src/service/server.rs
View file

@ -1,4 +1,3 @@
use futures::{SinkExt, StreamExt};
use thiserror::Error;
use super::*;
@ -81,42 +80,42 @@ where
}
/// Helper function to get the next message from the stream
async fn expect_next_message<S>(
stream: &mut S,
async fn expect_next_message<T>(
transport: &mut T,
context: &str,
) -> Result<ClientJsonRpcMessage, ServerError>
where
S: StreamExt<Item = ClientJsonRpcMessage> + Unpin,
T: Transport<RoleServer>,
{
stream
.next()
transport
.receive()
.await
.ok_or_else(|| ServerError::ConnectionClosed(context.to_string()))
}
/// Helper function to expect a request from the stream
async fn expect_request<S>(
stream: &mut S,
async fn expect_request<T>(
transport: &mut T,
context: &str,
) -> Result<(ClientRequest, RequestId), ServerError>
where
S: StreamExt<Item = ClientJsonRpcMessage> + Unpin,
T: Transport<RoleServer>,
{
let msg = expect_next_message(stream, context).await?;
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_request()
.ok_or(ServerError::ExpectedInitRequest(Some(msg_clone)))
}
/// Helper function to expect a notification from the stream
async fn expect_notification<S>(
stream: &mut S,
async fn expect_notification<T>(
transport: &mut T,
context: &str,
) -> Result<ClientNotification, ServerError>
where
S: StreamExt<Item = ClientJsonRpcMessage> + Unpin,
T: Transport<RoleServer>,
{
let msg = expect_next_message(stream, context).await?;
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_notification()
.ok_or(ServerError::ExpectedInitNotification(Some(msg_clone)))
@ -132,9 +131,7 @@ where
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
{
let (sink, stream) = transport.into_transport();
let mut sink = Box::pin(sink);
let mut stream = Box::pin(stream);
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// Convert ServerError to std::io::Error, then to E
@ -146,7 +143,7 @@ where
};
// Get initialize request
let (request, id) = expect_request(&mut stream, "initialized request")
let (request, id) = expect_request(&mut transport, "initialized request")
.await
.map_err(handle_server_error)?;
@ -173,7 +170,8 @@ where
));
}
Err(e) => {
sink.send(ServerJsonRpcMessage::error(e.clone(), id))
transport
.send(ServerJsonRpcMessage::error(e.clone(), id))
.await?;
return Err(handle_server_error(ServerError::InitializeFailed(e)));
}
@ -190,14 +188,15 @@ where
_ => init_response.protocol_version,
};
init_response.protocol_version = protocol_version;
sink.send(ServerJsonRpcMessage::response(
ServerResult::InitializeResult(init_response),
id,
))
.await?;
transport
.send(ServerJsonRpcMessage::response(
ServerResult::InitializeResult(init_response),
id,
))
.await?;
// Wait for initialize notification
let notification = expect_notification(&mut stream, "initialize notification")
let notification = expect_notification(&mut transport, "initialize notification")
.await
.map_err(handle_server_error)?;
let ClientNotification::InitializedNotification(_) = notification else {
@ -207,7 +206,7 @@ where
};
let _ = service.handle_notification(notification).await;
// Continue processing service
serve_inner(service, (sink, stream), peer, peer_rx, ct).await
serve_inner(service, transport, peer, peer_rx, ct).await
}
macro_rules! method {

161
crates/rmcp/src/transport.rs
View file

@ -1,15 +1,41 @@
//! # Transport
//! The transport type must implemented [`IntoTransport`] trait, which allow split into a sink and a stream.
//! The transport type must implemented [`Transport`] trait, which allow it send message concurrently and receive message sequentially.
//！
//! ## Standard Transport Types
//! There are 3 pairs of standard transport types:
//!
//! For client, the sink item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage) and stream item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage)
//! | transport | client | server |
//! |:-: |:-: |:-: |
//! | std IO | [`child_process::TokioChildProcess`] | [`io::stdio`] |
//! | streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::session::create_session`] |
//! | sse | [`sse_client::SseClientTransport`] | [`sse_server::SseServer`] |
//!
//! For server, the sink item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage) and stream item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage)
//！## Helper Transport Types
//! Thers are several helper transport types that can help you to create transport quickly.
//!
//! ## These types is automatically implemented [`IntoTransport`] trait
//! 1. For type that already implement both [`Sink`] and [`Stream`] trait, they are automatically implemented [`IntoTransport`] trait
//! 2. For tuple of sink `Tx` and stream `Rx`, type `(Tx, Rx)` are automatically implemented [`IntoTransport`] trait
//! 3. For type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait, they are automatically implemented [`IntoTransport`] trait
//! 4. For tuple of [`tokio::io::AsyncRead`] `R `and [`tokio::io::AsyncWrite`] `W`, type `(R, W)` are automatically implemented [`IntoTransport`] trait
//! ### [Worker Transport](`worker::WorkerTransport`)
//! Which allows you to run a worker and process messages in another tokio task.
//!
//! ### [Async Read/Write Transport](`async_rw::AsyncRwTransport`)
//! You need to enable `transport-async-rw` feature to use this transport.
//!
//! This transport is used to create a transport from a byte stream which implemented [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`].
//!
//! This could be very helpful when you want to create a transport from a byte stream, such as a file or a tcp connection.
//!
//! ### [Sink/Stream Transport](`sink_stream::SinkStreamTransport`)
//! This transport is used to create a transport from a sink and a stream.
//!
//! This could be very helpful when you want to create a transport from a duplex object stream, such as a websocket connection.
//!
//! ## [IntoTransport](`IntoTransport`) trait
//! [`IntoTransport`] is a helper trait that implicitly convert a type into a transport type.
//!
//! ### These types is automatically implemented [`IntoTransport`] trait
//! 1. A type that already implement both [`futures::Sink`] and [`futures::Stream`] trait, or a tuple `(Tx, Rx)` where `Tx` is [`futures::Sink`] and `Rx` is [`futures::Stream`].
//! 2. A type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait. or a tuple `(R, W)` where `R` is [`tokio::io::AsyncRead`] and `W` is [`tokio::io::AsyncWrite`].
//! 3. A type that implement [Worker](`worker::Worker`) trait.
//! 4. A type that implement [`Transport`] trait.
//!
//! ## Examples
//!
@ -38,97 +64,114 @@
//! }
//! ```
use futures::{Sink, Stream};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
#[cfg(feature = "transport-child-process")]
pub mod child_process;
#[cfg(feature = "transport-child-process")]
pub use child_process::TokioChildProcess;
pub mod sink_stream;
#[cfg(feature = "transport-async-rw")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-async-rw")))]
pub mod async_rw;
#[cfg(feature = "transport-worker")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-worker")))]
pub mod worker;
#[cfg(feature = "transport-worker")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-worker")))]
pub use worker::WorkerTransport;
#[cfg(feature = "transport-child-process")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-child-process")))]
pub mod child_process;
#[cfg(feature = "transport-child-process")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-child-process")))]
pub use child_process::TokioChildProcess;
#[cfg(feature = "transport-io")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]
pub mod io;
#[cfg(feature = "transport-io")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]
pub use io::stdio;
#[cfg(feature = "__transport-sse")]
pub mod sse;
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-sse")))]
pub mod sse_client;
#[cfg(feature = "__transport-sse")]
pub use sse::SseTransport;
#[cfg(all(feature = "__transport-sse", feature = "__auth"))]
pub mod sse_auth;
#[cfg(all(feature = "__transport-sse", feature = "__auth"))]
pub use sse_auth::{AuthorizedSseClient, create_authorized_transport};
// #[cfg(feature = "tower")]
// pub mod tower;
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-sse")))]
pub use sse_client::SseClientTransport;
#[cfg(feature = "transport-sse-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-server")))]
pub mod sse_server;
#[cfg(feature = "transport-sse-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-sse-server")))]
pub use sse_server::SseServer;
#[cfg(feature = "__auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "__auth")))]
pub mod auth;
#[cfg(feature = "__auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "__auth")))]
pub use auth::{AuthError, AuthorizationManager, AuthorizationSession, AuthorizedHttpClient};
// #[cfg(feature = "transport-ws")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-ws")))]
// pub mod ws;
#[cfg(feature = "transport-streamable-http-server-session")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server-session")))]
pub mod streamable_http_server;
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub use streamable_http_server::axum::StreamableHttpServer;
#[cfg(feature = "__transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-streamable-http-client")))]
pub mod streamable_http_client;
#[cfg(feature = "__transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-streamable-http-client")))]
pub use streamable_http_client::StreamableHttpClientTransport;
/// Common use codes
pub mod common;
pub trait Transport<R>: Send
where
R: ServiceRole,
{
type Error;
/// Send a message to the transport
///
/// Notice that the future returned by this function should be `Send` and `'static`.
/// It's because the sending message could be executed concurrently.
///
fn send(
&mut self,
item: TxJsonRpcMessage<R>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static;
/// Receive a message from the transport, this operation is sequential.
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<R>>> + Send;
/// Close the transport
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
pub trait IntoTransport<R, E, A>: Send + 'static
where
R: ServiceRole,
E: std::error::Error + Send + 'static,
{
fn into_transport(
self,
) -> (
impl Sink<TxJsonRpcMessage<R>, Error = E> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<R>> + Send + 'static,
);
fn into_transport(self) -> impl Transport<R, Error = E> + 'static;
}
pub enum TransportAdapterStreamSink {}
impl<Role, Rx, Tx, E> IntoTransport<Role, E, TransportAdapterStreamSink> for (Tx, Rx)
pub enum TransportAdapterIdentity {}
impl<R, T, E> IntoTransport<R, E, TransportAdapterIdentity> for T
where
Role: ServiceRole,
Tx: Sink<TxJsonRpcMessage<Role>, Error = E> + Send + 'static,
Rx: Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
E: std::error::Error + Send + 'static,
{
fn into_transport(
self,
) -> (
impl Sink<TxJsonRpcMessage<Role>, Error = E> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
) {
self
}
}
impl<R, T, E> IntoTransport<R, E, ()> for T
where
T: Stream<Item = RxJsonRpcMessage<R>> + Sink<TxJsonRpcMessage<R>, Error = E> + Send + 'static,
T: Transport<R, Error = E> + Send + 'static,
R: ServiceRole,
E: std::error::Error + Send + 'static,
{
fn into_transport(
self,
) -> (
impl Sink<TxJsonRpcMessage<R>, Error = E> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<R>> + Send + 'static,
) {
use futures::StreamExt;
self.split()
fn into_transport(self) -> impl Transport<R, Error = E> + 'static {
self
}
}

382
crates/rmcp/src/transport/async_rw.rs Normal file
View file

@ -0,0 +1,382 @@
use std::{marker::PhantomData, sync::Arc};
// use crate::schema::*;
use futures::{SinkExt, StreamExt};
use serde::{Serialize, de::DeserializeOwned};
use thiserror::Error;
use tokio::{
io::{AsyncRead, AsyncWrite},
sync::Mutex,
};
use tokio_util::{
bytes::{Buf, BufMut, BytesMut},
codec::{Decoder, Encoder, FramedRead, FramedWrite},
};
use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub enum TransportAdapterAsyncRW {}
impl<Role, R, W> IntoTransport<Role, std::io::Error, TransportAdapterAsyncRW> for (R, W)
where
Role: ServiceRole,
R: AsyncRead + Send + 'static + Unpin,
W: AsyncWrite + Send + 'static + Unpin,
{
fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
AsyncRwTransport::new(self.0, self.1)
}
}
pub enum TransportAdapterAsyncCombinedRW {}
impl<Role, S> IntoTransport<Role, std::io::Error, TransportAdapterAsyncCombinedRW> for S
where
Role: ServiceRole,
S: AsyncRead + AsyncWrite + Send + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
IntoTransport::<Role, std::io::Error, TransportAdapterAsyncRW>::into_transport(
tokio::io::split(self),
)
}
}
pub struct AsyncRwTransport<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
read: FramedRead<R, JsonRpcMessageCodec<RxJsonRpcMessage<Role>>>,
write: Arc<Mutex<FramedWrite<W, JsonRpcMessageCodec<TxJsonRpcMessage<Role>>>>>,
}
impl<Role: ServiceRole, R, W> AsyncRwTransport<Role, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new(read: R, write: W) -> Self {
let read = FramedRead::new(
read,
JsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::default(),
);
let write = Arc::new(Mutex::new(FramedWrite::new(
write,
JsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::default(),
)));
Self { read, write }
}
}
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
impl<R, W> AsyncRwTransport<crate::RoleClient, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new_client(read: R, write: W) -> Self {
Self::new(read, write)
}
}
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
impl<R, W> AsyncRwTransport<crate::RoleServer, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new_server(read: R, write: W) -> Self {
Self::new(read, write)
}
}
impl<Role: ServiceRole, R, W> Transport<Role> for AsyncRwTransport<Role, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
type Error = std::io::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<Role>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let lock = self.write.clone();
async move {
let mut write = lock.lock().await;
write.send(item).await.map_err(Into::into)
}
}
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
let next = self.read.next();
async {
next.await.and_then(|e| {
e.inspect_err(|e| {
tracing::error!("Error reading from stream: {}", e);
})
.ok()
})
}
}
async fn close(&mut self) -> Result<(), Self::Error> {
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct JsonRpcMessageCodec<T> {
_marker: PhantomData<fn() -> T>,
next_index: usize,
max_length: usize,
is_discarding: bool,
}
impl<T> Default for JsonRpcMessageCodec<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> JsonRpcMessageCodec<T> {
pub fn new() -> Self {
Self {
_marker: PhantomData,
next_index: 0,
max_length: usize::MAX,
is_discarding: false,
}
}
pub fn new_with_max_length(max_length: usize) -> Self {
Self {
max_length,
..Self::new()
}
}
pub fn max_length(&self) -> usize {
self.max_length
}
}
fn without_carriage_return(s: &[u8]) -> &[u8] {
if let Some(&b'\r') = s.last() {
&s[..s.len() - 1]
} else {
s
}
}
#[derive(Debug, Error)]
pub enum JsonRpcMessageCodecError {
#[error("max line length exceeded")]
MaxLineLengthExceeded,
#[error("serde error {0}")]
Serde(#[from] serde_json::Error),
#[error("io error {0}")]
Io(#[from] std::io::Error),
}
impl From<JsonRpcMessageCodecError> for std::io::Error {
fn from(value: JsonRpcMessageCodecError) -> Self {
match value {
JsonRpcMessageCodecError::MaxLineLengthExceeded => {
std::io::Error::new(std::io::ErrorKind::InvalidData, value)
}
JsonRpcMessageCodecError::Serde(e) => e.into(),
JsonRpcMessageCodecError::Io(e) => e,
}
}
}
impl<T: DeserializeOwned> Decoder for JsonRpcMessageCodec<T> {
type Item = T;
type Error = JsonRpcMessageCodecError;
fn decode(
&mut self,
buf: &mut BytesMut,
) -> Result<Option<Self::Item>, JsonRpcMessageCodecError> {
loop {
// Determine how far into the buffer we'll search for a newline. If
// there's no max_length set, we'll read to the end of the buffer.
let read_to = std::cmp::min(self.max_length.saturating_add(1), buf.len());
let newline_offset = buf[self.next_index..read_to]
.iter()
.position(|b| *b == b'\n');
match (self.is_discarding, newline_offset) {
(true, Some(offset)) => {
// If we found a newline, discard up to that offset and
// then stop discarding. On the next iteration, we'll try
// to read a line normally.
buf.advance(offset + self.next_index + 1);
self.is_discarding = false;
self.next_index = 0;
}
(true, None) => {
// Otherwise, we didn't find a newline, so we'll discard
// everything we read. On the next iteration, we'll continue
// discarding up to max_len bytes unless we find a newline.
buf.advance(read_to);
self.next_index = 0;
if buf.is_empty() {
return Ok(None);
}
}
(false, Some(offset)) => {
// Found a line!
let newline_index = offset + self.next_index;
self.next_index = 0;
let line = buf.split_to(newline_index + 1);
let line = &line[..line.len() - 1];
let line = without_carriage_return(line);
let item =
serde_json::from_slice(line).map_err(JsonRpcMessageCodecError::Serde)?;
return Ok(Some(item));
}
(false, None) if buf.len() > self.max_length => {
// Reached the maximum length without finding a
// newline, return an error and start discarding on the
// next call.
self.is_discarding = true;
return Err(JsonRpcMessageCodecError::MaxLineLengthExceeded);
}
(false, None) => {
// We didn't find a line or reach the length limit, so the next
// call will resume searching at the current offset.
self.next_index = read_to;
return Ok(None);
}
}
}
}
fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, JsonRpcMessageCodecError> {
Ok(match self.decode(buf)? {
Some(frame) => Some(frame),
None => {
self.next_index = 0;
// No terminating newline - return remaining data, if any
if buf.is_empty() || buf == &b"\r"[..] {
None
} else {
let line = buf.split_to(buf.len());
let line = without_carriage_return(&line);
let item =
serde_json::from_slice(line).map_err(JsonRpcMessageCodecError::Serde)?;
Some(item)
}
}
})
}
}
impl<T: Serialize> Encoder<T> for JsonRpcMessageCodec<T> {
type Error = JsonRpcMessageCodecError;
fn encode(&mut self, item: T, buf: &mut BytesMut) -> Result<(), JsonRpcMessageCodecError> {
serde_json::to_writer(buf.writer(), &item)?;
buf.put_u8(b'\n');
Ok(())
}
}
#[cfg(test)]
mod test {
use futures::{Sink, Stream};
use super::*;
fn from_async_read<T: DeserializeOwned, R: AsyncRead>(reader: R) -> impl Stream<Item = T> {
FramedRead::new(reader, JsonRpcMessageCodec::<T>::default()).filter_map(|result| {
if let Err(e) = &result {
tracing::error!("Error reading from stream: {}", e);
}
futures::future::ready(result.ok())
})
}
fn from_async_write<T: Serialize, W: AsyncWrite + Send>(
writer: W,
) -> impl Sink<T, Error = std::io::Error> {
FramedWrite::new(writer, JsonRpcMessageCodec::<T>::default()).sink_map_err(Into::into)
}
#[tokio::test]
async fn test_decode() {
use futures::StreamExt;
use tokio::io::BufReader;
let data = r#"{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":1}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":2}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":3}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":4}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":5}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":6}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":7}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":8}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":9}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":10}
"#;
let mut cursor = BufReader::new(data.as_bytes());
let mut stream = from_async_read::<serde_json::Value, _>(&mut cursor);
for i in 1..=10 {
let item = stream.next().await.unwrap();
assert_eq!(
item,
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": if i % 2 != 0 { [42, 23] } else { [23, 42] },
"id": i,
})
);
}
}
#[tokio::test]
async fn test_encode() {
let test_messages = vec![
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": [42, 23],
"id": 1,
}),
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": [23, 42],
"id": 2,
}),
];
// Create a buffer to write to
let mut buffer = Vec::new();
let mut writer = from_async_write(&mut buffer);
// Write the test messages
for message in test_messages.iter() {
writer.send(message.clone()).await.unwrap();
}
writer.close().await.unwrap();
drop(writer);
// Parse the buffer back into lines and check each one
let output = String::from_utf8_lossy(&buffer);
let mut lines = output.lines();
for expected_message in test_messages {
let line = lines.next().unwrap();
let parsed_message: serde_json::Value = serde_json::from_str(line).unwrap();
assert_eq!(parsed_message, expected_message);
}
// Make sure there are no extra lines
assert!(lines.next().is_none());
}
}

42
crates/rmcp/src/transport/auth.rs
View file

@ -13,9 +13,42 @@ use oauth2::{
use reqwest::{Client as HttpClient, IntoUrl, StatusCode, Url, header::AUTHORIZATION};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::sync::RwLock;
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, error};
/// sse client with oauth2 authorization
#[derive(Clone)]
pub struct AuthClient<C> {
pub http_client: C,
pub auth_manager: Arc<Mutex<AuthorizationManager>>,
}
impl<C: std::fmt::Debug> std::fmt::Debug for AuthClient<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AuthorizedClient")
.field("http_client", &self.http_client)
.field("auth_manager", &"...")
.finish()
}
}
impl<C> AuthClient<C> {
/// create new authorized sse client
pub fn new(http_client: C, auth_manager: AuthorizationManager) -> Self {
Self {
http_client,
auth_manager: Arc::new(Mutex::new(auth_manager)),
}
}
}
impl<C> AuthClient<C> {
pub fn get_access_token(&self) -> impl Future<Output = Result<String, AuthError>> + Send {
let auth_manager = self.auth_manager.clone();
async move { auth_manager.lock().await.get_access_token().await }
}
}
/// Auth error
#[derive(Debug, Error)]
pub enum AuthError {
@ -798,4 +831,11 @@ impl OAuthState {
}
}
}
pub fn into_authorization_manager(self) -> Option<AuthorizationManager> {
match self {
OAuthState::Authorized(manager) => Some(manager),
_ => None,
}
}
}

14
crates/rmcp/src/transport/child_process.rs
View file

@ -1,11 +1,10 @@
use futures::{Sink, Stream};
use tokio::{
io::AsyncRead,
process::{ChildStdin, ChildStdout},
};
use super::IntoTransport;
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
use super::{IntoTransport, Transport};
use crate::service::ServiceRole;
pub(crate) fn child_process(
mut child: tokio::process::Child,
@ -77,13 +76,8 @@ impl TokioChildProcess {
}
impl<R: ServiceRole> IntoTransport<R, std::io::Error, ()> for TokioChildProcess {
fn into_transport(
self,
) -> (
impl Sink<TxJsonRpcMessage<R>, Error = std::io::Error> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<R>> + Send + 'static,
) {
IntoTransport::<R, std::io::Error, super::io::TransportAdapterAsyncRW>::into_transport(
fn into_transport(self) -> impl Transport<R, Error = std::io::Error> + 'static {
IntoTransport::<R, std::io::Error, super::async_rw::TransportAdapterAsyncRW>::into_transport(
self.split(),
)
}

14
crates/rmcp/src/transport/common.rs
View file

@ -3,3 +3,17 @@
feature = "transport-sse-server"
))]
pub mod axum;
pub mod http_header;
#[cfg(feature = "reqwest")]
#[cfg_attr(docsrs, doc(cfg(feature = "reqwest")))]
mod reqwest;
#[cfg(feature = "client-side-sse")]
#[cfg_attr(docsrs, doc(cfg(feature = "client-side-sse")))]
pub mod client_side_sse;
#[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
pub mod auth;

7
crates/rmcp/src/transport/common/auth.rs Normal file
View file

@ -0,0 +1,7 @@
#[cfg(feature = "__transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-streamable-http-client")))]
mod streamable_http_client;
#[cfg(feature = "__transport-sse")]
#[cfg_attr(docsrs, doc(cfg(feature = "__transport-sse")))]
mod sse_client;

43
crates/rmcp/src/transport/common/auth/sse_client.rs Normal file
View file

@ -0,0 +1,43 @@
use crate::transport::{
auth::AuthClient,
sse_client::{SseClient, SseTransportError},
};
impl<C> SseClient for AuthClient<C>
where
C: SseClient,
{
type Error = SseTransportError<C::Error>;
async fn post_message(
&self,
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
mut auth_token: Option<String>,
) -> Result<(), SseTransportError<Self::Error>> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message(uri, message, auth_token)
.await
.map_err(SseTransportError::Client)
}
async fn get_stream(
&self,
uri: std::sync::Arc<str>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
) -> Result<
crate::transport::common::client_side_sse::BoxedSseResponse,
SseTransportError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream(uri, last_event_id, auth_token)
.await
.map_err(SseTransportError::Client)
}
}

64
crates/rmcp/src/transport/common/auth/streamable_http_client.rs Normal file
View file

@ -0,0 +1,64 @@
use crate::transport::{
auth::AuthClient,
streamable_http_client::{StreamableHttpClient, StreamableHttpError},
};
impl<C> StreamableHttpClient for AuthClient<C>
where
C: StreamableHttpClient + Send + Sync,
{
type Error = StreamableHttpError<C::Error>;
async fn delete_session(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
mut auth_token: Option<String>,
) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError<Self::Error>>
{
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.delete_session(uri, session_id, auth_token)
.await
.map_err(StreamableHttpError::Client)
}
async fn get_stream(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
) -> Result<
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream(uri, session_id, last_event_id, auth_token)
.await
.map_err(StreamableHttpError::Client)
}
async fn post_message(
&self,
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
session_id: Option<std::sync::Arc<str>>,
mut auth_token: Option<String>,
) -> Result<
crate::transport::streamable_http_client::StreamableHttpPostResponse,
StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message(uri, message, session_id, auth_token)
.await
.map_err(StreamableHttpError::Client)
}
}

236
crates/rmcp/src/transport/common/client_side_sse.rs Normal file
View file

@ -0,0 +1,236 @@
use std::{
pin::Pin,
sync::Arc,
task::{Poll, ready},
time::Duration,
};
use futures::{Stream, stream::BoxStream};
use sse_stream::{Error as SseError, Sse};
use crate::model::ServerJsonRpcMessage;
pub type BoxedSseResponse = BoxStream<'static, Result<Sse, SseError>>;
pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync {
fn retry(&self, current_times: usize) -> Option<Duration>;
}
#[derive(Debug, Clone)]
pub struct FixedInterval {
pub max_times: Option<usize>,
pub duration: Duration,
}
impl SseRetryPolicy for FixedInterval {
fn retry(&self, current_times: usize) -> Option<Duration> {
if let Some(max_times) = self.max_times {
if current_times >= max_times {
return None;
}
}
Some(self.duration)
}
}
impl FixedInterval {
pub const DEFAULT_MIN_DURATION: Duration = Duration::from_millis(1000);
}
impl Default for FixedInterval {
fn default() -> Self {
Self {
max_times: None,
duration: Self::DEFAULT_MIN_DURATION,
}
}
}
#[derive(Debug, Clone)]
pub struct ExponentialBackoff {
pub max_times: Option<usize>,
pub base_duration: Duration,
}
impl ExponentialBackoff {
pub const DEFAULT_DURATION: Duration = Duration::from_millis(1000);
}
impl Default for ExponentialBackoff {
fn default() -> Self {
Self {
max_times: None,
base_duration: Self::DEFAULT_DURATION,
}
}
}
impl SseRetryPolicy for ExponentialBackoff {
fn retry(&self, current_times: usize) -> Option<Duration> {
if let Some(max_times) = self.max_times {
if current_times >= max_times {
return None;
}
}
Some(self.base_duration * (2u32.pow(current_times as u32)))
}
}
pub(crate) trait SseStreamReconnect {
type Error: std::error::Error;
type Future: Future<Output = Result<BoxedSseResponse, Self::Error>> + Send;
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future;
}
pin_project_lite::pin_project! {
pub(crate) struct SseAutoReconnectStream<R>
where R: SseStreamReconnect
{
retry_policy: Arc<dyn SseRetryPolicy>,
last_event_id: Option<String>,
server_retry_interval: Option<Duration>,
connector: R,
#[pin]
state: SseAutoReconnectStreamState<R::Future>,
}
}
impl<R: SseStreamReconnect> SseAutoReconnectStream<R> {
pub fn new(
stream: BoxedSseResponse,
connector: R,
retry_policy: Arc<dyn SseRetryPolicy>,
) -> Self {
Self {
retry_policy,
last_event_id: None,
server_retry_interval: None,
connector,
state: SseAutoReconnectStreamState::Connected { stream },
}
}
}
pin_project_lite::pin_project! {
#[project = SseAutoReconnectStreamStateProj]
pub enum SseAutoReconnectStreamState<F> {
Connected {
#[pin]
stream: BoxedSseResponse,
},
Retrying {
retry_times: usize,
#[pin]
retrying: F,
},
WaitingNextRetry {
#[pin]
sleep: tokio::time::Sleep,
retry_times: usize,
},
Terminated,
}
}
impl<R> Stream for SseAutoReconnectStream<R>
where
R: SseStreamReconnect,
{
type Item = Result<ServerJsonRpcMessage, R::Error>;
fn poll_next(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut().project();
// let this_state = this.state.as_mut().project()
let state = this.state.as_mut().project();
let next_state = match state {
SseAutoReconnectStreamStateProj::Connected { stream } => {
match ready!(stream.poll_next(cx)) {
Some(Ok(sse)) => {
if let Some(new_server_retry) = sse.retry {
*this.server_retry_interval =
Some(Duration::from_millis(new_server_retry));
}
if let Some(event_id) = sse.id {
*this.last_event_id = Some(event_id);
}
if let Some(data) = sse.data {
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
Err(e) => {
// not sure should this be a hard error
tracing::warn!("failed to deserialize server message: {e}");
return self.poll_next(cx);
}
Ok(message) => {
return Poll::Ready(Some(Ok(message)));
}
};
} else {
return self.poll_next(cx);
}
}
Some(Err(e)) => {
tracing::warn!("sse stream error: {e}");
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
SseAutoReconnectStreamState::Retrying {
retry_times: 0,
retrying,
}
}
None => {
tracing::debug!("sse stream terminated");
return Poll::Ready(None);
}
}
}
SseAutoReconnectStreamStateProj::Retrying {
retry_times,
retrying,
} => {
let retry_result = ready!(retrying.poll(cx));
match retry_result {
Ok(new_stream) => SseAutoReconnectStreamState::Connected { stream: new_stream },
Err(e) => {
tracing::debug!("retry sse stream error: {e}");
*retry_times += 1;
if let Some(interval) = this.retry_policy.retry(*retry_times) {
let interval = this
.server_retry_interval
.map(|server_retry_interval| server_retry_interval.max(interval))
.unwrap_or(interval);
let sleep = tokio::time::sleep(interval);
SseAutoReconnectStreamState::WaitingNextRetry {
sleep,
retry_times: *retry_times,
}
} else {
tracing::error!("sse stream error: {e}, max retry times reached");
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
}
}
}
}
SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => {
ready!(sleep.poll(cx));
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
let retry_times = *retry_times;
SseAutoReconnectStreamState::Retrying {
retry_times,
retrying,
}
}
SseAutoReconnectStreamStateProj::Terminated => {
return Poll::Ready(None);
}
};
// update the state
this.state.set(next_state);
self.poll_next(cx)
}
}

4
crates/rmcp/src/transport/common/http_header.rs Normal file
View file

@ -0,0 +1,4 @@
pub const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
pub const HEADER_LAST_EVENT_ID: &str = "Last-Event-Id";
pub const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
pub const JSON_MIME_TYPE: &str = "application/json";

Some files were not shown because too many files have changed in this diff Show more