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

fix(streamable-http): gracefully shutdown while client connected (#494)

Browse source 
* fix(streamable-http): gracefully shutdown while client connected

* fix: adviced comments

* fix: windows test build
This commit is contained in:
Jonson Petard 2025-12-02 09:22:52 +08:00 • committed by GitHub
commit 3dee024325
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

19
crates/rmcp/src/transport/common/server_side_http.rs
View file

@ -6,6 +6,7 @@ use http::Response;
use http_body::Body;
use http_body_util::{BodyExt, Empty, Full, combinators::BoxBody};
use sse_stream::{KeepAlive, Sse, SseBody};
use tokio_util::sync::CancellationToken;
use super::http_header::EVENT_STREAM_MIME_TYPE;
use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
@ -65,20 +66,26 @@ pub struct ServerSseMessage {
pub(crate) fn sse_stream_response(
stream: impl futures::Stream<Item = ServerSseMessage> + Send + Sync + 'static,
keep_alive: Option<Duration>,
ct: CancellationToken,
) -> Response<BoxBody<Bytes, Infallible>> {
use futures::StreamExt;
let stream = SseBody::new(stream.map(|message| {
let data = serde_json::to_string(&message.message).expect("valid message");
let mut sse = Sse::default().data(data);
sse.id = message.event_id;
Result::<Sse, Infallible>::Ok(sse)
}));
let stream = stream
.map(|message| {
let data = serde_json::to_string(&message.message).expect("valid message");
let mut sse = Sse::default().data(data);
sse.id = message.event_id;
Result::<Sse, Infallible>::Ok(sse)
})
.take_until(async move { ct.cancelled().await });
let stream = SseBody::new(stream);
let stream = match keep_alive {
Some(duration) => stream
.with_keep_alive::<TokioTimer>(KeepAlive::new().interval(duration))
.boxed(),
None => stream.boxed(),
};
Response::builder()
.status(http::StatusCode::OK)
.header(http::header::CONTENT_TYPE, EVENT_STREAM_MIME_TYPE)

27
crates/rmcp/src/transport/streamable_http_server/tower.rs
View file

@ -6,6 +6,7 @@ use http::{Method, Request, Response, header::ALLOW};
use http_body::Body;
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use super::session::SessionManager;
use crate::{
@ -33,6 +34,11 @@ pub struct StreamableHttpServerConfig {
pub sse_keep_alive: Option<Duration>,
/// If true, the server will create a session for each request and keep it alive.
pub stateful_mode: bool,
/// Cancellation token for the Streamable HTTP server.
///
/// When this token is cancelled, all active sessions are terminated and
/// the server stops accepting new requests.
pub cancellation_token: CancellationToken,
}
impl Default for StreamableHttpServerConfig {
@ -40,6 +46,7 @@ impl Default for StreamableHttpServerConfig {
Self {
sse_keep_alive: Some(Duration::from_secs(15)),
stateful_mode: true,
cancellation_token: CancellationToken::new(),
}
}
}
@ -209,7 +216,11 @@ where
.resume(&session_id, last_event_id)
.await
.map_err(internal_error_response("resume session"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
} else {
// create standalone stream
let stream = self
@ -217,7 +228,11 @@ where
.create_standalone_stream(&session_id)
.await
.map_err(internal_error_response("create standalone stream"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
}
@ -307,7 +322,11 @@ where
.create_stream(&session_id, message)
.await
.map_err(internal_error_response("get session"))?;
Ok(sse_stream_response(stream, self.config.sse_keep_alive))
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
ClientJsonRpcMessage::Notification(_)
| ClientJsonRpcMessage::Response(_)
@ -380,6 +399,7 @@ where
}
}),
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
);
response.headers_mut().insert(
@ -413,6 +433,7 @@ where
}
}),
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
ClientJsonRpcMessage::Notification(_notification) => {