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

feat: stateless mode of streamable http client (#233)

Browse source
This commit is contained in:
4t145 2025-05-28 19:37:18 +08:00 • committed by GitHub
commit ff71a52615
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 117 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

36
crates/rmcp/src/transport/common/client_side_sse.rs
View file

@ -76,6 +76,28 @@ impl SseRetryPolicy for ExponentialBackoff {
} }
} }
#[derive(Debug, Clone, Copy, Default)]
pub struct NeverRetry;
impl SseRetryPolicy for NeverRetry {
fn retry(&self, _current_times: usize) -> Option<Duration> {
None
}
}
#[derive(Debug, Default)]
pub struct NeverReconnect<E> {
error: Option<E>,
}
impl<E: std::error::Error + Send> SseStreamReconnect for NeverReconnect<E> {
type Error = E;
type Future = futures::future::Ready<Result<BoxedSseResponse, Self::Error>>;
fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future {
futures::future::ready(Err(self.error.take().expect("should not be called again")))
}
}
pub(crate) trait SseStreamReconnect { pub(crate) trait SseStreamReconnect {
type Error: std::error::Error; type Error: std::error::Error;
type Future: Future<Output = Result<BoxedSseResponse, Self::Error>> + Send; type Future: Future<Output = Result<BoxedSseResponse, Self::Error>> + Send;
@ -111,6 +133,20 @@ impl<R: SseStreamReconnect> SseAutoReconnectStream<R> {
} }
} }
impl<E: std::error::Error + Send> SseAutoReconnectStream<NeverReconnect<E>> {
pub fn never_reconnect(stream: BoxedSseResponse, error_when_reconnect: E) -> Self {
Self {
retry_policy: Arc::new(NeverRetry),
last_event_id: None,
server_retry_interval: None,
connector: NeverReconnect {
error: Some(error_when_reconnect),
},
state: SseAutoReconnectStreamState::Connected { stream },
}
}
}
pin_project_lite::pin_project! { pin_project_lite::pin_project! {
#[project = SseAutoReconnectStreamStateProj] #[project = SseAutoReconnectStreamStateProj]
pub enum SseAutoReconnectStreamState<F> { pub enum SseAutoReconnectStreamState<F> {

141
crates/rmcp/src/transport/streamable_http_client.rs
View file

@ -1,6 +1,6 @@
use std::{borrow::Cow, sync::Arc, time::Duration}; use std::{borrow::Cow, sync::Arc, time::Duration};
use futures::{StreamExt, future::BoxFuture, stream::BoxStream}; use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream};
pub use sse_stream::Error as SseError; pub use sse_stream::Error as SseError;
use sse_stream::Sse; use sse_stream::Sse;
use thiserror::Error; use thiserror::Error;
@ -193,8 +193,7 @@ impl<C: StreamableHttpClient + Default> StreamableHttpClientWorker<C> {
client: C::default(), client: C::default(),
config: StreamableHttpClientTransportConfig { config: StreamableHttpClientTransportConfig {
uri: url.into(), uri: url.into(),
retry_config: Arc::new(ExponentialBackoff::default()), ..Default::default()
channel_buffer_capacity: 16,
}, },
} }
} }
@ -208,7 +207,9 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> { impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
async fn execute_sse_stream( async fn execute_sse_stream(
sse_stream: SseAutoReconnectStream<StreamableHttpClientReconnect<C>>, sse_stream: impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>>
+ Send
+ 'static,
sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>, sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
ct: CancellationToken, ct: CancellationToken,
) -> Result<(), StreamableHttpError<C::Error>> { ) -> Result<(), StreamableHttpError<C::Error>> {
@ -277,16 +278,19 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
.map_err(WorkerQuitReason::fatal_context( .map_err(WorkerQuitReason::fatal_context(
"process initialize response", "process initialize response",
))?; ))?;
let Some(session_id) = session_id else { let session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
return Err(WorkerQuitReason::fatal( Some(session_id.into())
"missing session id in initialize response", } else {
"process initialize response", if !self.config.allow_stateless {
)); return Err(WorkerQuitReason::fatal(
"missing session id in initialize response",
"process initialize response",
));
}
None
}; };
let session_id: Arc<str> = session_id.into();
// delete session when drop guard is dropped // delete session when drop guard is dropped
{ if let Some(session_id) = &session_id {
let ct = transport_task_ct.clone(); let ct = transport_task_ct.clone();
let client = self.client.clone(); let client = self.client.clone();
let session_id = session_id.clone(); let session_id = session_id.clone();
@ -322,7 +326,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
.post_message( .post_message(
config.uri.clone(), config.uri.clone(),
initialized_notification.message, initialized_notification.message,
Some(session_id.clone()), session_id.clone(),
None, None,
) )
.await .await
@ -340,38 +344,40 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
StreamResult(Result<(), StreamableHttpError<E>>), StreamResult(Result<(), StreamableHttpError<E>>),
} }
let mut streams = tokio::task::JoinSet::new(); let mut streams = tokio::task::JoinSet::new();
match self if let Some(session_id) = &session_id {
.client match self
.get_stream(config.uri.clone(), session_id.clone(), None, None) .client
.await .get_stream(config.uri.clone(), session_id.clone(), None, None)
{ .await
Ok(stream) => { {
let sse_stream = SseAutoReconnectStream::new( Ok(stream) => {
stream, let sse_stream = SseAutoReconnectStream::new(
StreamableHttpClientReconnect { stream,
client: self.client.clone(), StreamableHttpClientReconnect {
session_id: session_id.clone(), client: self.client.clone(),
uri: config.uri.clone(), session_id: session_id.clone(),
}, uri: config.uri.clone(),
self.config.retry_config.clone(), },
); self.config.retry_config.clone(),
streams.spawn(Self::execute_sse_stream( );
sse_stream, streams.spawn(Self::execute_sse_stream(
sse_worker_tx.clone(), sse_stream,
transport_task_ct.child_token(), sse_worker_tx.clone(),
)); transport_task_ct.child_token(),
tracing::debug!("got common stream"); ));
} tracing::debug!("got common stream");
Err(StreamableHttpError::SeverDoesNotSupportSse) => { }
tracing::debug!("server doesn't support sse, skip common stream"); Err(StreamableHttpError::SeverDoesNotSupportSse) => {
} tracing::debug!("server doesn't support sse, skip common stream");
Err(e) => { }
// fail to get common stream Err(e) => {
tracing::error!("fail to get common stream: {e}"); // fail to get common stream
return Err(WorkerQuitReason::fatal( tracing::error!("fail to get common stream: {e}");
"fail to get general purpose event stream", return Err(WorkerQuitReason::fatal(
"get general purpose event stream", "fail to get general purpose event stream",
)); "get general purpose event stream",
));
}
} }
} }
loop { loop {
@ -407,7 +413,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
let WorkerSendRequest { message, responder } = send_request; let WorkerSendRequest { message, responder } = send_request;
let response = self let response = self
.client .client
.post_message(config.uri.clone(), message, Some(session_id.clone()), None) .post_message(config.uri.clone(), message, session_id.clone(), None)
.await; .await;
let send_result = match response { let send_result = match response {
Err(e) => Err(e), Err(e) => Err(e),
@ -420,20 +426,32 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
Ok(()) Ok(())
} }
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => { Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
let sse_stream = SseAutoReconnectStream::new( if let Some(session_id) = &session_id {
stream, let sse_stream = SseAutoReconnectStream::new(
StreamableHttpClientReconnect { stream,
client: self.client.clone(), StreamableHttpClientReconnect {
session_id: session_id.clone(), client: self.client.clone(),
uri: config.uri.clone(), session_id: session_id.clone(),
}, uri: config.uri.clone(),
self.config.retry_config.clone(), },
); self.config.retry_config.clone(),
streams.spawn(Self::execute_sse_stream( );
sse_stream, streams.spawn(Self::execute_sse_stream(
sse_worker_tx.clone(), sse_stream,
transport_task_ct.child_token(), sse_worker_tx.clone(),
)); transport_task_ct.child_token(),
));
} else {
let sse_stream = SseAutoReconnectStream::never_reconnect(
stream,
StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
transport_task_ct.child_token(),
));
}
tracing::trace!("got new sse stream"); tracing::trace!("got new sse stream");
Ok(()) Ok(())
} }
@ -470,6 +488,8 @@ pub struct StreamableHttpClientTransportConfig {
pub uri: Arc<str>, pub uri: Arc<str>,
pub retry_config: Arc<dyn SseRetryPolicy>, pub retry_config: Arc<dyn SseRetryPolicy>,
pub channel_buffer_capacity: usize, pub channel_buffer_capacity: usize,
/// if true, the transport will not require a session to be established
pub allow_stateless: bool,
} }
impl StreamableHttpClientTransportConfig { impl StreamableHttpClientTransportConfig {
@ -487,6 +507,7 @@ impl Default for StreamableHttpClientTransportConfig {
uri: "localhost".into(), uri: "localhost".into(),
retry_config: Arc::new(ExponentialBackoff::default()), retry_config: Arc::new(ExponentialBackoff::default()),
channel_buffer_capacity: 16, channel_buffer_capacity: 16,
allow_stateless: true,
} }
} }
} }