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
11 changed files with 953 additions and 182 deletions
Download patch file Download diff file Expand all files Collapse all files

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,