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

feat!: add distributed SSE event store (#1024)

Browse source
This commit is contained in:
Dale Seo 2026-07-23 11:02:27 -04:00 • committed by GitHub
commit aad4d4e4c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 953 additions and 182 deletions
Download patch file Download diff file Expand all files Collapse all files

10
crates/rmcp/Cargo.toml
View file

@ -414,6 +414,16 @@ required-features = [
]
path = "tests/test_streamable_http_session_store.rs"
[[test]]
name = "test_streamable_http_event_store"
required-features = [
"client",
"server",
"transport-streamable-http-client-reqwest",
"transport-streamable-http-server",
]
path = "tests/test_streamable_http_event_store.rs"
[[test]]
name = "test_streamable_http_connection_reuse"
required-features = [

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

@ -31,7 +31,7 @@ where
async fn get_stream(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
session_id: Option<std::sync::Arc<str>>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -50,7 +50,7 @@ where
async fn get_stream_with_max_sse_event_size(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
session_id: Option<std::sync::Arc<str>>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,

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

@ -288,6 +288,7 @@ pin_project_lite::pin_project! {
where R: SseStreamReconnect
{
retry_policy: Arc<dyn SseRetryPolicy>,
reconnect_only_after_event_id: bool,
last_event_id: Option<String>,
server_retry_interval: Option<Duration>,
connector: R,
@ -304,6 +305,22 @@ impl<R: SseStreamReconnect> SseAutoReconnectStream<R> {
) -> Self {
Self {
retry_policy,
reconnect_only_after_event_id: false,
last_event_id: None,
server_retry_interval: None,
connector,
state: SseAutoReconnectStreamState::Connected { stream },
}
}
pub fn new_after_event_id(
stream: BoxedSseResponse,
connector: R,
retry_policy: Arc<dyn SseRetryPolicy>,
) -> Self {
Self {
retry_policy,
reconnect_only_after_event_id: true,
last_event_id: None,
server_retry_interval: None,
connector,
@ -317,6 +334,7 @@ impl<E: std::error::Error + Send> SseAutoReconnectStream<NeverReconnect<E>> {
pub(crate) fn never_reconnect(stream: BoxedSseResponse, error_when_reconnect: E) -> Self {
Self {
retry_policy: Arc::new(NeverRetry),
reconnect_only_after_event_id: false,
last_event_id: None,
server_retry_interval: None,
connector: NeverReconnect {
@ -409,6 +427,10 @@ where
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err));
}
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err));
}
this.connector
.handle_stream_error(&e, this.last_event_id.as_deref());
let retrying = this
@ -420,6 +442,13 @@ where
}
}
None => {
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
tracing::debug!(
"sse response ended before an event ID was received; cannot resume"
);
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(None);
}
// Per SEP-1699, a graceful stream close is
// reconnectable. If the server sent a `retry` field
// we MUST wait that long before reconnecting.
@ -686,4 +715,24 @@ mod tests {
&& attempts.load(Ordering::Relaxed) == 0
);
}
#[tokio::test]
async fn response_without_event_id_does_not_reconnect() {
let attempts = Arc::new(AtomicUsize::new(0));
let connector = CountingReconnect {
attempts: attempts.clone(),
};
let stream = SseAutoReconnectStream::new_after_event_id(
futures::stream::empty().boxed(),
connector,
Arc::new(FixedInterval {
max_times: Some(1),
duration: Duration::ZERO,
}),
);
let mut stream = std::pin::pin!(stream);
assert!(stream.next().await.is_none());
assert_eq!(attempts.load(Ordering::Relaxed), 0);
}
}

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

@ -52,7 +52,7 @@ impl StreamableHttpClient for reqwest::Client {
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -71,7 +71,7 @@ impl StreamableHttpClient for reqwest::Client {
async fn get_stream_with_max_sse_event_size(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -79,8 +79,10 @@ impl StreamableHttpClient for reqwest::Client {
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
let mut request_builder = self
.get(uri.as_ref())
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "))
.header(HEADER_SESSION_ID, session_id.as_ref());
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
if let Some(session_id) = session_id {
request_builder = request_builder.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);
}

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

@ -119,6 +119,15 @@ impl ServerSseMessage {
retry: Some(retry),
}
}
/// Create a retry hint without changing the client's last event ID.
pub fn retry(retry: Duration) -> Self {
Self {
event_id: None,
message: None,
retry: Some(retry),
}
}
}
pub(crate) fn sse_stream_response(

11
crates/rmcp/src/transport/common/unix_socket.rs
View file

@ -376,7 +376,7 @@ impl StreamableHttpClient for UnixSocketHttpClient {
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -396,7 +396,7 @@ impl StreamableHttpClient for UnixSocketHttpClient {
async fn get_stream_with_max_sse_event_size(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_token: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -410,8 +410,11 @@ impl StreamableHttpClient for UnixSocketHttpClient {
.header(
http::header::ACCEPT,
format!("{EVENT_STREAM_MIME_TYPE}, {JSON_MIME_TYPE}"),
)
.header(HEADER_SESSION_ID, session_id.as_ref());
);
if let Some(session_id) = session_id {
builder = builder.header(HEADER_SESSION_ID, session_id.as_ref());
}
if let Some(last_id) = last_event_id {
builder = builder.header(HEADER_LAST_EVENT_ID, last_id);

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

@ -348,10 +348,14 @@ pub trait StreamableHttpClient: Clone + Send + 'static {
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> impl Future<Output = Result<(), StreamableHttpError<Self::Error>>> + Send + '_;
/// Open an SSE stream, optionally scoped to a legacy session.
///
/// `session_id` is `None` when resuming a stateless response using only
/// `last_event_id`.
fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -375,7 +379,7 @@ pub trait StreamableHttpClient: Clone + Send + 'static {
fn get_stream_with_max_sse_event_size(
&self,
uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -399,7 +403,7 @@ pub struct RetryConfig {
struct StreamableHttpClientReconnect<C> {
pub client: C,
pub session_id: Arc<str>,
pub session_id: Option<Arc<str>>,
pub uri: Arc<str>,
pub auth_header: Option<String>,
pub custom_headers: HashMap<HeaderName, HeaderValue>,
@ -550,37 +554,6 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
Ok(())
}
/// Convert a raw SSE stream into a JSON-RPC message stream without
/// reconnection logic.
fn raw_sse_to_jsonrpc(
stream: BoxedSseStream,
) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
{
stream.filter_map(|event| async {
match event {
Err(e) => Some(Err(StreamableHttpError::Sse(e))),
Ok(sse) => {
let is_message =
matches!(sse.event.as_deref(), None | Some("") | Some("message"));
if !is_message {
return None;
}
let data = sse.data?;
if data.trim().is_empty() {
return None;
}
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
Ok(msg) => Some(Ok(msg)),
Err(e) => {
tracing::debug!("failed to deserialize server message: {e}");
None
}
}
}
}
})
}
/// Convert an SSE stream into JSON-RPC messages with reconnect semantics.
///
/// This is used for request-scoped SSE responses as well as the standalone
@ -590,7 +563,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
fn reconnecting_sse_to_jsonrpc(
stream: BoxedSseStream,
client: C,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
uri: Arc<str>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
@ -598,7 +571,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
retry_config: Arc<dyn SseRetryPolicy>,
) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
{
SseAutoReconnectStream::new(
SseAutoReconnectStream::new_after_event_id(
stream,
StreamableHttpClientReconnect {
client,
@ -614,10 +587,8 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
/// Convert a POST response SSE stream into JSON-RPC messages.
///
/// Stateful sessions can resume via GET when the response stream closes
/// before the server sends the matching JSON-RPC response. Stateless
/// transports do not have enough state to resume, so they keep the raw
/// SSE-to-JSON-RPC mapping.
/// Request-scoped streams resume via GET once the server has supplied an
/// event ID. The session header remains optional for stateless transports.
fn response_sse_to_jsonrpc(
stream: BoxedSseStream,
session_id: Option<Arc<str>>,
@ -628,20 +599,17 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
max_sse_event_size: usize,
retry_config: Arc<dyn SseRetryPolicy>,
) -> BoxStream<'static, Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> {
match session_id {
Some(session_id) => Self::reconnecting_sse_to_jsonrpc(
stream,
client,
session_id,
uri,
auth_header,
custom_headers,
max_sse_event_size,
retry_config,
)
.boxed(),
None => Self::raw_sse_to_jsonrpc(stream).boxed(),
}
Self::reconnecting_sse_to_jsonrpc(
stream,
client,
session_id,
uri,
auth_header,
custom_headers,
max_sse_event_size,
retry_config,
)
.boxed()
}
async fn execute_sse_stream(
@ -709,7 +677,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
let result = match client
.get_stream_with_max_sse_event_size(
uri,
session_id.clone(),
Some(session_id.clone()),
None,
auth_header,
protocol_headers.clone(),
@ -722,7 +690,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
stream,
StreamableHttpClientReconnect {
client,
session_id,
session_id: Some(session_id),
uri: reconnect_uri,
auth_header: reconnect_auth_header,
custom_headers: protocol_headers,
@ -1559,7 +1527,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
/// async fn get_stream(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _session_id: Option<Arc<str>>,
/// _last_event_id: Option<String>,
/// _auth_header: Option<String>,
/// _custom_headers: HashMap<HeaderName, HeaderValue>,
@ -1646,7 +1614,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientTransport<C> {
/// async fn get_stream(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _session_id: Option<Arc<str>>,
/// _last_event_id: Option<String>,
/// _auth_header: Option<String>,
/// _custom_headers: HashMap<HeaderName, HeaderValue>,
@ -1778,11 +1746,110 @@ impl Default for StreamableHttpClientTransportConfig {
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use serde_json::json;
use super::*;
use crate::model::{ListToolsResult, NumberOrString, ServerResult, Tool};
type ReconnectAttempt = (Option<String>, Option<String>);
#[derive(Clone, Default)]
struct StatelessReconnectClient {
reconnects: Arc<Mutex<Vec<ReconnectAttempt>>>,
}
impl StreamableHttpClient for StatelessReconnectClient {
type Error = std::io::Error;
async fn post_message(
&self,
_uri: Arc<str>,
_message: ClientJsonRpcMessage,
_session_id: Option<Arc<str>>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
Err(StreamableHttpError::UnexpectedServerResponse(
"unexpected POST".into(),
))
}
async fn delete_session(
&self,
_uri: Arc<str>,
_session_id: Arc<str>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(), StreamableHttpError<Self::Error>> {
Ok(())
}
async fn get_stream(
&self,
_uri: Arc<str>,
session_id: Option<Arc<str>>,
last_event_id: Option<String>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<BoxedSseStream, StreamableHttpError<Self::Error>> {
self.reconnects
.lock()
.expect("lock reconnects")
.push((session_id.map(|id| id.to_string()), last_event_id));
let response = ServerJsonRpcMessage::response(
ServerResult::ListToolsResult(ListToolsResult::default()),
NumberOrString::Number(1),
);
Ok(futures::stream::once(async move {
Ok(Sse {
event: None,
data: Some(serde_json::to_string(&response).expect("serialize response")),
id: Some("event-1".into()),
retry: None,
})
})
.boxed())
}
}
#[tokio::test]
async fn stateless_response_reconnects_with_last_event_id() {
let initial = futures::stream::iter([Ok(Sse {
event: None,
data: None,
id: Some("event-0".into()),
retry: Some(0),
})])
.boxed();
let client = StatelessReconnectClient::default();
let reconnects = client.reconnects.clone();
let stream =
StreamableHttpClientWorker::<StatelessReconnectClient>::response_sse_to_jsonrpc(
initial,
None,
client,
Arc::from("http://localhost/mcp"),
None,
HashMap::new(),
DEFAULT_MAX_SSE_EVENT_SIZE,
Arc::new(ExponentialBackoff {
max_times: Some(1),
base_duration: Duration::ZERO,
}),
);
let mut stream = std::pin::pin!(stream);
let message = stream.next().await.expect("replayed response").unwrap();
assert!(matches!(message, ServerJsonRpcMessage::Response(_)));
assert_eq!(
reconnects.lock().expect("lock reconnects").as_slice(),
&[(None, Some("event-0".into()))]
);
}
fn tool(name: &'static str, annotation: serde_json::Value) -> Tool {
let schema = json!({
"type": "object",

12
crates/rmcp/src/transport/streamable_http_server/session.rs
View file

@ -20,6 +20,8 @@
//! Implement the [`SessionManager`] trait to back sessions with a database,
//! Redis, or any other external store.
use std::sync::Arc;
use futures::Stream;
pub use crate::transport::common::server_side_http::{ServerSseMessage, SessionId};
@ -32,7 +34,10 @@ pub mod local;
pub mod never;
pub mod store;
pub use store::{SessionState, SessionStore, SessionStoreError};
pub use store::{
EventId, EventStore, EventStoreError, EventStream, SessionState, SessionStore,
SessionStoreError, StreamId,
};
/// Extension marker inserted into the `initialize` request extensions during a
/// session restore replay. Handlers can check for its presence to distinguish a
@ -151,4 +156,9 @@ pub trait SessionManager: Send + Sync + 'static {
) -> impl Future<Output = Result<RestoreOutcome<Self::Transport>, Self::Error>> + Send {
futures::future::ready(Ok(RestoreOutcome::NotSupported))
}
/// Return the shared event store used for resumable SSE streams.
fn event_store(&self) -> Option<Arc<dyn EventStore>> {
None
}
}

164
crates/rmcp/src/transport/streamable_http_server/session/local.rs
View file

@ -1,6 +1,7 @@
use std::{
collections::{HashMap, HashSet, VecDeque},
num::ParseIntError,
sync::Arc,
time::{Duration, Instant},
};
@ -32,6 +33,15 @@ use crate::{
pub struct LocalSessionManager {
pub sessions: tokio::sync::RwLock<HashMap<SessionId, LocalSessionHandle>>,
pub session_config: SessionConfig,
event_store: Option<Arc<dyn EventStore>>,
}
impl LocalSessionManager {
/// Configure this session manager to use a shared event store.
pub fn with_event_store(mut self, event_store: Arc<dyn EventStore>) -> Self {
self.event_store = Some(event_store);
self
}
}
#[derive(Debug, Error)]
@ -49,7 +59,11 @@ impl SessionManager for LocalSessionManager {
type Transport = WorkerTransport<LocalSessionWorker>;
async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> {
let id = session_id();
let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
let (handle, worker) = create_local_session_with_event_store(
id.clone(),
self.session_config.clone(),
self.event_store.clone(),
);
self.sessions.write().await.insert(id.clone(), handle);
Ok((id, WorkerTransport::spawn(worker)))
}
@ -95,15 +109,7 @@ impl SessionManager for LocalSessionManager {
let receiver = handle.establish_request_wise_channel().await?;
let http_request_id = receiver.http_request_id;
handle.push_message(message, http_request_id).await?;
let priming = self.session_config.sse_retry.map(|retry| {
let event_id = match http_request_id {
Some(id) => format!("0/{id}"),
None => "0".into(),
};
ServerSseMessage::priming(event_id, retry)
});
Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner)))
Ok(ReceiverStream::new(receiver.inner))
}
async fn create_standalone_stream(
@ -123,12 +129,19 @@ impl SessionManager for LocalSessionManager {
id: &SessionId,
last_event_id: String,
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
if let Some(event_store) = &self.event_store {
let stream = event_store
.replay_events_after(&last_event_id)
.await
.map_err(SessionError::EventStore)?;
return Ok(stream.left_stream());
}
let sessions = self.sessions.read().await;
let handle = sessions
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let receiver = handle.resume(last_event_id.parse()?).await?;
Ok(ReceiverStream::new(receiver.inner))
Ok(ReceiverStream::new(receiver.inner).right_stream())
}
async fn accept_message(
@ -153,10 +166,18 @@ impl SessionManager for LocalSessionManager {
// A concurrent request already restored this session.
return Ok(RestoreOutcome::AlreadyPresent);
}
let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
let (handle, worker) = create_local_session_with_event_store(
id.clone(),
self.session_config.clone(),
self.event_store.clone(),
);
sessions.insert(id, handle);
Ok(RestoreOutcome::Restored(WorkerTransport::spawn(worker)))
}
fn event_store(&self) -> Option<Arc<dyn EventStore>> {
self.event_store.clone()
}
}
/// `<index>/request_id>`
@ -209,7 +230,9 @@ impl std::str::FromStr for EventId {
}
}
use super::{RestoreOutcome, ServerSseMessage, SessionManager};
use super::{
EventStore, EventStoreError, RestoreOutcome, ServerSseMessage, SessionManager, StreamId,
};
struct CachedTx {
tx: Sender<ServerSseMessage>,
@ -217,6 +240,8 @@ struct CachedTx {
http_request_id: Option<HttpRequestId>,
capacity: usize,
starting_index: usize,
stream_id: StreamId,
event_store: Option<Arc<dyn EventStore>>,
}
impl CachedTx {
@ -224,6 +249,8 @@ impl CachedTx {
tx: Sender<ServerSseMessage>,
http_request_id: Option<HttpRequestId>,
starting_index: usize,
stream_id: StreamId,
event_store: Option<Arc<dyn EventStore>>,
) -> Self {
Self {
cache: VecDeque::with_capacity(tx.capacity()),
@ -231,10 +258,16 @@ impl CachedTx {
tx,
http_request_id,
starting_index,
stream_id,
event_store,
}
}
fn new_common(tx: Sender<ServerSseMessage>) -> Self {
Self::new(tx, None, 0)
fn new_common(
tx: Sender<ServerSseMessage>,
session_id: &SessionId,
event_store: Option<Arc<dyn EventStore>>,
) -> Self {
Self::new(tx, None, 0, format!("{session_id}:common"), event_store)
}
fn next_event_id(&self) -> EventId {
@ -253,16 +286,31 @@ impl CachedTx {
}
}
async fn send(&mut self, message: ServerJsonRpcMessage) {
let event_id = self.next_event_id();
let message = ServerSseMessage::new(event_id.to_string(), message);
self.cache_and_send(message).await;
async fn send(&mut self, message: ServerJsonRpcMessage) -> Result<(), SessionError> {
self.store_cache_and_send(ServerSseMessage::from_message(message))
.await
}
async fn send_priming(&mut self, retry: Duration) {
let event_id = self.next_event_id();
let message = ServerSseMessage::priming(event_id.to_string(), retry);
self.cache_and_send(message).await;
async fn send_priming(&mut self, retry: Duration) -> Result<(), SessionError> {
self.store_cache_and_send(ServerSseMessage::retry(retry))
.await
}
async fn store_cache_and_send(
&mut self,
mut event: ServerSseMessage,
) -> Result<(), SessionError> {
let event_id = if let Some(event_store) = &self.event_store {
event_store
.store_event(&self.stream_id, &event)
.await
.map_err(SessionError::EventStore)?
} else {
self.next_event_id().to_string()
};
event.event_id = Some(event_id);
self.cache_and_send(event).await;
Ok(())
}
async fn cache_and_send(&mut self, message: ServerSseMessage) {
@ -330,6 +378,7 @@ pub struct LocalSessionWorker {
shadow_txs: Vec<Sender<ServerSseMessage>>,
event_rx: Receiver<SessionEvent>,
session_config: SessionConfig,
event_store: Option<Arc<dyn EventStore>>,
}
impl LocalSessionWorker {
@ -353,6 +402,8 @@ pub enum SessionError {
InvalidEventId,
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Event store error: {0}")]
EventStore(#[source] EventStoreError),
}
impl From<SessionError> for std::io::Error {
@ -450,12 +501,21 @@ impl LocalSessionWorker {
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let http_request_id = self.next_http_request_id();
let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let starting_index = usize::from(self.session_config.sse_retry.is_some());
let mut cached_tx = CachedTx::new(
tx,
Some(http_request_id),
0,
uuid::Uuid::new_v4().to_string(),
self.event_store.clone(),
);
if let Some(retry) = self.session_config.sse_retry {
cached_tx.send_priming(retry).await?;
}
self.tx_router.insert(
http_request_id,
HttpRequestWise {
resources: Default::default(),
tx: CachedTx::new(tx, Some(http_request_id), starting_index),
tx: cached_tx,
completed_at: None,
},
);
@ -548,7 +608,7 @@ impl LocalSessionWorker {
match outbound_channel {
OutboundChannel::RequestWise { id, close } => {
if let Some(request_wise) = self.tx_router.get_mut(&id) {
request_wise.tx.send(message).await;
request_wise.tx.send(message).await?;
if close {
if let Some(channel) = self.tx_router.remove(&id) {
for resource in channel.resources {
@ -560,7 +620,7 @@ impl LocalSessionWorker {
return Err(SessionError::ChannelClosed(Some(id)));
}
}
OutboundChannel::Common => self.common.send(message).await,
OutboundChannel::Common => self.common.send(message).await?,
}
Ok(())
}
@ -593,10 +653,20 @@ impl LocalSessionWorker {
inner: rx,
})
}
None => self.resume_or_shadow_common(last_event_id.index).await,
None => {
self.resume_or_shadow_common(Some(last_event_id.index))
.await
}
}
}
async fn establish_common_channel(
&mut self,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let last_event_index = self.event_store.is_none().then_some(0);
self.resume_or_shadow_common(last_event_index).await
}
/// Resume the common channel, or create a shadow stream if the primary is
/// still active.
///
@ -611,7 +681,7 @@ impl LocalSessionWorker {
/// killing each other by repeatedly replacing the common channel sender.
async fn resume_or_shadow_common(
&mut self,
last_event_index: usize,
last_event_index: Option<usize>,
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let is_replacing_dead_primary = self.common.tx.is_closed();
let capacity = if is_replacing_dead_primary {
@ -624,9 +694,9 @@ impl LocalSessionWorker {
// Primary common channel is dead — replace it.
tracing::debug!("Replacing dead common channel with new primary");
self.common.tx = tx;
// Replay cached messages from where the client left off so
// server-initiated requests and notifications are not lost.
self.common.sync(last_event_index).await?;
if let Some(last_event_index) = last_event_index {
self.common.sync(last_event_index).await?;
}
} else {
// Primary common channel is still active. Create a shadow stream
// that stays alive via SSE keep-alive but doesn't receive
@ -668,7 +738,7 @@ impl LocalSessionWorker {
// Send priming event if retry interval is specified
if let Some(interval) = retry_interval {
request_wise.tx.send_priming(interval).await;
request_wise.tx.send_priming(interval).await?;
}
// Close the stream by dropping the sender
@ -685,7 +755,7 @@ impl LocalSessionWorker {
None => {
// Send priming event if retry interval is specified
if let Some(interval) = retry_interval {
self.common.send_priming(interval).await;
self.common.send_priming(interval).await?;
}
// Close the stream by dropping the sender
@ -732,6 +802,9 @@ pub enum SessionEvent {
retry_interval: Option<Duration>,
responder: oneshot::Sender<Result<(), SessionError>>,
},
EstablishCommonChannel {
responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
},
}
#[derive(Debug, Clone)]
@ -820,13 +893,7 @@ impl LocalSessionHandle {
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.event_tx
.send(SessionEvent::Resume {
last_event_id: EventId {
http_request_id: None,
index: 0,
},
responder: tx,
})
.send(SessionEvent::EstablishCommonChannel { responder: tx })
.await
.map_err(|_| SessionError::SessionServiceTerminated)?;
rx.await
@ -1085,6 +1152,10 @@ impl Worker for LocalSessionWorker {
let handle_result = self.establish_request_wise_channel().await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::EstablishCommonChannel { responder }) => {
let handle_result = self.establish_common_channel().await;
let _ = responder.send(handle_result);
}
InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
id,
responder,
@ -1179,11 +1250,19 @@ impl Default for SessionConfig {
pub fn create_local_session(
id: impl Into<SessionId>,
config: SessionConfig,
) -> (LocalSessionHandle, LocalSessionWorker) {
create_local_session_with_event_store(id, config, None)
}
fn create_local_session_with_event_store(
id: impl Into<SessionId>,
config: SessionConfig,
event_store: Option<Arc<dyn EventStore>>,
) -> (LocalSessionHandle, LocalSessionWorker) {
let id = id.into();
let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
let common = CachedTx::new_common(common_tx);
let common = CachedTx::new_common(common_tx, &id, event_store.clone());
tracing::info!(session_id = ?id, "create new session");
let handle = LocalSessionHandle {
event_tx,
@ -1198,6 +1277,7 @@ pub fn create_local_session(
shadow_txs: Vec::new(),
event_rx,
session_config: config.clone(),
event_store,
};
(handle, session_worker)
}

20
crates/rmcp/src/transport/streamable_http_server/session/never.rs
View file

@ -1,7 +1,9 @@
use std::sync::Arc;
use futures::Stream;
use thiserror::Error;
use super::{ServerSseMessage, SessionId, SessionManager};
use super::{EventStore, ServerSseMessage, SessionId, SessionManager};
use crate::{
RoleServer,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
@ -14,7 +16,17 @@ use crate::{
pub struct ErrorSessionManagementNotSupported;
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct NeverSessionManager {}
pub struct NeverSessionManager {
event_store: Option<Arc<dyn EventStore>>,
}
impl NeverSessionManager {
/// Configure resumable SSE storage without enabling sessions.
pub fn with_event_store(mut self, event_store: Arc<dyn EventStore>) -> Self {
self.event_store = Some(event_store);
self
}
}
#[non_exhaustive]
pub enum NeverTransport {}
impl Transport<RoleServer> for NeverTransport {
@ -107,4 +119,8 @@ impl SessionManager for NeverSessionManager {
) -> impl Future<Output = Result<(), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn event_store(&self) -> Option<Arc<dyn EventStore>> {
self.event_store.clone()
}
}

59
crates/rmcp/src/transport/streamable_http_server/session/store.rs
View file

@ -1,4 +1,61 @@
use crate::model::InitializeRequestParams;
use std::pin::Pin;
use futures::Stream;
use crate::{
model::InitializeRequestParams, transport::common::server_side_http::ServerSseMessage,
};
/// An opaque identifier for a persisted SSE event.
pub type EventId = String;
/// An opaque identifier for an SSE stream.
pub type StreamId = String;
/// A stream of persisted SSE events in delivery order.
pub type EventStream = Pin<Box<dyn Stream<Item = ServerSseMessage> + Send + Sync + 'static>>;
/// Type alias for boxed event store errors.
pub type EventStoreError = Box<dyn std::error::Error + Send + Sync + 'static>;
/// Persistent storage for resumable Streamable HTTP events.
///
/// Implementations typically use a database or distributed log shared by all
/// server instances. Event IDs must be globally unique across all streams, and
/// events must be committed before [`EventStore::store_event`] returns so the
/// returned ID is safe to send to a client.
#[async_trait::async_trait]
pub trait EventStore: Send + Sync + 'static {
/// Persist an event and return the opaque ID clients should receive.
///
/// The store assigns a globally unique ID and must retain its association
/// with `stream_id` so a later replay only returns events from that stream.
async fn store_event(
&self,
stream_id: &str,
event: &ServerSseMessage,
) -> Result<EventId, EventStoreError>;
/// Return events strictly after `last_event_id` in delivery order.
///
/// Implementations must locate the stream from the globally unique event
/// ID and yield only later events from that stream, with their originally
/// assigned event IDs.
///
/// A finite stream enables reconnect-and-poll behavior. Implementations
/// backed by a distributed log may keep the stream open to deliver new
/// events as they are appended by any server instance.
async fn replay_events_after(
&self,
last_event_id: &str,
) -> Result<EventStream, EventStoreError>;
}
impl std::fmt::Debug for dyn EventStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("<EventStore>")
}
}
/// State persisted to an external store for cross-instance session recovery.
///

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

@ -19,7 +19,8 @@ use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use super::session::{
RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker, SessionState, SessionStore,
EventStore, EventStoreError, RestoreOutcome, SessionId, SessionManager, SessionRestoreMarker,
SessionState, SessionStore,
};
use crate::{
RoleServer,
@ -49,6 +50,7 @@ use crate::{
/// Default maximum POST request body size (4 MiB).
pub(crate) const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024;
const STATELESS_STREAM_CHANNEL_CAPACITY: usize = 16;
#[non_exhaustive]
#[derive(Debug, Clone)]
@ -321,6 +323,21 @@ fn method_not_allowed_response() -> BoxResponse {
.expect("valid response")
}
async fn persist_and_forward_event(
event_store: &dyn EventStore,
stream_id: &str,
mut event: ServerSseMessage,
output: &mut Option<tokio::sync::mpsc::Sender<ServerSseMessage>>,
) -> Result<(), EventStoreError> {
event.event_id = Some(event_store.store_event(stream_id, &event).await?);
if let Some(sender) = output {
if sender.send(event).await.is_err() {
*output = None;
}
}
Ok(())
}
fn invalid_request_jsonrpc_response(
id: Option<RequestId>,
message: impl Into<Cow<'static, str>>,
@ -920,6 +937,98 @@ where
(self.service_factory)()
}
fn persisted_stateless_stream(
&self,
first: Option<ServerJsonRpcMessage>,
mut receiver: tokio::sync::mpsc::Receiver<ServerJsonRpcMessage>,
request_ct: CancellationToken,
event_store: Arc<dyn EventStore>,
) -> ReceiverStream<ServerSseMessage> {
let (sender, output) = tokio::sync::mpsc::channel(STATELESS_STREAM_CHANNEL_CAPACITY);
let stream_id = uuid::Uuid::new_v4().to_string();
let retry = self.config.sse_retry;
let server_ct = self.config.cancellation_token.child_token();
tokio::spawn(async move {
let mut sender = Some(sender);
if let Some(retry) = retry {
if let Err(error) = persist_and_forward_event(
event_store.as_ref(),
&stream_id,
ServerSseMessage::retry(retry),
&mut sender,
)
.await
{
tracing::error!(%stream_id, %error, "failed to persist SSE priming event");
request_ct.cancel();
return;
}
}
let mut first = first;
loop {
let message = if let Some(message) = first.take() {
Some(message)
} else {
tokio::select! {
message = receiver.recv() => message,
_ = server_ct.cancelled() => {
request_ct.cancel();
None
}
}
};
let Some(message) = message else {
break;
};
tracing::trace!(?message);
if let Err(error) = persist_and_forward_event(
event_store.as_ref(),
&stream_id,
ServerSseMessage::from_message(message),
&mut sender,
)
.await
{
tracing::error!(%stream_id, %error, "failed to persist SSE event");
request_ct.cancel();
break;
}
}
});
ReceiverStream::new(output)
}
fn stateless_sse_response(
&self,
first: Option<ServerJsonRpcMessage>,
receiver: tokio::sync::mpsc::Receiver<ServerJsonRpcMessage>,
request_ct: CancellationToken,
) -> BoxResponse {
if let Some(event_store) = self.session_manager.event_store() {
let stream = self.persisted_stateless_stream(first, receiver, request_ct, event_store);
sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
)
} else {
let stream = futures::stream::iter(first)
.chain(ReceiverStream::new(receiver))
.map(|message| {
tracing::trace!(?message);
ServerSseMessage::from_message(message)
});
sse_stream_response(
CancelOnDisconnect::new(stream, request_ct),
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
)
}
}
// The HTTP status must be known before opening an SSE stream.
async fn serve_negotiated_request_directly(
&self,
@ -979,19 +1088,7 @@ where
return jsonrpc_message_response(first, true);
}
// The handler may still be streaming, so guard the response: dropping it
// (client disconnect) must cancel the handler.
let stream = futures::stream::once(async move { first })
.chain(ReceiverStream::new(receiver))
.map(|message| {
tracing::trace!(?message);
ServerSseMessage::from_message(message)
});
Ok(sse_stream_response(
CancelOnDisconnect::new(stream, request_ct),
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
Ok(self.stateless_sse_response(Some(first), receiver, request_ct))
}
/// Returns the cached input schema for `name`, constructing a service once
@ -1235,15 +1332,18 @@ where
return response;
}
let method = request.method().clone();
let allowed_methods = match self.config.legacy_session_mode {
true => "GET, POST, DELETE",
false => "POST",
let supports_stateless_replay = self.session_manager.event_store().is_some();
let allowed_methods = match (self.config.legacy_session_mode, supports_stateless_replay) {
(true, _) => "GET, POST, DELETE",
(false, true) => "GET, POST",
(false, false) => "POST",
};
let result = match (method, self.config.legacy_session_mode) {
(Method::POST, _) => self.handle_post(request).await,
// if legacy session mode is disabled, we don't support GET or DELETE because there is no session
(Method::GET, true) => self.handle_get(request).await,
(Method::DELETE, true) => self.handle_delete(request).await,
let result = match method {
Method::POST => self.handle_post(request).await,
Method::GET if self.config.legacy_session_mode || supports_stateless_replay => {
self.handle_get(request).await
}
Method::DELETE if self.config.legacy_session_mode => self.handle_delete(request).await,
_ => {
// Handle other methods or return an error
let response = Response::builder()
@ -1264,9 +1364,6 @@ where
B: Body + Send + 'static,
B::Error: Display,
{
if !is_legacy_request(None, request.headers())? {
return Ok(method_not_allowed_response());
}
// check accept header
if !request
.headers()
@ -1284,6 +1381,32 @@ where
)
.expect("valid response"));
}
let request_uses_legacy_protocol = is_legacy_request(None, request.headers())?;
let legacy_request = self.config.legacy_session_mode && request_uses_legacy_protocol;
if !legacy_request {
let Some(last_event_id) = request
.headers()
.get(HEADER_LAST_EVENT_ID)
.and_then(|value| value.to_str().ok())
else {
return Ok(method_not_allowed_response());
};
let Some(event_store) = self.session_manager.event_store() else {
return Ok(method_not_allowed_response());
};
let stream = match event_store.replay_events_after(last_event_id).await {
Ok(stream) => stream,
Err(error) => {
tracing::warn!(%error, "stateless SSE resume failed, returning empty stream");
Box::pin(futures::stream::empty())
}
};
return Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
));
}
// check session id
let session_id = request
.headers()
@ -1361,7 +1484,11 @@ where
.await
.map_err(internal_error_response("create standalone stream"))?;
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage::priming("0", retry);
let priming = if self.session_manager.event_store().is_some() {
ServerSseMessage::retry(retry)
} else {
ServerSseMessage::priming("0", retry)
};
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
@ -1658,10 +1785,9 @@ where
request.request.extensions_mut().insert(part);
let (transport, mut receiver) =
OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
// Give this stateless request its own cancellation token so a
// client disconnect can cancel the in-flight handler (#857). A
// stateless request is one-shot (no session, no resumption), so a
// dropped response is terminal and safe to cancel.
// Give this stateless request its own cancellation token so an
// unpersisted response can cancel the in-flight handler on
// disconnect (#857).
let request_ct = CancellationToken::new();
let service =
serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
@ -1710,35 +1836,10 @@ where
.body(Full::new(Bytes::from(body)).boxed())
.expect("valid response"))
} else {
// The handler emitted an intermediate message and is still
// running, so guard the streamed sequence too: dropping it
// (client disconnect) must cancel the handler.
let first = futures::stream::once(async move {
ServerSseMessage::from_message(message)
});
let remaining = ReceiverStream::new(receiver).map(|message| {
tracing::trace!(?message);
ServerSseMessage::from_message(message)
});
Ok(sse_stream_response(
CancelOnDisconnect::new(first.chain(remaining), request_ct),
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
Ok(self.stateless_sse_response(Some(message), receiver, request_ct))
}
} else {
// SSE mode (default): cancel the handler if the client
// disconnects (drops the response stream) before it completes.
let stream = ReceiverStream::new(receiver).map(|message| {
tracing::trace!(?message);
ServerSseMessage::from_message(message)
});
let stream = CancelOnDisconnect::new(stream, request_ct);
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
Ok(self.stateless_sse_response(None, receiver, request_ct))
}
}
ClientJsonRpcMessage::Notification(_notification) => {
@ -1821,16 +1922,11 @@ where
}
pin_project! {
/// Wraps a stateless SSE response stream so a client disconnect cancels the
/// in-flight request.
/// Cancels an unpersisted stateless request when its response is dropped.
///
/// A stateless streamable-HTTP request is one-shot: it has no session and no
/// resumption, so a dropped response stream means the client is gone for
/// good. When the stream is dropped *before* it ends naturally, the request's
/// cancellation token is fired, which stops the dedicated `serve_directly`
/// loop and cancels the handler's `RequestContext::ct` (see #857). If the
/// stream ends naturally (the request completed), the guard is disarmed so
/// normal completion cancels nothing.
/// Persisted requests keep running so another connection can resume them.
/// Without an event store, dropping the stream fires the request's
/// cancellation token. Natural completion disarms the guard.
struct CancelOnDisconnect<S> {
#[pin]
inner: S,

372
crates/rmcp/tests/test_streamable_http_event_store.rs Normal file
View file

@ -0,0 +1,372 @@
#![cfg(all(
feature = "client",
feature = "server",
feature = "transport-streamable-http-client-reqwest",
feature = "transport-streamable-http-server",
not(feature = "local")
))]
use std::{
collections::HashMap,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
};
use futures::StreamExt;
use rmcp::{
ErrorData, ServerHandler,
model::{
CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
ProgressNotificationParam, ServerCapabilities, ServerInfo,
},
service::RequestContext,
transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService,
session::{
EventStore, EventStoreError, EventStream, ServerSseMessage, SessionState, SessionStore,
SessionStoreError, local::LocalSessionManager, never::NeverSessionManager,
},
},
};
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
#[derive(Clone, Default)]
struct InMemorySessionStore(Arc<RwLock<HashMap<String, SessionState>>>);
#[async_trait::async_trait]
impl SessionStore for InMemorySessionStore {
async fn load(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
Ok(self.0.read().await.get(session_id).cloned())
}
async fn store(&self, session_id: &str, state: &SessionState) -> Result<(), SessionStoreError> {
self.0
.write()
.await
.insert(session_id.to_owned(), state.clone());
Ok(())
}
async fn delete(&self, session_id: &str) -> Result<(), SessionStoreError> {
self.0.write().await.remove(session_id);
Ok(())
}
}
#[derive(Clone)]
struct StoredEvent {
stream_id: String,
event: ServerSseMessage,
}
#[derive(Clone, Default)]
struct InMemoryEventStore {
events: Arc<RwLock<Vec<StoredEvent>>>,
next_id: Arc<AtomicU64>,
}
#[async_trait::async_trait]
impl EventStore for InMemoryEventStore {
async fn store_event(
&self,
stream_id: &str,
event: &ServerSseMessage,
) -> Result<String, EventStoreError> {
let event_id = format!("event-{}", self.next_id.fetch_add(1, Ordering::Relaxed));
let mut event = event.clone();
event.event_id = Some(event_id.clone());
self.events.write().await.push(StoredEvent {
stream_id: stream_id.to_owned(),
event,
});
Ok(event_id)
}
async fn replay_events_after(
&self,
last_event_id: &str,
) -> Result<EventStream, EventStoreError> {
let events = self.events.read().await;
let last_index = events
.iter()
.position(|stored| stored.event.event_id.as_deref() == Some(last_event_id))
.ok_or_else(|| std::io::Error::other("event not found"))?;
let stream_id = events[last_index].stream_id.clone();
let replay = events
.iter()
.skip(last_index + 1)
.filter(|stored| stored.stream_id == stream_id)
.map(|stored| stored.event.clone())
.collect::<Vec<_>>();
Ok(Box::pin(futures::stream::iter(replay)))
}
}
#[derive(Clone)]
struct ProgressServer;
impl ServerHandler for ProgressServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<rmcp::RoleServer>,
) -> Result<CallToolResponse, ErrorData> {
if request.name == "progress" || request.name == "slow-progress" {
let progress_token = context
.meta
.get_progress_token()
.expect("request includes progressToken");
context
.peer
.notify_progress(
ProgressNotificationParam::new(progress_token, 50.0)
.with_total(100.0)
.with_message("working"),
)
.await
.expect("progress notification is delivered");
}
if request.name == "slow-progress" {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into())
}
}
async fn spawn_server(
session_store: Arc<dyn SessionStore>,
event_store: Arc<dyn EventStore>,
cancellation_token: &CancellationToken,
legacy_session_mode: bool,
) -> anyhow::Result<(String, tokio::task::JoinHandle<()>)> {
let config = {
let mut config = StreamableHttpServerConfig::default();
config.sse_keep_alive = None;
config.legacy_session_mode = legacy_session_mode;
config.cancellation_token = cancellation_token.child_token();
config.session_store = Some(session_store);
config
};
let router = if legacy_session_mode {
let session_manager =
Arc::new(LocalSessionManager::default().with_event_store(event_store));
let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config);
axum::Router::new().nest_service("/mcp", service)
} else {
let session_manager =
Arc::new(NeverSessionManager::default().with_event_store(event_store));
let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config);
axum::Router::new().nest_service("/mcp", service)
};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let address = listener.local_addr()?;
let handle = tokio::spawn({
let cancellation_token = cancellation_token.clone();
async move {
let _ = axum::serve(listener, router)
.with_graceful_shutdown(async move { cancellation_token.cancelled_owned().await })
.await;
}
});
Ok((format!("http://{address}/mcp"), handle))
}
fn event_id_containing<'a>(body: &'a str, needle: &str) -> Option<&'a str> {
body.split("\n\n")
.find(|event| event.contains(needle))?
.lines()
.find_map(|line| line.strip_prefix("id: "))
}
#[tokio::test]
async fn restored_instance_replays_events_from_shared_store() -> anyhow::Result<()> {
let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default());
let event_store: Arc<dyn EventStore> = Arc::new(InMemoryEventStore::default());
let http = reqwest::Client::new();
let cancellation_a = CancellationToken::new();
let (url_a, server_a) = spawn_server(
session_store.clone(),
event_store.clone(),
&cancellation_a,
true,
)
.await?;
let initialize = http
.post(&url_a)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
.send()
.await?;
let session_id = initialize
.headers()
.get("mcp-session-id")
.expect("initialize returns a session ID")
.to_str()?
.to_owned();
let _ = initialize.text().await?;
let initialized_status = http
.post(&url_a)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("Mcp-Session-Id", &session_id)
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
.send()
.await?
.status();
assert_eq!(initialized_status, reqwest::StatusCode::ACCEPTED);
let original_body = http
.post(&url_a)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("Mcp-Session-Id", &session_id)
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"progress","arguments":{},"_meta":{"progressToken":"progress-1"}}}"#)
.send()
.await?
.text()
.await?;
let progress_event_id = event_id_containing(&original_body, "notifications/progress")
.expect("progress event has a persisted event ID")
.to_owned();
cancellation_a.cancel();
server_a.await?;
let cancellation_b = CancellationToken::new();
let (url_b, server_b) = spawn_server(session_store, event_store, &cancellation_b, true).await?;
let replay = http
.get(&url_b)
.header("Accept", "text/event-stream")
.header("Mcp-Session-Id", &session_id)
.header("Mcp-Protocol-Version", "2025-06-18")
.header("Last-Event-ID", progress_event_id)
.send()
.await?;
assert_eq!(replay.status(), reqwest::StatusCode::OK);
let replay_body = replay.text().await?;
assert!(
replay_body.contains(r#""id":2"#),
"instance B should replay the final response stored by instance A: {replay_body}"
);
cancellation_b.cancel();
server_b.await?;
Ok(())
}
#[tokio::test]
async fn stateless_instance_replays_events_from_shared_store() -> anyhow::Result<()> {
let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default());
let event_store = Arc::new(InMemoryEventStore::default());
let http = reqwest::Client::new();
let cancellation_a = CancellationToken::new();
let (url_a, server_a) = spawn_server(
session_store.clone(),
event_store.clone(),
&cancellation_a,
false,
)
.await?;
let original = http
.post(&url_a)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("MCP-Protocol-Version", "2026-07-28")
.header("Mcp-Method", "tools/call")
.header("Mcp-Name", "slow-progress")
.body(
r#"{
"jsonrpc":"2.0",
"id":2,
"method":"tools/call",
"params":{
"name":"slow-progress",
"arguments":{},
"_meta":{
"progressToken":"progress-1",
"io.modelcontextprotocol/protocolVersion":"2026-07-28",
"io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"},
"io.modelcontextprotocol/clientCapabilities":{}
}
}
}"#,
)
.send()
.await?;
assert_eq!(original.status(), reqwest::StatusCode::OK);
assert!(
original.headers().get("mcp-session-id").is_none(),
"stateless response must not create a session"
);
let mut body = original.bytes_stream();
let mut received = String::new();
let progress_event_id = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let chunk = body
.next()
.await
.expect("response remains open until progress arrives")?;
received.push_str(&String::from_utf8_lossy(&chunk));
if let Some(event_id) = event_id_containing(&received, "notifications/progress") {
return Ok::<_, reqwest::Error>(event_id.to_owned());
}
}
})
.await??;
drop(body);
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let stored_response = event_store.events.read().await.iter().any(|stored| {
stored.event.message.as_ref().is_some_and(|message| {
matches!(
message.as_ref(),
rmcp::model::ServerJsonRpcMessage::Response(_)
)
})
});
if stored_response {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
})
.await?;
cancellation_a.cancel();
server_a.await?;
let cancellation_b = CancellationToken::new();
let (url_b, server_b) =
spawn_server(session_store, event_store, &cancellation_b, false).await?;
let replay = http
.get(&url_b)
.header("Accept", "text/event-stream")
.header("MCP-Protocol-Version", "2026-07-28")
.header("Last-Event-ID", progress_event_id)
.send()
.await?;
assert_eq!(replay.status(), reqwest::StatusCode::OK);
let replay_body = replay.text().await?;
assert!(
replay_body.contains(r#""id":2"#),
"instance B should replay the stateless response stored by instance A: {replay_body}"
);
cancellation_b.cancel();
server_b.await?;
Ok(())
}

4
crates/rmcp/tests/test_streamable_http_stale_session.rs
View file

@ -160,7 +160,7 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient {
async fn get_stream(
&self,
_uri: Arc<str>,
session_id: Arc<str>,
session_id: Option<Arc<str>>,
_last_event_id: Option<String>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
@ -168,7 +168,7 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient {
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
StreamableHttpError<Self::Error>,
> {
if session_id.as_ref() == "session-1" {
if session_id.as_deref() == Some("session-1") {
let cancel = self.stale_stream_cancelled.clone();
Ok(Box::pin(stream::once(async move {
cancel.cancelled_owned().await;