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

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

7
crates/rmcp/src/transport/common/reqwest.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;

82
crates/rmcp/src/transport/common/reqwest/sse_client.rs Normal file
View file

@ -0,0 +1,82 @@
use std::sync::Arc;
use futures::StreamExt;
use reqwest::header::ACCEPT;
use sse_stream::SseStream;
use crate::transport::{
SseClientTransport,
common::http_header::{EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID},
sse_client::{SseClient, SseClientConfig, SseTransportError},
};
impl SseClient for reqwest::Client {
type Error = reqwest::Error;
async fn post_message(
&self,
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
auth_token: Option<String>,
) -> Result<(), SseTransportError<Self::Error>> {
let mut request_builder = self.post(uri.as_ref()).json(&message);
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
request_builder
.send()
.await
.and_then(|resp| resp.error_for_status())
.map_err(SseTransportError::from)
.map(drop)
}
async fn get_stream(
&self,
uri: std::sync::Arc<str>,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> Result<
crate::transport::common::client_side_sse::BoxedSseResponse,
SseTransportError<Self::Error>,
> {
let mut request_builder = self
.get(uri.as_ref())
.header(ACCEPT, EVENT_STREAM_MIME_TYPE);
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
let response = request_builder.send().await?;
let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) {
Some(ct) => {
if !ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) {
return Err(SseTransportError::UnexpectedContentType(Some(ct.clone())));
}
}
None => {
return Err(SseTransportError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
}
}
impl SseClientTransport<reqwest::Client> {
pub async fn start(
uri: impl Into<Arc<str>>,
) -> Result<Self, SseTransportError<reqwest::Error>> {
SseClientTransport::start_with_client(
reqwest::Client::default(),
SseClientConfig {
uri: uri.into(),
..Default::default()
},
)
.await
}
}

138
crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs Normal file
View file

@ -0,0 +1,138 @@
use std::sync::Arc;
use futures::{StreamExt, stream::BoxStream};
use reqwest::header::ACCEPT;
use sse_stream::{Sse, SseStream};
use crate::{
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::{
common::http_header::{
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
},
streamable_http_client::*,
},
};
impl StreamableHttpClient for reqwest::Client {
type Error = reqwest::Error;
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
let mut request_builder = self
.get(uri.as_ref())
.header(ACCEPT, EVENT_STREAM_MIME_TYPE)
.header(HEADER_SESSION_ID, session_id.as_ref());
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
let response = request_builder.send().await?;
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
return Err(StreamableHttpError::SeverDoesNotSupportSse);
}
let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) {
Some(ct) => {
if !ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) {
return Err(StreamableHttpError::UnexpectedContentType(Some(
String::from_utf8_lossy(ct.as_bytes()).to_string(),
)));
}
}
None => {
return Err(StreamableHttpError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
}
async fn delete_session(
&self,
uri: Arc<str>,
session: Arc<str>,
auth_token: Option<String>,
) -> Result<(), StreamableHttpError<Self::Error>> {
let mut request_builder = self.delete(uri.as_ref());
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
let response = request_builder
.header(HEADER_SESSION_ID, session.as_ref())
.send()
.await?;
// if method no allowed
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
tracing::debug!("this server doesn't support deleting session");
return Ok(());
}
let _response = response.error_for_status()?;
Ok(())
}
async fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_token: Option<String>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
let mut request = self
.post(uri.as_ref())
.header(ACCEPT, EVENT_STREAM_MIME_TYPE)
.header(ACCEPT, JSON_MIME_TYPE);
if let Some(auth_header) = auth_token {
request = request.bearer_auth(auth_header);
}
if let Some(session_id) = session_id {
request = request.header(HEADER_SESSION_ID, session_id.as_ref());
}
let response = request.json(&message).send().await?;
if response.status() == reqwest::StatusCode::ACCEPTED {
return Ok(StreamableHttpPostResponse::Accepted);
}
let content_type = response.headers().get(reqwest::header::CONTENT_TYPE);
let session_id = response.headers().get(HEADER_SESSION_ID);
let session_id = session_id
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
match content_type {
Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
}
Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
let message: ServerJsonRpcMessage = response.json().await?;
Ok(StreamableHttpPostResponse::Json(message, session_id))
}
_ => {
// unexpected content type
tracing::error!("unexpected content type: {:?}", content_type);
Err(StreamableHttpError::UnexpectedContentType(
content_type.map(|ct| String::from_utf8_lossy(ct.as_bytes()).to_string()),
))
}
}
}
}
impl StreamableHttpClientTransport<reqwest::Client> {
pub fn from_uri(uri: impl Into<Arc<str>>) -> Self {
StreamableHttpClientTransport::with_client(
reqwest::Client::default(),
StreamableHttpClientTransportConfig {
uri: uri.into(),
..Default::default()
},
)
}
}

308
crates/rmcp/src/transport/io.rs
View file

@ -1,314 +1,6 @@
use std::marker::PhantomData;
// use crate::schema::*;
use futures::{Sink, SinkExt, Stream, StreamExt};
use serde::{Serialize, de::DeserializeOwned};
use thiserror::Error;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_util::{
bytes::{Buf, BufMut, BytesMut},
codec::{Decoder, Encoder, FramedRead, FramedWrite},
};
use super::IntoTransport;
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
#[cfg(feature = "transport-io")]
/// # StdIO Transport
///
/// Create a pair of [`tokio::io::Stdin`] and [`tokio::io::Stdout`].
pub fn stdio() -> (tokio::io::Stdin, tokio::io::Stdout) {
(tokio::io::stdin(), tokio::io::stdout())
}
pub enum TransportAdapterAsyncRW {}
impl<Role, R, W> IntoTransport<Role, std::io::Error, TransportAdapterAsyncRW> for (R, W)
where
Role: ServiceRole,
R: AsyncRead + Send + 'static,
W: AsyncWrite + Send + 'static,
{
fn into_transport(
self,
) -> (
impl Sink<TxJsonRpcMessage<Role>, Error = std::io::Error> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
) {
(from_async_write(self.1), from_async_read(self.0))
}
}
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 Sink<TxJsonRpcMessage<Role>, Error = std::io::Error> + Send + 'static,
impl Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
) {
IntoTransport::<Role, std::io::Error, TransportAdapterAsyncRW>::into_transport(
tokio::io::split(self),
)
}
}
pub 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())
})
}
pub fn from_async_write<T: Serialize, W: AsyncWrite>(
writer: W,
) -> impl Sink<T, Error = std::io::Error> {
FramedWrite::new(writer, JsonRpcMessageCodec::<T>::default()).sink_map_err(Into::into)
}
#[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 super::*;
#[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());
}
}

77
crates/rmcp/src/transport/sink_stream.rs Normal file
View file

@ -0,0 +1,77 @@
use std::sync::Arc;
use futures::{Sink, Stream};
use tokio::sync::Mutex;
use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub struct SinkStreamTransport<Si, St> {
stream: St,
sink: Arc<Mutex<Si>>,
}
impl<Si, St> SinkStreamTransport<Si, St> {
pub fn new(sink: Si, stream: St) -> Self {
Self {
stream,
sink: Arc::new(Mutex::new(sink)),
}
}
}
impl<Role: ServiceRole, Si, St> Transport<Role> for SinkStreamTransport<Si, St>
where
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
{
type Error = Si::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<Role>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
use futures::SinkExt;
let lock = self.sink.clone();
async move {
let mut write = lock.lock().await;
write.send(item).await
}
}
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
use futures::StreamExt;
self.stream.next()
}
async fn close(&mut self) -> Result<(), Self::Error> {
Ok(())
}
}
pub enum TransportAdapterSinkStream {}
impl<Role, Si, St> IntoTransport<Role, Si::Error, TransportAdapterSinkStream> for (Si, St)
where
Role: ServiceRole,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin + 'static,
Si::Error: std::error::Error + Send + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = Si::Error> + 'static {
SinkStreamTransport::new(self.0, self.1)
}
}
pub enum TransportAdapterAsyncCombinedRW {}
impl<Role, S> IntoTransport<Role, S::Error, TransportAdapterAsyncCombinedRW> for S
where
Role: ServiceRole,
S: Sink<TxJsonRpcMessage<Role>> + Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
S::Error: std::error::Error + Send + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = S::Error> + 'static {
use futures::StreamExt;
IntoTransport::<Role, S::Error, TransportAdapterSinkStream>::into_transport(self.split())
}
}

388
crates/rmcp/src/transport/sse.rs
View file

@ -1,388 +0,0 @@
//！ reference: https://html.spec.whatwg.org/multipage/server-sent-events.html
use std::{collections::VecDeque, sync::Arc, time::Duration};
use futures::{FutureExt, Sink, Stream, StreamExt, future::BoxFuture, stream::BoxStream};
use reqwest::{
Client as HttpClient, IntoUrl, Url,
header::{ACCEPT, HeaderValue},
};
use sse_stream::{Error as SseError, Sse, SseStream};
use thiserror::Error;
use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
const MIME_TYPE: &str = "text/event-stream";
const HEADER_LAST_EVENT_ID: &str = "Last-Event-ID";
#[derive(Error, Debug)]
pub enum SseTransportError<E: std::error::Error + Send + Sync + 'static> {
#[error("SSE error: {0}")]
Sse(#[from] SseError),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Transport error: {0}")]
Transport(E),
#[error("unexpected end of stream")]
UnexpectedEndOfStream,
#[error("Url error: {0}")]
Url(#[from] url::ParseError),
#[error("Unexpected content type: {0:?}")]
UnexpectedContentType(Option<HeaderValue>),
}
type SseStreamFuture<E> =
BoxFuture<'static, Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<E>>>;
enum SseTransportState<E: std::error::Error + Send + Sync + 'static> {
Connected(BoxStream<'static, Result<Sse, SseError>>),
Retrying {
times: usize,
fut: SseStreamFuture<E>,
},
Fatal {
reason: String,
},
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct SseTransportRetryConfig {
pub max_times: Option<usize>,
pub min_duration: Duration,
}
impl SseTransportRetryConfig {
pub const DEFAULT_MIN_DURATION: Duration = Duration::from_millis(1000);
}
impl Default for SseTransportRetryConfig {
fn default() -> Self {
Self {
max_times: None,
min_duration: Self::DEFAULT_MIN_DURATION,
}
}
}
impl From<reqwest::Error> for SseTransportError<reqwest::Error> {
fn from(e: reqwest::Error) -> Self {
SseTransportError::Transport(e)
}
}
pub trait SseClient<E: std::error::Error + Send + Sync>: Clone + Send + Sync + 'static {
fn connect(&self, last_event_id: Option<String>) -> SseStreamFuture<E>;
fn post(
&self,
endpoint: &str,
message: ClientJsonRpcMessage,
) -> BoxFuture<'static, Result<(), SseTransportError<E>>>;
}
pub struct RetryConfig {
pub max_times: Option<usize>,
pub min_duration: Duration,
}
#[derive(Clone)]
pub struct ReqwestSseClient {
http_client: HttpClient,
sse_url: Url,
}
impl ReqwestSseClient {
pub fn new<U>(url: U) -> Result<Self, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let url = url.into_url()?;
Ok(Self {
http_client: HttpClient::default(),
sse_url: url,
})
}
pub async fn new_with_timeout<U>(
url: U,
timeout: Duration,
) -> Result<Self, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let mut client = HttpClient::builder();
client = client.timeout(timeout);
let client = client.build()?;
let url = url.into_url()?;
Ok(Self {
http_client: client,
sse_url: url,
})
}
pub async fn new_with_client<U>(
url: U,
client: HttpClient,
) -> Result<Self, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let url = url.into_url()?;
Ok(Self {
http_client: client,
sse_url: url,
})
}
}
impl SseClient<reqwest::Error> for ReqwestSseClient {
fn connect(&self, last_event_id: Option<String>) -> SseStreamFuture<reqwest::Error> {
let client = self.http_client.clone();
let sse_url = self.sse_url.as_ref().to_string();
let last_event_id = last_event_id.clone();
let fut = async move {
let mut request_builder = client.get(&sse_url).header(ACCEPT, MIME_TYPE);
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
let response = request_builder.send().await?;
let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) {
Some(ct) => {
if !ct.as_bytes().starts_with(MIME_TYPE.as_bytes()) {
return Err(SseTransportError::UnexpectedContentType(Some(ct.clone())));
}
}
None => {
return Err(SseTransportError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
};
fut.boxed()
}
fn post(
&self,
session_id: &str,
message: ClientJsonRpcMessage,
) -> BoxFuture<'static, Result<(), SseTransportError<reqwest::Error>>> {
let client = self.http_client.clone();
let sse_url = self.sse_url.clone();
let session_id = session_id.to_string();
Box::pin(async move {
let uri = sse_url.join(&session_id).map_err(SseTransportError::from)?;
let request_builder = client.post(uri.as_ref()).json(&message);
request_builder
.send()
.await
.and_then(|resp| resp.error_for_status())
.map_err(SseTransportError::from)
.map(drop)
})
}
}
/// # Transport for client sse
///
/// Call [`SseTransport::start`] to create a new transport from url.
///
/// Call [`SseTransport::start_with_client`] to create a new transport with a customized reqwest client.
pub struct SseTransport<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> {
client: Arc<C>,
state: SseTransportState<E>,
last_event_id: Option<String>,
recommended_retry_duration_ms: Option<u64>,
session_id: String,
#[allow(clippy::type_complexity)]
request_queue: VecDeque<tokio::sync::oneshot::Receiver<Result<(), SseTransportError<E>>>>,
pub retry_config: SseTransportRetryConfig,
}
impl SseTransport<ReqwestSseClient, reqwest::Error> {
pub async fn start<U>(
url: U,
) -> Result<SseTransport<ReqwestSseClient, reqwest::Error>, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let client = ReqwestSseClient::new(url)?;
SseTransport::start_with_client(client).await
}
}
impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> SseTransport<C, E> {
pub async fn start_with_client(client: C) -> Result<Self, SseTransportError<E>> {
let mut event_stream = client.connect(None).await?;
let mut last_event_id = None;
let mut retry = None;
let session_id = loop {
let next_event = event_stream
.next()
.await
.ok_or(SseTransportError::UnexpectedEndOfStream)??;
if let Some(id) = next_event.id {
last_event_id = Some(id);
}
if let Some(retry_ms) = next_event.retry {
retry = Some(retry_ms);
}
if let Some("endpoint") = next_event.event.as_deref() {
break next_event.data.unwrap_or_default();
}
};
Ok(SseTransport {
client: Arc::new(client),
state: SseTransportState::Connected(Box::pin(event_stream)),
last_event_id,
recommended_retry_duration_ms: retry,
session_id,
request_queue: Default::default(),
retry_config: Default::default(),
})
}
fn retry_connection(&self) -> SseStreamFuture<E> {
let retry_duration = {
let recommended_retry_duration = self
.recommended_retry_duration_ms
.map(Duration::from_millis);
let config_retry_duration = self.retry_config.min_duration;
recommended_retry_duration
.map(|d: Duration| d.max(config_retry_duration))
.unwrap_or(config_retry_duration)
};
std::thread::sleep(retry_duration);
let last_event_id = self.last_event_id.clone();
self.client.connect(last_event_id)
}
}
impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> Stream for SseTransport<C, E> {
type Item = ServerJsonRpcMessage;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
let retry_config = self.retry_config;
match &mut self.state {
SseTransportState::Connected(event_stream) => {
let event = std::task::ready!(event_stream.poll_next_unpin(cx));
match event {
Some(Ok(event)) => {
if let Some(retry) = event.retry {
self.as_mut().recommended_retry_duration_ms = Some(retry);
}
if let Some(id) = event.id {
self.as_mut().last_event_id = Some(id);
}
if let Some(data) = event.data {
match serde_json::from_str(&data) {
Ok(message) => std::task::Poll::Ready(Some(message)),
Err(e) => {
tracing::error!(error = %e, "failed to parse json rpc request");
self.poll_next(cx)
}
}
} else {
self.poll_next(cx)
}
}
Some(Err(e)) => {
tracing::error!(error = %e, "sse event stream encounter an error");
let fut = self.retry_connection();
self.as_mut().state = SseTransportState::Retrying { times: 1, fut };
self.poll_next(cx)
}
None => std::task::Poll::Ready(None),
}
}
SseTransportState::Retrying { fut, times } => {
let retry_result = std::task::ready!(fut.poll_unpin(cx));
match retry_result {
Ok(stream) => {
self.as_mut().state = SseTransportState::Connected(stream);
self.poll_next(cx)
}
Err(e) => {
tracing::warn!(error = %e, "retrying failed");
if let Some(max_retry_times) = retry_config.max_times {
if *times >= max_retry_times {
self.as_mut().state = SseTransportState::Fatal {
reason: format!("retrying failed after {} times: {}", times, e),
};
return self.poll_next(cx);
}
}
let times = *times + 1;
let fut = self.retry_connection();
self.as_mut().state = SseTransportState::Retrying { times, fut };
self.poll_next(cx)
}
}
}
SseTransportState::Fatal { reason } => {
tracing::error!("sse transport fatal error: {}", reason);
std::task::Poll::Ready(None)
}
}
}
}
impl<C: SseClient<E>, E: std::error::Error + Send + Sync + 'static> Sink<ClientJsonRpcMessage>
for SseTransport<C, E>
{
type Error = SseTransportError<E>;
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
const QUEUE_SIZE: usize = 16;
if self.as_mut().request_queue.len() >= QUEUE_SIZE {
std::task::ready!(
self.as_mut()
.request_queue
.front_mut()
.expect("queue is not empty")
.poll_unpin(cx)
)
.expect("sender shall not drop")?;
}
std::task::Poll::Ready(Ok(()))
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ClientJsonRpcMessage,
) -> Result<(), Self::Error> {
let client = self.client.clone();
let session_id = self.session_id.clone();
let (tx, rx) = tokio::sync::oneshot::channel();
let session_id = session_id.clone();
tokio::spawn(async move {
let result = { client.post(&session_id, item).await };
let _ = tx.send(result);
});
self.request_queue.push_back(rx);
Ok(())
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
let queue = &mut self.as_mut().request_queue;
while let Some(fut) = queue.front_mut() {
std::task::ready!(fut.poll_unpin(cx)).expect("sender shall not drop")?;
queue.pop_front();
}
std::task::Poll::Ready(Ok(()))
}
fn poll_close(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.poll_flush(cx)
}
}

183
crates/rmcp/src/transport/sse_auth.rs
View file

@ -1,183 +0,0 @@
use std::sync::Arc;
use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
use reqwest::{
Client as HttpClient, IntoUrl, Url,
header::{ACCEPT, AUTHORIZATION},
};
use sse_stream::{Error as SseError, Sse, SseStream};
use tokio::sync::Mutex;
use super::{
auth::{AuthError, AuthorizationManager, OAuthState},
sse::{SseClient, SseTransport, SseTransportError, SseTransportRetryConfig},
};
use crate::model::ClientJsonRpcMessage;
// SSE MIME type
const MIME_TYPE: &str = "text/event-stream";
const HEADER_LAST_EVENT_ID: &str = "Last-Event-ID";
/// sse client with oauth2 authorization
#[derive(Clone)]
pub struct AuthorizedSseClient {
http_client: HttpClient,
sse_url: Url,
auth_manager: Arc<Mutex<AuthorizationManager>>,
_retry_config: Option<SseTransportRetryConfig>, // TODO retry config may be used by authorized transport for token
}
impl AuthorizedSseClient {
/// create new authorized sse client
pub fn new<U>(
url: U,
auth_manager: Arc<Mutex<AuthorizationManager>>,
retry_config: Option<SseTransportRetryConfig>,
) -> Result<Self, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let url = url.into_url().map_err(SseTransportError::from)?;
Ok(Self {
http_client: HttpClient::default(),
sse_url: url,
auth_manager,
_retry_config: retry_config,
})
}
/// create authorized sse client with custom http client
pub async fn new_with_client<U>(
url: U,
client: HttpClient,
auth_manager: Arc<Mutex<AuthorizationManager>>,
retry_config: Option<SseTransportRetryConfig>,
) -> Result<Self, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
let url = url.into_url().map_err(SseTransportError::from)?;
Ok(Self {
http_client: client,
sse_url: url,
auth_manager,
_retry_config: retry_config,
})
}
}
impl SseClient<reqwest::Error> for AuthorizedSseClient {
fn connect(
&self,
last_event_id: Option<String>,
) -> BoxFuture<
'static,
Result<BoxStream<'static, Result<Sse, SseError>>, SseTransportError<reqwest::Error>>,
> {
let client = self.http_client.clone();
let sse_url = self.sse_url.as_ref().to_string();
let last_event_id = last_event_id.clone();
let auth_manager = self.auth_manager.clone();
let fut = async move {
// get access token
let token = auth_manager.lock().await.get_access_token().await?;
// build request
let mut request_builder = client
.get(&sse_url)
.header(ACCEPT, MIME_TYPE)
.header(AUTHORIZATION, format!("Bearer {}", token));
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
let response = request_builder.send().await?;
let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) {
Some(ct) => {
if !ct.as_bytes().starts_with(MIME_TYPE.as_bytes()) {
return Err(SseTransportError::UnexpectedContentType(Some(ct.clone())));
}
}
None => {
return Err(SseTransportError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
};
Box::pin(fut)
}
fn post(
&self,
session_id: &str,
message: ClientJsonRpcMessage,
) -> BoxFuture<'static, Result<(), SseTransportError<reqwest::Error>>> {
let client = self.http_client.clone();
let sse_url = self.sse_url.clone();
let session_id = session_id.to_string();
let auth_manager = self.auth_manager.clone();
Box::pin(async move {
// get access token
let token = auth_manager
.lock()
.await
.get_access_token()
.await
.map_err(SseTransportError::<reqwest::Error>::from)?;
let uri = sse_url.join(&session_id).map_err(SseTransportError::from)?;
let request_builder = client
.post(uri.as_ref())
.header(AUTHORIZATION, format!("Bearer {}", token))
.json(&message);
request_builder
.send()
.await
.and_then(|resp| resp.error_for_status())
.map_err(SseTransportError::from)
.map(drop)
})
}
}
impl From<AuthError> for SseTransportError<reqwest::Error> {
fn from(err: AuthError) -> Self {
SseTransportError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
err.to_string(),
))
}
}
/// create authorized sse transport
pub async fn create_authorized_transport<U>(
url: U,
auth_state: OAuthState,
retry_config: Option<SseTransportRetryConfig>,
) -> Result<SseTransport<AuthorizedSseClient, reqwest::Error>, SseTransportError<reqwest::Error>>
where
U: IntoUrl,
{
match auth_state {
OAuthState::Authorized(auth_manager) => {
let auth_manager = Arc::new(Mutex::new(auth_manager));
let client = AuthorizedSseClient::new(url, auth_manager, retry_config)?;
let mut transport = SseTransport::start_with_client(client).await?;
transport.retry_config = retry_config.unwrap_or_default();
Ok(transport)
}
_ => Err(SseTransportError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
"Not authorized".to_string(),
))),
}
}

169
crates/rmcp/src/transport/sse_client.rs Normal file
View file

@ -0,0 +1,169 @@
//！ reference: https://html.spec.whatwg.org/multipage/server-sent-events.html
use std::{pin::Pin, sync::Arc};
use futures::{StreamExt, future::BoxFuture};
use reqwest::header::HeaderValue;
use sse_stream::Error as SseError;
use thiserror::Error;
use super::{
Transport,
common::client_side_sse::{BoxedSseResponse, SseRetryPolicy, SseStreamReconnect},
};
use crate::{
RoleClient,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::common::client_side_sse::SseAutoReconnectStream,
};
#[derive(Error, Debug)]
pub enum SseTransportError<E: std::error::Error + Send + Sync + 'static> {
#[error("SSE error: {0}")]
Sse(#[from] SseError),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Client error: {0}")]
Client(E),
#[error("unexpected end of stream")]
UnexpectedEndOfStream,
#[error("Url error: {0}")]
Url(#[from] url::ParseError),
#[error("Unexpected content type: {0:?}")]
UnexpectedContentType(Option<HeaderValue>),
#[error("Tokio join error: {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
#[error("Transport terminated")]
TransportTerminated,
#[cfg(feature = "__auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "__auth")))]
#[error("Auth error: {0}")]
Auth(#[from] crate::transport::auth::AuthError),
}
impl From<reqwest::Error> for SseTransportError<reqwest::Error> {
fn from(e: reqwest::Error) -> Self {
SseTransportError::Client(e)
}
}
pub trait SseClient: Clone + Send + Sync + 'static {
type Error: std::error::Error + Send + Sync + 'static;
fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
auth_token: Option<String>,
) -> impl Future<Output = Result<(), SseTransportError<Self::Error>>> + Send + '_;
fn get_stream(
&self,
uri: Arc<str>,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> impl Future<Output = Result<BoxedSseResponse, SseTransportError<Self::Error>>> + Send + '_;
}
struct SseClientReconnect<C> {
pub client: C,
pub uri: Arc<str>,
}
impl<C: SseClient> SseStreamReconnect for SseClientReconnect<C> {
type Error = SseTransportError<C::Error>;
type Future = BoxFuture<'static, Result<BoxedSseResponse, Self::Error>>;
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future {
let client = self.client.clone();
let uri = self.uri.clone();
let last_event_id = last_event_id.map(|s| s.to_owned());
Box::pin(async move { client.get_stream(uri, last_event_id, None).await })
}
}
type ServerMessageStream<C> = Pin<Box<SseAutoReconnectStream<SseClientReconnect<C>>>>;
pub struct SseClientTransport<C: SseClient> {
client: C,
config: SseClientConfig,
post_uri: Arc<str>,
stream: Option<ServerMessageStream<C>>,
}
impl<C: SseClient> Transport<RoleClient> for SseClientTransport<C> {
type Error = SseTransportError<C::Error>;
async fn receive(&mut self) -> Option<ServerJsonRpcMessage> {
self.stream.as_mut()?.next().await?.ok()
}
fn send(
&mut self,
item: crate::service::TxJsonRpcMessage<RoleClient>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let client = self.client.clone();
let uri = self.post_uri.clone();
async move { client.post_message(uri, item, None).await }
}
async fn close(&mut self) -> Result<(), Self::Error> {
self.stream.take();
Ok(())
}
}
impl<C: SseClient + std::fmt::Debug> std::fmt::Debug for SseClientTransport<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SseClientWorker")
.field("client", &self.client)
.field("config", &self.config)
.finish()
}
}
impl<C: SseClient> SseClientTransport<C> {
pub async fn start_with_client(
client: C,
config: SseClientConfig,
) -> Result<Self, SseTransportError<C::Error>> {
let mut sse_stream = client.get_stream(config.uri.clone(), None, None).await?;
// wait the endpoint event
let endpoint = loop {
let sse = sse_stream
.next()
.await
.ok_or(SseTransportError::UnexpectedEndOfStream)??;
let Some("endpoint") = sse.event.as_deref() else {
continue;
};
break sse.data.unwrap_or_default();
};
let post_uri: Arc<str> = format!(
"{}/{}",
config.uri.trim_end_matches("/"),
endpoint.trim_start_matches("/")
)
.into();
let stream = Box::pin(SseAutoReconnectStream::new(
sse_stream,
SseClientReconnect {
client: client.clone(),
uri: config.uri.clone(),
},
config.retry_policy.clone(),
));
Ok(Self {
client,
config,
post_uri,
stream: Some(stream),
})
}
}
#[derive(Debug, Clone)]
pub struct SseClientConfig {
pub uri: Arc<str>,
pub retry_policy: Arc<dyn SseRetryPolicy>,
}
impl Default for SseClientConfig {
fn default() -> Self {
Self {
uri: "".into(),
retry_policy: Arc::new(super::common::client_side_sse::FixedInterval::default()),
}
}
}

492
crates/rmcp/src/transport/streamable_http_client.rs Normal file
View file

@ -0,0 +1,492 @@
use std::{borrow::Cow, sync::Arc, time::Duration};
use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
pub use sse_stream::Error as SseError;
use sse_stream::Sse;
use thiserror::Error;
use tokio_util::sync::CancellationToken;
use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect};
use crate::{
RoleClient,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::{
common::client_side_sse::SseAutoReconnectStream,
worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport},
},
};
type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>;
#[derive(Error, Debug)]
pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
#[error("SSE error: {0}")]
Sse(#[from] SseError),
#[error("Io error: {0}")]
Io(#[from] std::io::Error),
#[error("Client error: {0}")]
Client(E),
#[error("unexpected end of stream")]
UnexpectedEndOfStream,
#[error("unexpected server response: {0}")]
UnexpectedServerResponse(Cow<'static, str>),
#[error("Unexpected content type: {0:?}")]
UnexpectedContentType(Option<String>),
#[error("Server does not support SSE")]
SeverDoesNotSupportSse,
#[error("Server does not support delete session")]
SeverDoesNotSupportDeleteSession,
#[error("Tokio join error: {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
#[error("Deserialize error: {0}")]
Deserialize(#[from] serde_json::Error),
#[error("Transport channel closed")]
TransportChannelClosed,
#[cfg(feature = "__auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "__auth")))]
#[error("Auth error: {0}")]
Auth(#[from] crate::transport::auth::AuthError),
}
impl From<reqwest::Error> for StreamableHttpError<reqwest::Error> {
fn from(e: reqwest::Error) -> Self {
StreamableHttpError::Client(e)
}
}
pub enum StreamableHttpPostResponse {
Accepted,
Json(ServerJsonRpcMessage, Option<String>),
Sse(BoxedSseStream, Option<String>),
}
impl std::fmt::Debug for StreamableHttpPostResponse {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Accepted => write!(f, "Accepted"),
Self::Json(arg0, arg1) => f.debug_tuple("Json").field(arg0).field(arg1).finish(),
Self::Sse(_, arg1) => f.debug_tuple("Sse").field(arg1).finish(),
}
}
}
impl StreamableHttpPostResponse {
pub async fn expect_initialized<E>(
self,
) -> Result<(ServerJsonRpcMessage, Option<String>), StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Json(message, session_id) => Ok((message, session_id)),
Self::Sse(mut stream, session_id) => {
let event =
stream
.next()
.await
.ok_or(StreamableHttpError::UnexpectedServerResponse(
"empty sse stream".into(),
))??;
let message: ServerJsonRpcMessage =
serde_json::from_str(&event.data.unwrap_or_default())?;
Ok((message, session_id))
}
_ => Err(StreamableHttpError::UnexpectedServerResponse(
"expect initialized, accepted".into(),
)),
}
}
pub fn expect_json<E>(self) -> Result<ServerJsonRpcMessage, StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Json(message, ..) => Ok(message),
got => Err(StreamableHttpError::UnexpectedServerResponse(
format!("expect json, got {got:?}").into(),
)),
}
}
pub fn expect_accepted<E>(self) -> Result<(), StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Accepted => Ok(()),
got => Err(StreamableHttpError::UnexpectedServerResponse(
format!("expect accepted, got {got:?}").into(),
)),
}
}
}
pub trait StreamableHttpClient: Clone + Send + 'static {
type Error: std::error::Error + Send + Sync + 'static;
fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_header: Option<String>,
) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
+ Send
+ '_;
fn delete_session(
&self,
uri: Arc<str>,
session_id: Arc<str>,
auth_header: Option<String>,
) -> impl Future<Output = Result<(), StreamableHttpError<Self::Error>>> + Send + '_;
fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
last_event_id: Option<String>,
auth_header: Option<String>,
) -> impl Future<
Output = Result<
BoxStream<'static, Result<Sse, SseError>>,
StreamableHttpError<Self::Error>,
>,
> + Send
+ '_;
}
pub struct RetryConfig {
pub max_times: Option<usize>,
pub min_duration: Duration,
}
struct StreamableHttpClientReconnect<C> {
pub client: C,
pub session_id: Arc<str>,
pub uri: Arc<str>,
}
impl<C: StreamableHttpClient> SseStreamReconnect for StreamableHttpClientReconnect<C> {
type Error = StreamableHttpError<C::Error>;
type Future = BoxFuture<'static, Result<BoxedSseStream, Self::Error>>;
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future {
let client = self.client.clone();
let uri = self.uri.clone();
let session_id = self.session_id.clone();
let last_event_id = last_event_id.map(|s| s.to_owned());
Box::pin(async move {
client
.get_stream(uri, session_id, last_event_id, None)
.await
})
}
}
#[derive(Debug, Clone, Default)]
pub struct StreamableHttpClientWorker<C: StreamableHttpClient> {
pub client: C,
pub config: StreamableHttpClientTransportConfig,
}
impl<C: StreamableHttpClient + Default> StreamableHttpClientWorker<C> {
pub fn new_simple(url: impl Into<Arc<str>>) -> Self {
Self {
client: C::default(),
config: StreamableHttpClientTransportConfig {
uri: url.into(),
retry_config: Arc::new(ExponentialBackoff::default()),
channel_buffer_capacity: 16,
},
}
}
}
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
pub fn new(client: C, config: StreamableHttpClientTransportConfig) -> Self {
Self { client, config }
}
}
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
async fn execute_sse_stream(
sse_stream: SseAutoReconnectStream<StreamableHttpClientReconnect<C>>,
sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
ct: CancellationToken,
) -> Result<(), StreamableHttpError<C::Error>> {
let mut sse_stream = std::pin::pin!(sse_stream);
loop {
let message = tokio::select! {
event = sse_stream.next() => {
event
}
_ = ct.cancelled() => {
tracing::debug!("cancelled");
break;
}
};
let Some(message) = message.transpose()? else {
break;
};
let yield_result = sse_worker_tx.send(message).await;
if yield_result.is_err() {
tracing::trace!("streamable http transport worker dropped, exiting");
break;
}
}
Ok(())
}
}
impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
type Role = RoleClient;
type Error = StreamableHttpError<C::Error>;
fn err_closed() -> Self::Error {
StreamableHttpError::TransportChannelClosed
}
fn err_join(e: tokio::task::JoinError) -> Self::Error {
StreamableHttpError::TokioJoinError(e)
}
fn config(&self) -> super::worker::WorkerConfig {
super::worker::WorkerConfig {
name: Some("StreamableHttpClientWorker".into()),
channel_buffer_capacity: self.config.channel_buffer_capacity,
}
}
async fn run(
self,
mut context: super::worker::WorkerContext<Self>,
) -> Result<(), WorkerQuitReason> {
let channel_buffer_capacity = self.config.channel_buffer_capacity;
let (sse_worker_tx, mut sse_worker_rx) =
tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity);
let config = self.config.clone();
let transport_task_ct = context.cancellation_token.clone();
let _drop_guard = transport_task_ct.clone().drop_guard();
let WorkerSendRequest {
responder,
message: initialize_request,
} = context.recv_from_handler().await?;
let _ = responder.send(Ok(()));
let (message, session_id) = self
.client
.post_message(config.uri.clone(), initialize_request, None, None)
.await
.map_err(WorkerQuitReason::fatal_context("send initialize request"))?
.expect_initialized::<Self::Error>()
.await
.map_err(WorkerQuitReason::fatal_context(
"process initialize response",
))?;
let Some(session_id) = session_id else {
return Err(WorkerQuitReason::fatal(
"missing session id in initialize response",
"process initialize response",
));
};
let session_id: Arc<str> = session_id.into();
// delete session when drop guard is dropped
{
let ct = transport_task_ct.clone();
let client = self.client.clone();
let session_id = session_id.clone();
let url = config.uri.clone();
tokio::spawn(async move {
ct.cancelled().await;
let delete_session_result =
client.delete_session(url, session_id.clone(), None).await;
match delete_session_result {
Ok(_) => {
tracing::info!(session_id = session_id.as_ref(), "delete session success")
}
Err(StreamableHttpError::SeverDoesNotSupportDeleteSession) => {
tracing::info!(
session_id = session_id.as_ref(),
"server doesn't support delete session"
)
}
Err(e) => {
tracing::error!(
session_id = session_id.as_ref(),
"fail to delete session: {e}"
);
}
};
});
}
context.send_to_handler(message).await?;
let initialized_notification = context.recv_from_handler().await?;
// expect a initialized response
self.client
.post_message(
config.uri.clone(),
initialized_notification.message,
Some(session_id.clone()),
None,
)
.await
.map_err(WorkerQuitReason::fatal_context(
"send initialized notification",
))?
.expect_accepted::<Self::Error>()
.map_err(WorkerQuitReason::fatal_context(
"process initialized notification response",
))?;
let _ = initialized_notification.responder.send(Ok(()));
enum Event<W: Worker, E: std::error::Error + Send + Sync + 'static> {
ClientMessage(WorkerSendRequest<W>),
ServerMessage(ServerJsonRpcMessage),
StreamResult(Result<(), StreamableHttpError<E>>),
}
let mut streams = tokio::task::JoinSet::new();
match self
.client
.get_stream(config.uri.clone(), session_id.clone(), None, None)
.await
{
Ok(stream) => {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: self.client.clone(),
session_id: session_id.clone(),
uri: config.uri.clone(),
},
self.config.retry_config.clone(),
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
transport_task_ct.child_token(),
));
tracing::debug!("got common stream");
}
Err(StreamableHttpError::SeverDoesNotSupportSse) => {
tracing::debug!("server doesn't support sse, skip common stream");
}
Err(e) => {
// fail to get common stream
tracing::error!("fail to get common stream: {e}");
return Err(WorkerQuitReason::fatal(
"fail to get general purpose event stream",
"get general purpose event stream",
));
}
}
loop {
let event = tokio::select! {
_ = transport_task_ct.cancelled() => {
tracing::debug!("cancelled");
return Err(WorkerQuitReason::Cancelled);
}
message = context.recv_from_handler() => {
let message = message?;
Event::ClientMessage(message)
},
message = sse_worker_rx.recv() => {
let Some(message) = message else {
tracing::trace!("transport dropped, exiting");
return Err(WorkerQuitReason::HandlerTerminated);
};
Event::ServerMessage(message)
},
terminated_stream = streams.join_next(), if !streams.is_empty() => {
match terminated_stream {
Some(result) => {
Event::StreamResult(result.map_err(StreamableHttpError::TokioJoinError).and_then(std::convert::identity))
}
None => {
continue
}
}
}
};
match event {
Event::ClientMessage(send_request) => {
let WorkerSendRequest { message, responder } = send_request;
let response = self
.client
.post_message(config.uri.clone(), message, Some(session_id.clone()), None)
.await;
let send_result = match response {
Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => {
tracing::trace!("client message accepted");
Ok(())
}
Ok(StreamableHttpPostResponse::Json(message, ..)) => {
context.send_to_handler(message).await?;
Ok(())
}
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: self.client.clone(),
session_id: session_id.clone(),
uri: config.uri.clone(),
},
self.config.retry_config.clone(),
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
transport_task_ct.child_token(),
));
tracing::trace!("got new sse stream");
Ok(())
}
};
let _ = responder.send(send_result);
}
Event::ServerMessage(json_rpc_message) => {
// send the message to the handler
context.send_to_handler(json_rpc_message).await?;
}
Event::StreamResult(result) => {
if result.is_err() {
tracing::warn!(
"sse client event stream terminated with error: {:?}",
result
);
}
}
}
}
}
}
pub type StreamableHttpClientTransport<C> = WorkerTransport<StreamableHttpClientWorker<C>>;
impl<C: StreamableHttpClient> StreamableHttpClientTransport<C> {
pub fn with_client(client: C, config: StreamableHttpClientTransportConfig) -> Self {
let worker = StreamableHttpClientWorker::new(client, config);
WorkerTransport::spawn(worker)
}
}
#[derive(Debug, Clone)]
pub struct StreamableHttpClientTransportConfig {
pub uri: Arc<str>,
pub retry_config: Arc<dyn SseRetryPolicy>,
pub channel_buffer_capacity: usize,
}
impl StreamableHttpClientTransportConfig {
pub fn with_uri(uri: impl Into<Arc<str>>) -> Self {
Self {
uri: uri.into(),
..Default::default()
}
}
}
impl Default for StreamableHttpClientTransportConfig {
fn default() -> Self {
Self {
uri: "localhost".into(),
retry_config: Arc::new(ExponentialBackoff::default()),
channel_buffer_capacity: 16,
}
}
}

2
crates/rmcp/src/transport/streamable_http_server.rs
View file

@ -1,3 +1,5 @@
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub mod axum;
pub mod session;
pub use session::{SessionConfig, create_session};

91
crates/rmcp/src/transport/streamable_http_server/axum.rs
View file

@ -1,4 +1,10 @@
use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
use std::{
collections::HashMap,
io,
net::{Ipv6Addr, SocketAddr, SocketAddrV6},
sync::Arc,
time::Duration,
};
use axum::{
Json, Router,
@ -15,30 +21,28 @@ use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use super::session::{
EventId, HEADER_LAST_EVENT_ID, Session, SessionTransport, StreamableHttpMessageReceiver,
};
use super::session::{EventId, SessionHandle, SessionWorker, StreamableHttpMessageReceiver};
use crate::{
RoleServer, Service,
model::ClientJsonRpcMessage,
transport::{
common::axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
streamable_http_server::session::HEADER_SESSION_ID,
transport::common::{
axum::{DEFAULT_AUTO_PING_INTERVAL, SessionId, session_id},
http_header::{HEADER_LAST_EVENT_ID, HEADER_SESSION_ID},
},
};
type SessionManager = Arc<tokio::sync::RwLock<HashMap<SessionId, Session>>>;
type SessionManager = Arc<tokio::sync::RwLock<HashMap<SessionId, SessionHandle>>>;
#[derive(Clone)]
struct App {
session_manager: SessionManager,
transport_tx: tokio::sync::mpsc::UnboundedSender<SessionTransport>,
transport_tx: tokio::sync::mpsc::UnboundedSender<SessionWorker>,
sse_ping_interval: Duration,
}
impl App {
pub fn new(
sse_ping_interval: Duration,
) -> (Self, tokio::sync::mpsc::UnboundedReceiver<SessionTransport>) {
) -> (Self, tokio::sync::mpsc::UnboundedReceiver<SessionWorker>) {
let (transport_tx, transport_rx) = tokio::sync::mpsc::unbounded_channel();
(
Self {
@ -82,7 +86,7 @@ async fn post_handler(
let session = sm
.get(session_id)
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
session.handle().clone()
session.clone()
};
// inject request part
message.insert_extension(parts);
@ -138,7 +142,7 @@ async fn post_handler(
return Err((StatusCode::GONE, "session terminated").into_response());
};
let response = session.handle().initialize(message).await.map_err(|e| {
let response = session.initialize(message).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to initialize: {e}"),
@ -169,21 +173,24 @@ async fn get_handler(
let last_event_id = header_map
.get(HEADER_LAST_EVENT_ID)
.and_then(|v| v.to_str().ok());
match last_event_id {
Some(last_event_id) => {
let last_event_id = last_event_id.parse::<EventId>().map_err(|e| {
(StatusCode::BAD_REQUEST, format!("invalid event_id {e}")).into_response()
})?;
let sm = app.session_manager.read().await;
let session = sm.get(session_id).ok_or_else(|| {
let session = {
let sm = app.session_manager.read().await;
sm.get(session_id)
.ok_or_else(|| {
(
StatusCode::NOT_FOUND,
format!("session {session_id} not found"),
)
.into_response()
})?
.clone()
};
match last_event_id {
Some(last_event_id) => {
let last_event_id = last_event_id.parse::<EventId>().map_err(|e| {
(StatusCode::BAD_REQUEST, format!("invalid event_id {e}")).into_response()
})?;
let handle = session.handle();
let receiver = handle.resume(last_event_id).await.map_err(|e| {
let receiver = session.resume(last_event_id).await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("resume error {e}"),
@ -194,16 +201,7 @@ async fn get_handler(
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(app.sse_ping_interval)))
}
None => {
let sm = app.session_manager.read().await;
let session = sm.get(session_id).ok_or_else(|| {
(
StatusCode::NOT_FOUND,
format!("session {session_id} not found"),
)
.into_response()
})?;
let handle = session.handle();
let receiver = handle.establish_common_channel().await.map_err(|e| {
let receiver = session.establish_common_channel().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("establish common channel error {e}"),
@ -227,18 +225,19 @@ async fn delete_handler(
let session_id = session_id
.to_str()
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
let mut sm = app.session_manager.write().await;
let session = sm
.remove(session_id)
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
let cancel_result = session.cancel().await.map_err(|e| {
let session = {
let mut sm = app.session_manager.write().await;
sm.remove(session_id)
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?
};
session.close().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("fail to cancel session {session_id}: tokio join error: {e}"),
)
.into_response()
})?;
tracing::info!(session_id, ?cancel_result, "session deleted");
tracing::debug!(session_id, "session deleted");
Ok(StatusCode::ACCEPTED)
} else {
Err((StatusCode::BAD_REQUEST, "missing session id").into_response())
@ -252,10 +251,20 @@ pub struct StreamableHttpServerConfig {
pub ct: CancellationToken,
pub sse_keep_alive: Option<Duration>,
}
impl Default for StreamableHttpServerConfig {
fn default() -> Self {
Self {
bind: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 80, 0, 0)),
path: "/".to_string(),
ct: CancellationToken::new(),
sse_keep_alive: None,
}
}
}
#[derive(Debug)]
pub struct StreamableHttpServer {
transport_rx: tokio::sync::mpsc::UnboundedReceiver<SessionTransport>,
transport_rx: tokio::sync::mpsc::UnboundedReceiver<SessionWorker>,
pub config: StreamableHttpServerConfig,
}
@ -263,9 +272,7 @@ impl StreamableHttpServer {
pub async fn serve(bind: SocketAddr) -> io::Result<Self> {
Self::serve_with_config(StreamableHttpServerConfig {
bind,
path: "/".to_string(),
ct: CancellationToken::new(),
sse_keep_alive: None,
..Default::default()
})
.await
}
@ -333,7 +340,7 @@ impl StreamableHttpServer {
self.config.ct.cancel();
}
pub async fn next_transport(&mut self) -> Option<SessionTransport> {
pub async fn next_transport(&mut self) -> Option<SessionWorker> {
self.transport_rx.recv().await
}
}

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