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

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()
},
)
}
}